package com.secureline.server;

import android.security.keystore.KeyGenParameterSpec;
import android.security.keystore.KeyProperties;
import android.util.Slog;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.cert.Certificate;
import java.security.spec.ECGenParameterSpec;
import java.util.ArrayList;
import java.util.List;

public class SecureLineDeviceKeyManager {

    private static final String TAG =
            "SecureLineDeviceKeyManager";

    private static final String DEVICE_KEY_ALIAS =
            "secureline_device_key";

    // ============================================================
    // existence
    // ============================================================

    public boolean deviceKeyExists() {

        try {

            KeyStore ks =
                    KeyStore.getInstance("AndroidKeyStore");

            ks.load(null);

            return ks.containsAlias(DEVICE_KEY_ALIAS);

        } catch (Throwable t) {

            Slog.e(TAG,
                    "KeyStore error",
                    t);

            return false;
        }
    }

    // ============================================================
    // delete
    // ============================================================

    public void deleteDeviceKey() {

        try {

            KeyStore ks =
                    KeyStore.getInstance("AndroidKeyStore");

            ks.load(null);

            ks.deleteEntry(DEVICE_KEY_ALIAS);

            Slog.w(TAG,
                    "Device key deleted");

        } catch (Throwable t) {

            Slog.e(TAG,
                    "deleteDeviceKey failed",
                    t);
        }
    }

    // ============================================================
    // generate attested key
    // ============================================================

    public AttestationResult generateKey(byte[] challenge)
            throws Exception {

        KeyPairGenerator kpg =
                KeyPairGenerator.getInstance(
                        KeyProperties.KEY_ALGORITHM_EC,
                        "AndroidKeyStore"
                );

        KeyGenParameterSpec.Builder builder =
                new KeyGenParameterSpec.Builder(
                        DEVICE_KEY_ALIAS,
                        KeyProperties.PURPOSE_SIGN
                )
                        .setAlgorithmParameterSpec(
                                new ECGenParameterSpec("secp256r1")
                        )
                        .setDigests(
                                KeyProperties.DIGEST_SHA256
                        )
                        .setUserAuthenticationRequired(false)
                        .setUnlockedDeviceRequired(false)
                        .setAttestationChallenge(challenge);

        try {

            builder.setIsStrongBoxBacked(true);

            kpg.initialize(builder.build());

        } catch (Throwable strongBoxError) {

            Slog.w(TAG,
                    "StrongBox unavailable → fallback TEE");

            builder.setIsStrongBoxBacked(false);

            kpg.initialize(builder.build());
        }

        kpg.generateKeyPair();

        return getAttestation();
    }

    // ============================================================
    // attestation
    // ============================================================

    public AttestationResult getAttestation()
            throws Exception {

        KeyStore ks =
                KeyStore.getInstance("AndroidKeyStore");

        ks.load(null);

        Certificate[] chain =
                ks.getCertificateChain(DEVICE_KEY_ALIAS);

        if (chain == null || chain.length == 0) {

            throw new SecurityException(
                    "Attestation chain missing"
            );
        }

        List<byte[]> certChain =
                new ArrayList<>();

        for (Certificate cert : chain) {

            certChain.add(
                    cert.getEncoded()
            );
        }

        PublicKey publicKey =
                ks.getCertificate(DEVICE_KEY_ALIAS)
                        .getPublicKey();

        return new AttestationResult(
                publicKey.getEncoded(),
                certChain
        );
    }

    // ============================================================
    // signing
    // ============================================================

    public byte[] sign(byte[] data)
            throws Exception {

        KeyStore ks =
                KeyStore.getInstance("AndroidKeyStore");

        ks.load(null);

        PrivateKey privateKey =
                (PrivateKey)
                        ks.getKey(DEVICE_KEY_ALIAS, null);

        Signature sig =
                Signature.getInstance("SHA256withECDSA");

        sig.initSign(privateKey);

        sig.update(data);

        return sig.sign();
    }

    // ============================================================
    // fingerprint
    // ============================================================

    public String getPublicKeyFingerprint() {

        try {

            PublicKey pk = getPublicKey();

            if (pk == null)
                return "null";

            java.security.MessageDigest md =
                    java.security.MessageDigest.getInstance(
                            "SHA-256"
                    );

            byte[] digest =
                    md.digest(pk.getEncoded());

            StringBuilder sb =
                    new StringBuilder();

            for (byte b : digest) {

                sb.append(
                        String.format("%02x", b)
                );
            }

            return sb.toString();

        } catch (Throwable t) {

            Slog.e(TAG,
                    "fingerprint error",
                    t);

            return "error";
        }
    }

    public PublicKey getPublicKey()
            throws Exception {

        KeyStore ks =
                KeyStore.getInstance("AndroidKeyStore");

        ks.load(null);

        Certificate cert =
                ks.getCertificate(DEVICE_KEY_ALIAS);

        return cert != null
                ? cert.getPublicKey()
                : null;
    }

    // ============================================================

    public static class AttestationResult {

        public final byte[] publicKey;

        public final List<byte[]> certificateChain;

        public AttestationResult(
                byte[] publicKey,
                List<byte[]> chain
        ) {

            this.publicKey = publicKey;
            this.certificateChain = chain;
        }
    }
}