package club.clubk.ktag.apps.mqttbroker;

import android.content.Context;
import android.util.Log;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.math.BigInteger;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.cert.Certificate;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Date;

public class SSLCertificateManager {
    private static final String TAG = "SSLCertManager";
    private static final String KEYSTORE_FILE = "mqtt_broker.p12";
    private static final String KEYSTORE_PASSWORD = "mqtt_broker_password";
    private static final String KEY_ALIAS = "mqtt_broker_key";

    public static File getKeystoreFile(Context context) {
        return new File(context.getFilesDir(), KEYSTORE_FILE);
    }

    public static String getKeystorePassword() {
        return KEYSTORE_PASSWORD;
    }

    public static String getKeyAlias() {
        return KEY_ALIAS;
    }

    public static boolean keystoreExists(Context context) {
        return getKeystoreFile(context).exists();
    }

    public static void generateSelfSignedCertificate(Context context) throws Exception {
        Log.d(TAG, "Generating self-signed certificate...");

        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
        keyGen.initialize(2048, new SecureRandom());
        KeyPair keyPair = keyGen.generateKeyPair();
        PrivateKey privateKey = keyPair.getPrivate();

        long now = System.currentTimeMillis();
        Date notBefore = new Date(now - 24 * 60 * 60 * 1000);
        Date notAfter = new Date(now + 365L * 24 * 60 * 60 * 1000);

        BigInteger serial = BigInteger.valueOf(now);
        String dn = "CN=MQTT Broker, O=KTag Apps, C=US";
        byte[] certDer = buildSelfSignedCertificateDer(
                keyPair, privateKey, serial, notBefore, notAfter, dn);

        CertificateFactory cf = CertificateFactory.getInstance("X.509");
        X509Certificate cert = (X509Certificate) cf.generateCertificate(
                new ByteArrayInputStream(certDer));

        cert.verify(keyPair.getPublic());

        KeyStore keyStore = KeyStore.getInstance("PKCS12");
        keyStore.load(null, null);

        Certificate[] chain = new Certificate[]{cert};
        keyStore.setKeyEntry(KEY_ALIAS, privateKey, KEYSTORE_PASSWORD.toCharArray(), chain);

        File keystoreFile = getKeystoreFile(context);
        try (FileOutputStream fos = new FileOutputStream(keystoreFile)) {
            keyStore.store(fos, KEYSTORE_PASSWORD.toCharArray());
        }

        Log.d(TAG, "Self-signed certificate generated and saved to " + keystoreFile.getAbsolutePath());
    }

    public static KeyStore loadKeystore(Context context) throws Exception {
        File keystoreFile = getKeystoreFile(context);
        if (!keystoreFile.exists()) {
            throw new Exception("Keystore file does not exist");
        }

        KeyStore keyStore = KeyStore.getInstance("PKCS12");
        try (FileInputStream fis = new FileInputStream(keystoreFile)) {
            keyStore.load(fis, KEYSTORE_PASSWORD.toCharArray());
        }

        return keyStore;
    }

    public static void deleteKeystore(Context context) {
        File keystoreFile = getKeystoreFile(context);
        if (keystoreFile.exists()) {
            if (keystoreFile.delete()) {
                Log.d(TAG, "Keystore deleted");
            } else {
                Log.e(TAG, "Failed to delete keystore");
            }
        }
    }

    public static String getCertificateInfo(Context context) {
        try {
            KeyStore keyStore = loadKeystore(context);
            X509Certificate cert = (X509Certificate) keyStore.getCertificate(KEY_ALIAS);

            if (cert != null) {
                StringBuilder info = new StringBuilder();
                info.append("Subject: ").append(cert.getSubjectDN().getName()).append("\n");
                info.append("Issuer: ").append(cert.getIssuerDN().getName()).append("\n");
                info.append("Valid from: ").append(cert.getNotBefore()).append("\n");
                info.append("Valid until: ").append(cert.getNotAfter()).append("\n");
                info.append("Serial: ").append(cert.getSerialNumber()).append("\n");
                return info.toString();
            }
        } catch (Exception e) {
            Log.e(TAG, "Error getting certificate info", e);
        }
        return "Certificate info not available";
    }

    // -----------------------------------------------------------------------
    // DER / ASN.1 certificate builder using only java.security APIs
    // -----------------------------------------------------------------------

    private static byte[] buildSelfSignedCertificateDer(
            KeyPair keyPair, PrivateKey privateKey,
            BigInteger serial, Date notBefore, Date notAfter,
            String distinguishedName) throws Exception {

        byte[] tbsCert = buildTbsCertificate(
                serial, notBefore, notAfter, distinguishedName,
                keyPair.getPublic().getEncoded());

        Signature sig = Signature.getInstance("SHA256withRSA");
        sig.initSign(privateKey);
        sig.update(tbsCert);
        byte[] signature = sig.sign();

        byte[] signatureAlgorithm = derSequence(new byte[][]{
                derOid(new int[]{1, 2, 840, 113549, 1, 1, 11}),
                derNull()
        });

        byte[] signatureBits = new byte[signature.length + 1];
        signatureBits[0] = 0x00;
        System.arraycopy(signature, 0, signatureBits, 1, signature.length);
        byte[] signatureBitString = derBitString(signatureBits);

        return derSequence(new byte[][]{tbsCert, signatureAlgorithm, signatureBitString});
    }

    private static byte[] buildTbsCertificate(
            BigInteger serial, Date notBefore, Date notAfter,
            String distinguishedName, byte[] publicKeyEncoded) throws Exception {

        byte[] version = derExplicitTag(0, derInteger(BigInteger.valueOf(2)));
        byte[] serialNumber = derInteger(serial);

        byte[] signatureAlgorithm = derSequence(new byte[][]{
                derOid(new int[]{1, 2, 840, 113549, 1, 1, 11}),
                derNull()
        });

        byte[] issuer = encodeDN(distinguishedName);
        byte[] subject = issuer;

        byte[] validity = derSequence(new byte[][]{
                derUtcTime(notBefore),
                derUtcTime(notAfter)
        });

        byte[] subjectPublicKeyInfo = publicKeyEncoded;

        return derSequence(new byte[][]{
                version, serialNumber, signatureAlgorithm,
                issuer, validity, subject, subjectPublicKeyInfo
        });
    }

    private static byte[] encodeDN(String dn) {
        String[] parts = dn.split(",");
        byte[][] rdnSets = new byte[parts.length][];
        for (int i = 0; i < parts.length; i++) {
            String part = parts[i].trim();
            int eq = part.indexOf('=');
            String key = part.substring(0, eq).trim();
            String value = part.substring(eq + 1).trim();
            rdnSets[i] = derSet(new byte[][]{
                    derSequence(new byte[][]{
                            derOidForAttributeType(key),
                            derUtf8String(value)
                    })
            });
        }
        return derSequence(rdnSets);
    }

    private static byte[] derOidForAttributeType(String type) {
        switch (type.toUpperCase()) {
            case "CN": return derOid(new int[]{2, 5, 4, 3});
            case "O":  return derOid(new int[]{2, 5, 4, 10});
            case "C":  return derOid(new int[]{2, 5, 4, 6});
            case "OU": return derOid(new int[]{2, 5, 4, 11});
            case "ST": return derOid(new int[]{2, 5, 4, 8});
            case "L":  return derOid(new int[]{2, 5, 4, 7});
            default:
                throw new IllegalArgumentException("Unsupported attribute type: " + type);
        }
    }

    private static byte[] derSequence(byte[][] items) {
        return derConstructed(0x30, items);
    }

    private static byte[] derSet(byte[][] items) {
        return derConstructed(0x31, items);
    }

    private static byte[] derConstructed(int tag, byte[][] items) {
        int totalLen = 0;
        for (byte[] item : items) totalLen += item.length;
        byte[] content = new byte[totalLen];
        int offset = 0;
        for (byte[] item : items) {
            System.arraycopy(item, 0, content, offset, item.length);
            offset += item.length;
        }
        return derWrap(tag, content);
    }

    private static byte[] derExplicitTag(int tagNumber, byte[] content) {
        return derWrap(0xA0 | tagNumber, content);
    }

    private static byte[] derInteger(BigInteger value) {
        byte[] encoded = value.toByteArray();
        return derWrap(0x02, encoded);
    }

    private static byte[] derBitString(byte[] content) {
        return derWrap(0x03, content);
    }

    private static byte[] derNull() {
        return new byte[]{0x05, 0x00};
    }

    private static byte[] derUtf8String(String s) {
        byte[] content = s.getBytes(java.nio.charset.StandardCharsets.UTF_8);
        return derWrap(0x0C, content);
    }

    @SuppressWarnings("deprecation")
    private static byte[] derUtcTime(Date date) {
        java.util.Calendar cal = java.util.Calendar.getInstance(java.util.TimeZone.getTimeZone("UTC"));
        cal.setTime(date);
        String s = String.format(java.util.Locale.US, "%02d%02d%02d%02d%02d%02dZ",
                cal.get(java.util.Calendar.YEAR) % 100,
                cal.get(java.util.Calendar.MONTH) + 1,
                cal.get(java.util.Calendar.DAY_OF_MONTH),
                cal.get(java.util.Calendar.HOUR_OF_DAY),
                cal.get(java.util.Calendar.MINUTE),
                cal.get(java.util.Calendar.SECOND));
        return derWrap(0x17, s.getBytes(java.nio.charset.StandardCharsets.US_ASCII));
    }

    private static byte[] derOid(int[] components) {
        java.io.ByteArrayOutputStream buf = new java.io.ByteArrayOutputStream();
        buf.write(40 * components[0] + components[1]);
        for (int i = 2; i < components.length; i++) {
            encodeOidComponent(buf, components[i]);
        }
        return derWrap(0x06, buf.toByteArray());
    }

    private static void encodeOidComponent(java.io.ByteArrayOutputStream buf, int value) {
        if (value < 128) {
            buf.write(value);
        } else {
            byte[] temp = new byte[5];
            int pos = temp.length;
            temp[--pos] = (byte) (value & 0x7F);
            value >>= 7;
            while (value > 0) {
                temp[--pos] = (byte) (0x80 | (value & 0x7F));
                value >>= 7;
            }
            buf.write(temp, pos, temp.length - pos);
        }
    }

    private static byte[] derWrap(int tag, byte[] content) {
        byte[] lenBytes = derLength(content.length);
        byte[] result = new byte[1 + lenBytes.length + content.length];
        result[0] = (byte) tag;
        System.arraycopy(lenBytes, 0, result, 1, lenBytes.length);
        System.arraycopy(content, 0, result, 1 + lenBytes.length, content.length);
        return result;
    }

    private static byte[] derLength(int length) {
        if (length < 128) {
            return new byte[]{(byte) length};
        } else if (length < 256) {
            return new byte[]{(byte) 0x81, (byte) length};
        } else if (length < 65536) {
            return new byte[]{(byte) 0x82, (byte) (length >> 8), (byte) length};
        } else {
            return new byte[]{(byte) 0x83, (byte) (length >> 16), (byte) (length >> 8), (byte) length};
        }
    }
}
