package minsecurity.identity.persist.sqlite.db;

import minsecurity.Common;
import minsecurity.certificate.cert.CertUtils;
import minsecurity.certificate.cert.Certificate;
import minsecurity.crypto.KeyUtils;
import minsecurity.crypto.PrivateKeyInterface;
import minsecurity.crypto.PublicKeyInterface;
import minsecurity.identity.Identity;
import minsecurity.identity.KeyParam;
import minsecurity.identity.persist.sqlite.Sqlite;

import java.sql.*;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;


public final class Db {
    private static Identity getIdentityByNameFromStorage(String name, Connection c) throws Exception {
        PreparedStatement pstmt = c.prepareStatement("SELECT * FROM identityinfo WHERE name = ? LIMIT 1");// LIMIT 1代替Queryrow
        pstmt.setString(1, name);
        ResultSet rs = pstmt.executeQuery();
        pstmt.close();
        return getIdentityFromSqlRow(rs);
    }

    private static Identity getIdentityFromSqlRow(ResultSet rs) throws Exception {
        if (!rs.isBeforeFirst())
            return null;

        String pubStr = "", priStr = "", pass = "", certStr = "", name = "", prikeyRawByte = "";
        int algo = 0, sign = 0, def = 0;

        while (rs.next()){
            pubStr = rs.getString("pubKey");
            priStr = rs.getString("priKey");
            pass = rs.getString("pass");
            certStr = rs.getString("cert");
            name = rs.getString("name");
            prikeyRawByte = rs.getString("prikey_raw_byte");
            algo = rs.getInt("pubkey_algo");
            sign = rs.getInt("signature_algo");
            def = rs.getInt("is_default");
        }

        byte[] pubByte = Base64.getDecoder().decode(pubStr);
        PublicKeyInterface pubKey = KeyUtils.unMarshalPublicKey(pubByte, algo);
        byte[] priByte = Base64.getDecoder().decode(priStr);
        PrivateKeyInterface priKey = KeyUtils.unMarshalPrivateKey(priByte, algo);

        Certificate cert = CertUtils.fromPem(certStr, null, Common.SM4ECB);
        byte[] priKeyByte = Base64.getDecoder().decode(prikeyRawByte);
        KeyParam keyParam = new KeyParam(algo, sign);
        return new Identity(name, keyParam, priKey, priKeyByte, pubKey, pass, cert, def == 1);
    }

    private static List<Identity> getIdentityFromSqlRows(ResultSet rs) throws Exception {
        if (!rs.isBeforeFirst())
            return new ArrayList<>();

        List<Identity> res = new ArrayList<>();
        String pubStr = "", priStr = "", pass = "", certStr = "", name = "", prikeyRawByte = "";
        int algo = 0, sign = 0, def = 0;

        while (rs.next()){
            pubStr = rs.getString("pubKey");
            priStr = rs.getString("priKey");
            pass = rs.getString("pass");
            certStr = rs.getString("cert");
            name = rs.getString("name");
            prikeyRawByte = rs.getString("prikey_raw_byte");
            algo = rs.getInt("pubkey_algo");
            sign = rs.getInt("signature_algo");
            def = rs.getInt("is_default");
            byte[] pubByte = Base64.getDecoder().decode(pubStr);
            PublicKeyInterface pubKey = KeyUtils.unMarshalPublicKey(pubByte, algo);
            byte[] priByte = Base64.getDecoder().decode(priStr);
            PrivateKeyInterface priKey = KeyUtils.unMarshalPrivateKey(priByte, algo);

            Certificate cert = CertUtils.fromPem(certStr, null, Common.SM4ECB);
            byte[] priKeyByte = Base64.getDecoder().decode(prikeyRawByte);
            KeyParam keyParam = new KeyParam(algo, sign);
            res.add(new Identity(name, keyParam, priKey, priKeyByte, pubKey, pass, cert, def == 1));
        }

        return res;
    }

    private static Identity getDefaultIdentityFromStorage(ResultSet rs) throws Exception{
        if (!rs.isBeforeFirst())
            return null;

        String pubStr = "", priStr = "", pass = "", certStr = "", name = "", prikeyRawByte = "";
        int algo = 0, sign = 0, def = 0;

        while (rs.next()){
            pubStr = rs.getString("pubKey");
            priStr = rs.getString("priKey");
            pass = rs.getString("pass");
            certStr = rs.getString("cert");
            name = rs.getString("name");
            prikeyRawByte = rs.getString("prikey_raw_byte");
            algo = rs.getInt("pubkey_algo");
            sign = rs.getInt("signature_algo");
            def = rs.getInt("is_default");
        }

        byte[] pubByte = Base64.getDecoder().decode(pubStr);
        PublicKeyInterface pubKey = KeyUtils.unMarshalPublicKey(pubByte, algo);
        byte[] priByte = Base64.getDecoder().decode(priStr);
        PrivateKeyInterface priKey = KeyUtils.unMarshalPrivateKey(priByte, algo);

        Certificate cert = CertUtils.fromPem(certStr, null, Common.SM4ECB);
        byte[] priKeyByte = Base64.getDecoder().decode(prikeyRawByte);
        KeyParam keyParam = new KeyParam(algo, sign);
        return new Identity(name, keyParam, priKey, priKeyByte, pubKey, pass, cert, def == 1);
    }


    private static void setDefaultIdentityFromStorage(String name, Connection c) throws Exception{
        PreparedStatement stmt = c.prepareStatement("update identityinfo set is_default = 1 where name= ?");
        stmt.setString(1, name);
        stmt.executeUpdate();
        stmt.close();
    }

    private static void cancelDefaultIdentityFromStorage(String name, Connection c) throws Exception{
        PreparedStatement stmt = c.prepareStatement("update identityinfo set is_default = 0 where name= ?");
        stmt.setString(1, name);
        stmt.executeUpdate();
        stmt.close();
    }

    public static List<Identity> getAllIdentityFromStorage(String name) throws Exception {
        Connection c = Sqlite.getInstance().getConn();
        Statement stmt = c.createStatement();
        ResultSet rs = stmt.executeQuery("SELECT * from identityinfo");
        stmt.close();
        c.close();
        List<Identity> res = getIdentityFromSqlRows(rs);
        return res;
    }

    public static Identity getIdentityByNameFromStorage(String name) throws Exception {
        Connection c = Sqlite.getInstance().getConn();
        c.close();
        return getIdentityByNameFromStorage(name, c);
    }

    private static void setDefaultIdentityByNameInStorage(String name) throws Exception {
        Connection c = Sqlite.getInstance().getConn();
        c.setAutoCommit(false);
        try{
            PreparedStatement pstmt = c.prepareStatement("select * from identityinfo where name = ? LIMIT 1");
            pstmt.setString(1, name);
            ResultSet rs = pstmt.executeQuery(String.format("select * from identityinfo where name = %s LIMIT 1", name));
            getDefaultIdentityFromStorage(rs);
            pstmt.close();

            PreparedStatement pstmt2 = c.prepareStatement("SELECT * from identityinfo where is_default = ?");
            pstmt2.setInt(1, 1);
            rs = pstmt2.executeQuery();
            Identity id = getDefaultIdentityFromStorage(rs);
            if (id != null){
                cancelDefaultIdentityFromStorage(id.getName(), c);
            }
            setDefaultIdentityFromStorage(name, c);
            pstmt2.close();
            c.commit();
            c.close();
        }catch (Exception ex){
            System.out.println(ex.getMessage());
            c.rollback();
            c.close();
            throw ex;
        }
    }

    public static void SetDefaultIdentityByNameInStorage(String name) throws Exception {
        for (int i = 0; i < 4; i++){
            try{
                setDefaultIdentityByNameInStorage(name);
                return;
            }catch (Exception ex){
                if (i == 3)
                    return;
                if (ex.getMessage().equals("database is locked"))
                    Thread.sleep(50);
                else
                    throw ex;
            }
        }
    }

    public static void deleteIdentityByName(String name) throws Exception{
        Connection c = Sqlite.getInstance().getConn();
        PreparedStatement pstmt = c.prepareStatement("elete from identityinfo where name=?");
        pstmt.setString(1, name);
        pstmt.executeUpdate();
        pstmt.close();
        c.close();
    }

    public static void persistIdentity(Identity identity) throws Exception{
        Identity id = getIdentityByNameFromStorage(identity.getName());
        if (id != null)
            throw new Exception("The name has existed");

        Connection c = Sqlite.getInstance().getConn();
        String pubStr = "", priStr = "", certStr, prikeyRawByte = "";
        int algo, sign;

        if (identity.hasPrivateKey()){
            byte[] priByte = identity.getPrikey().getBytes();
            priStr = Base64.getEncoder().encodeToString(priByte);
        }
        if (identity.getPubkey() != null){
            byte[] pubByte = identity.getPubkey().getBytes();
            pubStr = Base64.getEncoder().encodeToString(pubByte);
        }

        algo = identity.getKeyParam().PublicKeyAlgorithm;
        sign = identity.getKeyParam().SignatureAlgorithm;

        certStr = CertUtils.toPem(identity.getCert(), null, Common.SM4ECB);

        if (identity.getPrikeyRawByte() != null)
            prikeyRawByte = Base64.getEncoder().encodeToString(identity.getPrikeyRawByte());

        PreparedStatement pstmt = c.prepareStatement("INSERT INTO identityinfo(name, pubkey, prikey, pubkey_algo, signature_algo, pass, cert,prikey_raw_byte) values(?,?,?,?,?,?,?,?)");
        pstmt.setString(1, identity.getName());
        pstmt.setString(2, priStr);
        pstmt.setString(3, pubStr);
        pstmt.setInt(4, algo);
        pstmt.setInt(5, sign);
        pstmt.setString(6, identity.getPasswd());
        pstmt.setString(7, certStr);
        pstmt.setString(8, prikeyRawByte);

        pstmt.executeUpdate();
        pstmt.close();
        c.close();
    }

    public Identity getDefaultIdentityFromStorage() throws Exception{
        Connection c = Sqlite.getInstance().getConn();
        PreparedStatement pstmt = c.prepareStatement("SELECT * from identityinfo where is_default= ? LIMIT 1");
        pstmt.setInt(1, 1);
        ResultSet rs = pstmt.executeQuery();
        return getDefaultIdentityFromStorage(rs);
    }
}