package com.okithaki.autoapp;

import android.content.Context;
import android.content.SharedPreferences;
import android.os.Build;
import android.security.keystore.KeyGenParameterSpec;
import android.security.keystore.KeyInfo;
import android.security.keystore.KeyProperties;
import android.util.Base64;

import org.json.JSONObject;

import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.MGF1ParameterSpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.OAEPParameterSpec;
import javax.crypto.spec.PSource;

/**
 * OK ITHAKI V249389
 *
 * Per-device RSA private key remains non-exportable inside Android Keystore.
 * Each direct message uses a fresh AES-256-GCM content key. That content key is
 * wrapped once for the recipient and once for the sender's own device.
 */
public final class SecureMessageCryptoV249389 {
    private static final String KEYSTORE = "AndroidKeyStore";
    private static final String KEY_ALIAS = "ok_ithaki_secure_dm_rsa_v249389";
    private static final String ALGORITHM = "RSA-OAEP-256+A256GCM";
    private static final int MAX_PLAINTEXT_BYTES = 2_000_000;
    private static final SecureRandom RANDOM = new SecureRandom();

    private SecureMessageCryptoV249389() { }

    private static KeyStore keyStore() throws Exception {
        KeyStore store = KeyStore.getInstance(KEYSTORE);
        store.load(null);
        return store;
    }

    public static synchronized KeyPair ensureKeyPair() throws Exception {
        KeyStore store = keyStore();
        if (store.containsAlias(KEY_ALIAS)) {
            PrivateKey privateKey = (PrivateKey) store.getKey(KEY_ALIAS, null);
            PublicKey publicKey = store.getCertificate(KEY_ALIAS).getPublicKey();
            if (privateKey != null && publicKey != null) {
                return new KeyPair(publicKey, privateKey);
            }
            store.deleteEntry(KEY_ALIAS);
        }

        KeyPairGenerator generator = KeyPairGenerator.getInstance(
                KeyProperties.KEY_ALGORITHM_RSA,
                KEYSTORE
        );
        KeyGenParameterSpec spec = new KeyGenParameterSpec.Builder(
                KEY_ALIAS,
                KeyProperties.PURPOSE_ENCRYPT | KeyProperties.PURPOSE_DECRYPT
        )
                .setKeySize(2048)
                .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_RSA_OAEP)
                .setDigests(KeyProperties.DIGEST_SHA256, KeyProperties.DIGEST_SHA1)
                .setRandomizedEncryptionRequired(true)
                .build();
        generator.initialize(spec);
        return generator.generateKeyPair();
    }

    public static String publicKeyBase64() throws Exception {
        PublicKey key = ensureKeyPair().getPublic();
        return Base64.encodeToString(key.getEncoded(), Base64.NO_WRAP);
    }

    private static OAEPParameterSpec oaepSpec() {
        return new OAEPParameterSpec(
                "SHA-256",
                "MGF1",
                MGF1ParameterSpec.SHA1,
                PSource.PSpecified.DEFAULT
        );
    }

    private static PublicKey decodePublicKey(String base64) throws Exception {
        byte[] bytes = Base64.decode(base64 == null ? "" : base64, Base64.DEFAULT);
        X509EncodedKeySpec spec = new X509EncodedKeySpec(bytes);
        return KeyFactory.getInstance("RSA").generatePublic(spec);
    }

    private static byte[] wrapAesKey(SecretKey key, PublicKey publicKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey, oaepSpec());
        return cipher.doFinal(key.getEncoded());
    }

    private static SecretKey unwrapAesKey(byte[] wrapped, PrivateKey privateKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        cipher.init(Cipher.UNWRAP_MODE, privateKey, oaepSpec());
        return (SecretKey) cipher.unwrap(wrapped, "AES", Cipher.SECRET_KEY);
    }

    public static String encryptEnvelope(
            String plaintext,
            String recipientPublicKeyBase64,
            String aadJson
    ) throws Exception {
        byte[] plain = (plaintext == null ? "" : plaintext).getBytes(StandardCharsets.UTF_8);
        if (plain.length == 0) {
            throw new IllegalArgumentException("empty_secure_message");
        }
        if (plain.length > MAX_PLAINTEXT_BYTES) {
            throw new IllegalArgumentException("secure_message_too_large");
        }
        byte[] aad = (aadJson == null ? "{}" : aadJson).getBytes(StandardCharsets.UTF_8);

        KeyPair ownPair = ensureKeyPair();
        PublicKey peerPublicKey = decodePublicKey(recipientPublicKeyBase64);

        KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
        keyGenerator.init(256);
        SecretKey contentKey = keyGenerator.generateKey();

        byte[] nonce = new byte[12];
        RANDOM.nextBytes(nonce);
        Cipher aes = Cipher.getInstance("AES/GCM/NoPadding");
        aes.init(Cipher.ENCRYPT_MODE, contentKey, new GCMParameterSpec(128, nonce));
        aes.updateAAD(aad);
        byte[] ciphertext = aes.doFinal(plain);

        JSONObject envelope = new JSONObject();
        envelope.put("v", 1);
        envelope.put("alg", ALGORITHM);
        envelope.put("nonce", Base64.encodeToString(nonce, Base64.NO_WRAP));
        envelope.put("ciphertext", Base64.encodeToString(ciphertext, Base64.NO_WRAP));
        envelope.put("wrapped_peer", Base64.encodeToString(
                wrapAesKey(contentKey, peerPublicKey), Base64.NO_WRAP));
        envelope.put("wrapped_self", Base64.encodeToString(
                wrapAesKey(contentKey, ownPair.getPublic()), Base64.NO_WRAP));
        envelope.put("aad", aadJson == null ? "{}" : aadJson);
        return envelope.toString();
    }

    public static String decryptEnvelope(String envelopeJson) throws Exception {
        JSONObject envelope = new JSONObject(envelopeJson == null ? "{}" : envelopeJson);
        if (envelope.optInt("v", 0) != 1 || !ALGORITHM.equals(envelope.optString("alg", ""))) {
            throw new IllegalArgumentException("unsupported_secure_envelope");
        }
        byte[] nonce = Base64.decode(envelope.getString("nonce"), Base64.DEFAULT);
        byte[] ciphertext = Base64.decode(envelope.getString("ciphertext"), Base64.DEFAULT);
        byte[] aad = envelope.optString("aad", "{}").getBytes(StandardCharsets.UTF_8);
        PrivateKey privateKey = ensureKeyPair().getPrivate();

        SecretKey contentKey = null;
        String[] wrappedFields = new String[]{"wrapped_self", "wrapped_peer"};
        Exception last = null;
        for (String field : wrappedFields) {
            try {
                byte[] wrapped = Base64.decode(envelope.getString(field), Base64.DEFAULT);
                contentKey = unwrapAesKey(wrapped, privateKey);
                break;
            } catch (Exception error) {
                last = error;
            }
        }
        if (contentKey == null) {
            throw last == null ? new IllegalArgumentException("secure_key_unwrap_failed") : last;
        }

        Cipher aes = Cipher.getInstance("AES/GCM/NoPadding");
        aes.init(Cipher.DECRYPT_MODE, contentKey, new GCMParameterSpec(128, nonce));
        aes.updateAAD(aad);
        byte[] plain = aes.doFinal(ciphertext);
        return new String(plain, StandardCharsets.UTF_8);
    }

    public static JSONObject status(Context context) {
        JSONObject result = new JSONObject();
        try {
            KeyPair pair = ensureKeyPair();
            boolean hardwareBacked = false;
            String securityLevel = "software_or_unknown";
            try {
                KeyFactory factory = KeyFactory.getInstance(pair.getPrivate().getAlgorithm(), KEYSTORE);
                KeyInfo info = factory.getKeySpec(pair.getPrivate(), KeyInfo.class);
                if (Build.VERSION.SDK_INT >= 29) {
                    int level = info.getSecurityLevel();
                    hardwareBacked = level == KeyProperties.SECURITY_LEVEL_TRUSTED_ENVIRONMENT
                            || level == KeyProperties.SECURITY_LEVEL_STRONGBOX;
                    if (level == KeyProperties.SECURITY_LEVEL_STRONGBOX) {
                        securityLevel = "strongbox";
                    } else if (level == KeyProperties.SECURITY_LEVEL_TRUSTED_ENVIRONMENT) {
                        securityLevel = "trusted_environment";
                    }
                } else {
                    hardwareBacked = info.isInsideSecureHardware();
                    securityLevel = hardwareBacked ? "secure_hardware" : "software_or_unknown";
                }
            } catch (Exception ignored) { }
            result.put("ok", true);
            result.put("version", "V249389");
            result.put("algorithm", ALGORITHM);
            result.put("public_key_b64", Base64.encodeToString(pair.getPublic().getEncoded(), Base64.NO_WRAP));
            result.put("private_key_exportable", false);
            result.put("hardware_backed", hardwareBacked);
            result.put("security_level", securityLevel);
            result.put("android_keystore", true);
        } catch (Exception error) {
            try {
                result.put("ok", false);
                result.put("version", "V249389");
                result.put("error", String.valueOf(error.getMessage()));
            } catch (Exception ignored) { }
        }
        return result;
    }

    public static JSONObject getSettings(Context context) {
        SharedPreferences prefs = context.getSharedPreferences("ok_ithaki_prefs", Context.MODE_PRIVATE);
        JSONObject result = new JSONObject();
        try {
            result.put("auto_delete_seconds", prefs.getInt("secure_auto_delete_v249389", 86400));
            result.put("storage_mode", prefs.getString("secure_storage_mode_v249389", "encrypted_queue"));
            result.put("screenshot_block", prefs.getBoolean("secure_screenshot_block_v249389", true));
            result.put("recording_block", prefs.getBoolean("secure_recording_block_v249389", true));
            result.put("voice_mask_mode", prefs.getString("secure_voice_mask_mode_v249389", "light"));
            result.put("voice_mask_recorded_audio_active", false);
            result.put("voice_mask_live_calls_active", false);
        } catch (Exception ignored) { }
        return result;
    }

    public static JSONObject saveSettings(Context context, String json) {
        JSONObject input;
        JSONObject output = new JSONObject();
        try {
            input = new JSONObject(json == null ? "{}" : json);
            int ttl = input.optInt("auto_delete_seconds", 86400);
            if (!(ttl == 0 || ttl == 1 || ttl == 10 || ttl == 30 || ttl == 60
                    || ttl == 300 || ttl == 3600 || ttl == 86400 || ttl == 259200)) {
                ttl = 86400;
            }
            String storage = input.optString("storage_mode", "encrypted_queue");
            if (!"zero_online_only".equals(storage)) {
                storage = "encrypted_queue";
            }
            String mask = input.optString("voice_mask_mode", "light");
            if (!("none".equals(mask) || "light".equals(mask) || "deep".equals(mask)
                    || "high".equals(mask) || "robotic".equals(mask))) {
                mask = "light";
            }
            context.getSharedPreferences("ok_ithaki_prefs", Context.MODE_PRIVATE)
                    .edit()
                    .putInt("secure_auto_delete_v249389", ttl)
                    .putString("secure_storage_mode_v249389", storage)
                    .putBoolean("secure_screenshot_block_v249389", input.optBoolean("screenshot_block", true))
                    .putBoolean("secure_recording_block_v249389", input.optBoolean("recording_block", true))
                    .putString("secure_voice_mask_mode_v249389", mask)
                    .apply();
            output = getSettings(context);
            output.put("ok", true);
        } catch (Exception error) {
            try {
                output.put("ok", false);
                output.put("error", String.valueOf(error.getMessage()));
            } catch (Exception ignored) { }
        }
        return output;
    }
}
