package com.secureline.notes.crypto;

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

import java.nio.ByteBuffer;
import java.security.KeyStore;
import java.security.SecureRandom;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

public class CryptoManager {

    private static final String KEY_ALIAS = "secureline_notes_key";

    private static final int PBKDF2_ITERATIONS = 120000;
    private static final int KEY_SIZE = 256;

    private static final int IV_SIZE = 12;
    private static final int SALT_SIZE = 16;
    private static final int VERSION_SIZE = 4;

    private static SecretKey getDeviceKey() throws Exception {

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

        SecretKey existing = (SecretKey) ks.getKey(KEY_ALIAS, null);
        if (existing != null) return existing;

        KeyGenerator generator = KeyGenerator.getInstance(
                KeyProperties.KEY_ALGORITHM_AES,
                "AndroidKeyStore"
        );

        generator.init(
                new KeyGenParameterSpec.Builder(
                        KEY_ALIAS,
                        KeyProperties.PURPOSE_ENCRYPT | KeyProperties.PURPOSE_DECRYPT
                )
                        .setBlockModes(KeyProperties.BLOCK_MODE_GCM)
                        .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_NONE)
                        .setRandomizedEncryptionRequired(true)
                        .build()
        );

        return generator.generateKey();
    }

    public static byte[] encryptLocal(byte[] data) throws Exception {

        if (data == null || data.length == 0)
            throw new IllegalArgumentException("invalid data");

        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.ENCRYPT_MODE, getDeviceKey());

        byte[] iv = cipher.getIV();
        if (iv == null || iv.length != IV_SIZE)
            throw new IllegalStateException("invalid iv");

        byte[] encrypted = cipher.doFinal(data);

        ByteBuffer buffer = ByteBuffer.allocate(IV_SIZE + encrypted.length);
        buffer.put(iv);
        buffer.put(encrypted);

        return buffer.array();
    }

    public static byte[] decryptLocal(byte[] data) throws Exception {

        if (data == null || data.length <= IV_SIZE)
            throw new IllegalArgumentException("invalid data");

        ByteBuffer buffer = ByteBuffer.wrap(data);

        byte[] iv = new byte[IV_SIZE];
        buffer.get(iv);

        int remaining = buffer.remaining();
        if (remaining <= 0)
            throw new IllegalArgumentException("invalid ciphertext");

        byte[] ciphertext = new byte[remaining];
        buffer.get(ciphertext);

        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(
                Cipher.DECRYPT_MODE,
                getDeviceKey(),
                new GCMParameterSpec(128, iv)
        );

        return cipher.doFinal(ciphertext);
    }

    public static byte[] encryptExport(byte[] data, String password) throws Exception {

        // ✅ FIX: password validation
        if (password == null || password.length() < 8)
            throw new IllegalArgumentException("weak password");

        if (data == null || data.length == 0)
            throw new IllegalArgumentException("invalid data");

        byte[] salt = new byte[SALT_SIZE];
        new SecureRandom().nextBytes(salt);

        SecretKey key = deriveKey(password, salt);

        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.ENCRYPT_MODE, key);

        byte[] iv = cipher.getIV();
        if (iv == null || iv.length != IV_SIZE)
            throw new IllegalStateException("invalid iv");

        byte[] encrypted = cipher.doFinal(data);

        ByteBuffer buffer = ByteBuffer.allocate(
                VERSION_SIZE + SALT_SIZE + IV_SIZE + encrypted.length
        );

        buffer.putInt(1);
        buffer.put(salt);
        buffer.put(iv);
        buffer.put(encrypted);

        for (int i = 0; i < salt.length; i++) salt[i] = 0;

        return buffer.array();
    }

    public static byte[] decryptExport(byte[] data, String password) throws Exception {

        // ✅ FIX: password validation
        if (password == null || password.length() < 8)
            throw new IllegalArgumentException("weak password");

        if (data == null || data.length < (VERSION_SIZE + SALT_SIZE + IV_SIZE))
            throw new IllegalArgumentException("invalid data");

        ByteBuffer buffer = ByteBuffer.wrap(data);

        int version = buffer.getInt();
        if (version != 1)
            throw new IllegalArgumentException("invalid format");

        // ✅ FIX: buffer safety
        if (buffer.remaining() < (SALT_SIZE + IV_SIZE))
            throw new IllegalArgumentException("invalid data");

        byte[] salt = new byte[SALT_SIZE];
        buffer.get(salt);

        byte[] iv = new byte[IV_SIZE];
        buffer.get(iv);

        int remaining = buffer.remaining();
        if (remaining <= 0)
            throw new IllegalArgumentException("invalid ciphertext");

        byte[] ciphertext = new byte[remaining];
        buffer.get(ciphertext);

        SecretKey key = deriveKey(password, salt);

        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(
                Cipher.DECRYPT_MODE,
                key,
                new GCMParameterSpec(128, iv)
        );

        byte[] result = cipher.doFinal(ciphertext);

        for (int i = 0; i < salt.length; i++) salt[i] = 0;

        return result;
    }

    private static SecretKey deriveKey(String password, byte[] salt) throws Exception {

        SecretKeyFactory factory =
                SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");

        PBEKeySpec spec = new PBEKeySpec(
                password.toCharArray(),
                salt,
                PBKDF2_ITERATIONS,
                KEY_SIZE
        );

        byte[] keyBytes = factory.generateSecret(spec).getEncoded();

        try {
            return new SecretKeySpec(keyBytes, "AES");
        } finally {
            spec.clearPassword();
            for (int i = 0; i < keyBytes.length; i++) {
                keyBytes[i] = 0;
            }
        }
    }
}