package com.bytelyst.platform

import javax.crypto.Cipher
import javax.crypto.KeyGenerator
import javax.crypto.SecretKey
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.SecretKeySpec
import java.security.SecureRandom

/**
 * Encrypted field structure — wire-compatible with @bytelyst/field-encrypt (TypeScript).
 * All byte arrays are hex-encoded strings for JSON serialization.
 */
data class BLEncryptedField(
    /** Sentinel — always true for encrypted fields. */
    val __encrypted: Boolean = true,
    /** Schema version (currently 1). */
    val v: Int = 1,
    /** Algorithm identifier. */
    val alg: String = "aes-256-gcm",
    /** Ciphertext (hex-encoded). */
    val ct: String,
    /** Initialization vector (hex-encoded, 12 bytes = 24 hex chars). */
    val iv: String,
    /** GCM authentication tag (hex-encoded, 16 bytes = 32 hex chars). */
    val tag: String,
    /** DEK identifier — identifies which key was used for encryption. */
    val dekId: String,
)

/**
 * AES-256-GCM field-level encryption.
 *
 * Produces [BLEncryptedField] objects that are wire-compatible with the
 * TypeScript `@bytelyst/field-encrypt` package. Backends and native clients
 * can encrypt/decrypt the same fields interchangeably.
 *
 * Usage:
 * ```kotlin
 * val key = BLFieldEncrypt.generateKey()
 * val encrypted = BLFieldEncrypt.encrypt("sensitive data", key, "dek_user1_notes")
 * val decrypted = BLFieldEncrypt.decrypt(encrypted, key)
 * ```
 */
object BLFieldEncrypt {

    private const val ALGORITHM = "AES/GCM/NoPadding"
    private const val KEY_SIZE_BYTES = 32
    private const val KEY_SIZE_BITS = KEY_SIZE_BYTES * 8
    private const val IV_SIZE_BYTES = 12
    private const val TAG_SIZE_BITS = 128

    private val secureRandom = SecureRandom()

    // ── Encrypt ─────────────────────────────────────────────

    /**
     * Encrypt a plaintext string with AES-256-GCM.
     *
     * @param plaintext UTF-8 string to encrypt.
     * @param key 32-byte AES secret key.
     * @param dekId DEK identifier stored in the output for key lookup on decrypt.
     * @param aad Optional additional authenticated data (e.g., "userId:context").
     * @return [BLEncryptedField] with hex-encoded ciphertext, IV, and tag.
     * @throws IllegalArgumentException if key size is wrong.
     */
    fun encrypt(
        plaintext: String,
        key: SecretKey,
        dekId: String,
        aad: String? = null,
    ): BLEncryptedField {
        require(key.encoded.size == KEY_SIZE_BYTES) {
            "AES-256-GCM requires a $KEY_SIZE_BYTES-byte key, got ${key.encoded.size}"
        }

        val iv = ByteArray(IV_SIZE_BYTES).also { secureRandom.nextBytes(it) }

        val cipher = Cipher.getInstance(ALGORITHM)
        val spec = GCMParameterSpec(TAG_SIZE_BITS, iv)
        cipher.init(Cipher.ENCRYPT_MODE, key, spec)

        if (aad != null) {
            cipher.updateAAD(aad.toByteArray(Charsets.UTF_8))
        }

        // GCM output = ciphertext || tag (last 16 bytes)
        val ciphertextWithTag = cipher.doFinal(plaintext.toByteArray(Charsets.UTF_8))
        val tagOffset = ciphertextWithTag.size - TAG_SIZE_BITS / 8
        val ct = ciphertextWithTag.copyOfRange(0, tagOffset)
        val tag = ciphertextWithTag.copyOfRange(tagOffset, ciphertextWithTag.size)

        return BLEncryptedField(
            ct = ct.toHexString(),
            iv = iv.toHexString(),
            tag = tag.toHexString(),
            dekId = dekId,
        )
    }

    // ── Decrypt ─────────────────────────────────────────────

    /**
     * Decrypt a [BLEncryptedField] back to plaintext.
     *
     * @param field Encrypted field object (from Cosmos DB or API response).
     * @param key 32-byte AES secret key (must match the key used to encrypt).
     * @param aad Optional AAD (must match the AAD used during encryption).
     * @return Decrypted UTF-8 string.
     * @throws javax.crypto.AEADBadTagException if authentication fails (tampered data or wrong key).
     */
    fun decrypt(
        field: BLEncryptedField,
        key: SecretKey,
        aad: String? = null,
    ): String {
        require(key.encoded.size == KEY_SIZE_BYTES) {
            "AES-256-GCM requires a $KEY_SIZE_BYTES-byte key, got ${key.encoded.size}"
        }

        val iv = field.iv.hexToByteArray()
        val ct = field.ct.hexToByteArray()
        val tag = field.tag.hexToByteArray()

        // GCM expects ciphertext || tag as input
        val ciphertextWithTag = ct + tag

        val cipher = Cipher.getInstance(ALGORITHM)
        val spec = GCMParameterSpec(TAG_SIZE_BITS, iv)
        cipher.init(Cipher.DECRYPT_MODE, key, spec)

        if (aad != null) {
            cipher.updateAAD(aad.toByteArray(Charsets.UTF_8))
        }

        val plaintextBytes = cipher.doFinal(ciphertextWithTag)
        return String(plaintextBytes, Charsets.UTF_8)
    }

    // ── Key Generation ──────────────────────────────────────

    /**
     * Generate a random 32-byte AES-256 secret key.
     */
    fun generateKey(): SecretKey {
        val keyGen = KeyGenerator.getInstance("AES")
        keyGen.init(KEY_SIZE_BITS, secureRandom)
        return keyGen.generateKey()
    }

    /**
     * Create a [SecretKey] from a hex-encoded string (64 hex chars = 32 bytes).
     */
    fun keyFromHex(hex: String): SecretKey {
        val bytes = hex.hexToByteArray()
        require(bytes.size == KEY_SIZE_BYTES) {
            "AES-256-GCM requires a $KEY_SIZE_BYTES-byte key, got ${bytes.size}"
        }
        return SecretKeySpec(bytes, "AES")
    }

    // ── Secure Store Key Derivation ────────────────────────

    /**
     * Get or create a persistent encryption key in [BLSecureStore].
     *
     * On first call, generates a random 32-byte AES-256 key and stores it
     * as a hex string in EncryptedSharedPreferences. On subsequent calls,
     * loads the existing key. This provides a stable per-device DEK for
     * client-side encryption without requiring the backend to provision keys.
     *
     * @param store The [BLSecureStore] instance for the current app.
     * @param account Storage key name (default: `"field_encrypt_dek"`).
     * @return A 32-byte [SecretKey] backed by secure storage.
     */
    fun getOrCreateKey(store: BLSecureStore, account: String = "field_encrypt_dek"): SecretKey {
        val existingHex = store.read(account)
        if (existingHex != null) {
            return keyFromHex(existingHex)
        }

        val newKey = generateKey()
        val hex = newKey.encoded.toHexString()
        store.save(account, hex)
        return newKey
    }

    /**
     * Load an existing encryption key from [BLSecureStore] without creating one.
     *
     * @param store The [BLSecureStore] instance.
     * @param account Storage key name (default: `"field_encrypt_dek"`).
     * @return The stored [SecretKey], or `null` if none exists.
     */
    fun loadKey(store: BLSecureStore, account: String = "field_encrypt_dek"): SecretKey? {
        val hex = store.read(account) ?: return null
        return try {
            keyFromHex(hex)
        } catch (_: IllegalArgumentException) {
            null
        }
    }

    /**
     * Delete the stored encryption key from [BLSecureStore].
     *
     * @param store The [BLSecureStore] instance.
     * @param account Storage key name.
     * @return `true` if the key was deleted.
     */
    fun deleteKey(store: BLSecureStore, account: String = "field_encrypt_dek"): Boolean {
        return store.delete(account)
    }

    // ── Type Guard ──────────────────────────────────────────

    /**
     * Check if a JSON-like map represents an encrypted field.
     * Compatible with the TypeScript `isEncryptedField()` type guard.
     */
    fun isEncrypted(value: Map<String, Any?>?): Boolean {
        if (value == null) return false
        return value["__encrypted"] == true &&
            value["v"] != null &&
            value["alg"] != null &&
            value["ct"] != null &&
            value["iv"] != null &&
            value["tag"] != null &&
            value["dekId"] != null
    }

    /**
     * Check if a [BLEncryptedField] is valid.
     */
    fun isEncrypted(field: BLEncryptedField?): Boolean {
        return field?.__encrypted == true
    }
}

// ── Hex Helpers ─────────────────────────────────────────────

/** Hex-encode a byte array to a lowercase string. */
fun ByteArray.toHexString(): String =
    joinToString("") { "%02x".format(it) }

/** Decode a hex-encoded string to a byte array. */
fun String.hexToByteArray(): ByteArray {
    require(length % 2 == 0) { "Hex string must have even length, got $length" }
    return ByteArray(length / 2) { i ->
        val offset = i * 2
        substring(offset, offset + 2).toInt(16).toByte()
    }
}
