package app.closer.crypto

import java.util.Base64
import com.google.crypto.tink.KeysetHandle
import com.google.crypto.tink.aead.AesGcmKeyManager
import com.google.crypto.tink.subtle.AesGcmJce
import org.bouncycastle.crypto.generators.Argon2BytesGenerator
import org.bouncycastle.crypto.params.Argon2Parameters
import java.security.SecureRandom
import javax.inject.Inject
import javax.inject.Singleton

/**
 * Pure crypto helper -- no Firestore, no Android dependencies.
 * Handles: keyset generation, Argon2id KDF, couple-key wrap/unwrap, phrase generation.
 *
 * Argon2id params (m=46 MiB, t=3, p=1): ~2-3 s on a mid-range phone.
 * Run on Dispatchers.Default to avoid blocking the main thread.
 */
@Singleton
class RecoveryKeyManager @Inject constructor() {

    data class WrappedKey(
        val cipherB64: String,
        val saltB64: String,
        val params: String
    )

    fun generateRecoveryPhrase(): String {
        val random = SecureRandom()
        val indices = (0 until PHRASE_WORD_COUNT).map { random.nextInt(WORDLIST.size) }
        return indices.joinToString(" ") { WORDLIST[it] }
    }

    fun newCoupleKeyset(): KeysetHandle =
        KeysetHandle.generateNew(AesGcmKeyManager.aes256GcmTemplate())

    /**
     * Wraps [keyset] with Argon2id(passphrase, salt) using AES-256-GCM.
     * AAD is the fixed constant [WRAP_AAD] so the blob is portable across
     * invite -> couple transition without re-wrapping.
     */
    fun wrap(keyset: KeysetHandle, phrase: String): WrappedKey {
        val salt = ByteArray(SALT_BYTES).also { SecureRandom().nextBytes(it) }
        val wrapKey = deriveKey(phrase, salt)
        val keysetBytes = serializeKeyset(keyset)
        val cipherBytes = AesGcmJce(wrapKey).encrypt(keysetBytes, WRAP_AAD)
        return WrappedKey(
            cipherB64 = Base64.getEncoder().encodeToString(cipherBytes),
            saltB64 = Base64.getEncoder().encodeToString(salt),
            params = PARAMS_TAG
        )
    }

    /**
     * Unwraps a [WrappedKey] using [phrase]. Throws [javax.crypto.AEADBadTagException]
     * (wrapped as GeneralSecurityException) if the phrase is wrong.
     */
    fun unwrap(wrapped: WrappedKey, phrase: String): KeysetHandle {
        val salt = Base64.getDecoder().decode(wrapped.saltB64)
        val cipher = Base64.getDecoder().decode(wrapped.cipherB64)
        val wrapKey = deriveKey(phrase, salt)
        val keysetBytes = AesGcmJce(wrapKey).decrypt(cipher, WRAP_AAD)
        return deserializeKeyset(keysetBytes)
    }

    /** Serialize a KeysetHandle to bytes (cleartext JSON -- stored inside EncryptedSharedPreferences). */
    fun serializeKeyset(handle: KeysetHandle): ByteArray {
        val baos = java.io.ByteArrayOutputStream()
        com.google.crypto.tink.CleartextKeysetHandle.write(
            handle, com.google.crypto.tink.JsonKeysetWriter.withOutputStream(baos)
        )
        return baos.toByteArray()
    }

    fun deserializeKeyset(bytes: ByteArray): KeysetHandle =
        com.google.crypto.tink.CleartextKeysetHandle.read(
            com.google.crypto.tink.JsonKeysetReader.withBytes(bytes)
        )

    private fun deriveKey(phrase: String, salt: ByteArray): ByteArray {
        val params = Argon2Parameters.Builder(Argon2Parameters.ARGON2_id)
            .withSalt(salt)
            .withParallelism(ARGON2_PARALLELISM)
            .withMemoryAsKB(ARGON2_MEMORY_KB)
            .withIterations(ARGON2_ITERATIONS)
            .build()
        val gen = Argon2BytesGenerator()
        gen.init(params)
        val out = ByteArray(KEY_BYTES)
        gen.generateBytes(phrase.toCharArray(), out, 0, KEY_BYTES)
        return out
    }

    /**
     * Encrypts [phrase] with a key derived from [inviteCode] via Argon2id + AES-256-GCM.
     * Output format: base64(salt[16] || gcm_ciphertext). Argon2id makes offline brute-force
     * infeasible even for the 6-char code space (~30 bits entropy).
     */
    fun encryptPhraseWithCode(inviteCode: String, phrase: String): String {
        val salt = ByteArray(SALT_BYTES).also { SecureRandom().nextBytes(it) }
        val key = deriveKey(inviteCode, salt)
        val ct = AesGcmJce(key).encrypt(phrase.toByteArray(Charsets.UTF_8), PHRASE_CIPHER_AAD)
        return Base64.getEncoder().encodeToString(salt + ct)
    }

    /**
     * Decrypts a ciphertext produced by [encryptPhraseWithCode] using [inviteCode].
     * Throws if the code is wrong or the ciphertext is malformed.
     */
    fun decryptPhraseWithCode(inviteCode: String, ciphertextB64: String): String {
        val payload = Base64.getDecoder().decode(ciphertextB64)
        val salt = payload.copyOf(SALT_BYTES)
        val ct = payload.copyOfRange(SALT_BYTES, payload.size)
        val key = deriveKey(inviteCode, salt)
        return String(AesGcmJce(key).decrypt(ct, PHRASE_CIPHER_AAD), Charsets.UTF_8)
    }

    companion object {
        private const val PHRASE_WORD_COUNT = 10
        private const val SALT_BYTES = 16
        private const val KEY_BYTES = 32
        private const val ARGON2_MEMORY_KB = 46 * 1024
        private const val ARGON2_ITERATIONS = 3
        private const val ARGON2_PARALLELISM = 1
        private const val PARAMS_TAG = "argon2id;v=19;m=47104;t=3;p=1"
        private val WRAP_AAD = "closer_couple_key".toByteArray(Charsets.UTF_8)
        private val PHRASE_CIPHER_AAD = "closer_invite_phrase".toByteArray(Charsets.UTF_8)

        // 256-word list -> 10 words -> ~80 bits raw entropy; Argon2id makes brute-force infeasible.
        val WORDLIST = arrayOf(
            "able","acid","acre","aged","aide","also","army","atom",
            "baby","back","bake","ball","bank","barn","base","bath",
            "bead","beam","bear","beat","belt","best","bill","bird",
            "bite","blue","boat","bold","bone","book","boot","bore",
            "cage","cake","call","calm","camp","cape","card","care",
            "cart","cave","cell","cent","coal","coat","code","coin",
            "cold","come","cone","cord","core","corn","cost","crew",
            "dame","damp","dare","dark","dart","date","dead","deal",
            "dean","deer","dent","dice","disk","dish","dock","done",
            "dose","dove","down","draw","drip","drop","drum","dusk",
            "each","earn","east","edge","epic","even","exam","exit",
            "fact","fail","fair","fall","fame","farm","fast","fate",
            "feel","fill","film","find","fire","fish","five","flat",
            "flow","foam","food","foot","form","free","fuel","full",
            "gain","game","gate","gear","give","glad","glow","goal",
            "gold","good","gray","grow","gulf","gust","half","hall",
            "halt","hand","hard","harm","head","heat","help","high",
            "hill","hint","hold","hole","home","hook","hope","hour",
            "huge","hull","hunt","hurt","idea","idle","inch","iris",
            "jade","jail","jest","join","jury","just","keen","kept",
            "kind","king","knee","knew","know","land","lane","last",
            "late","lead","leaf","lean","left","less","life","like",
            "line","lion","list","live","load","lock","long","look",
            "loop","lord","lost","loud","love","made","mail","main",
            "make","mark","maze","mean","meet","mild","mind","mine",
            "miss","mode","more","most","move","much","must","name",
            "near","need","nest","news","next","nice","nine","node",
            "noon","norm","note","once","only","open","over","pack",
            "page","pain","pair","pale","park","part","past","path",
            "peak","pick","pine","plan","play","plus","poor","port",
            "post","pull","pure","race","rank","rate","read","real"
        )
    }
}
