package com.peachpasswords.vault

import org.bouncycastle.crypto.generators.HKDFBytesGenerator
import org.bouncycastle.crypto.digests.SHA256Digest
import org.bouncycastle.crypto.params.HKDFParameters
import java.math.BigInteger
import java.security.MessageDigest
import javax.crypto.spec.SecretKeySpec

private const val WORD_BITS = 11
private const val TOTAL_WORDS = 12

fun rootKeyToWords(rootKeyBytes: ByteArray): List<String> {
    require(rootKeyBytes.size == 32) { "root key must be 32 bytes" }
    val entropy16 = rootKeyBytes.copyOfRange(0, 16)
    val entropyHash = MessageDigest.getInstance("SHA-256").digest(entropy16)
    // Mask to a byte before the unsigned shift: a Byte >= 0x80 sign-extends to a
    // negative Int, so `toInt() ushr 4` would yield a 28-bit value, not the 4-bit
    // checksum. The browser reads a Uint8 (`recovery-words.ts:36`). Without the
    // mask, ~50% of root keys (SHA[0] >= 0x80) produce non-browser-parity words.
    val checksum = (entropyHash[0].toInt() and 0xff) ushr 4

    var bitsBig = BigInteger.ZERO
    for (i in 0 until 16) {
        bitsBig = bitsBig.shiftLeft(8).or(BigInteger.valueOf((rootKeyBytes[i].toInt() and 0xff).toLong()))
    }
    bitsBig = bitsBig.shiftLeft(4).or(BigInteger.valueOf(checksum.toLong()))

    val words = ArrayList<String>(TOTAL_WORDS)
    for (i in 0 until TOTAL_WORDS) {
        val shift = (TOTAL_WORDS - 1 - i) * WORD_BITS
        val index = bitsBig.shiftRight(shift).and(BigInteger.valueOf(0x7ffL)).toInt()
        words += WORD_LIST[index]
    }
    return words
}

fun wordsToRootKeyMaterial(words: List<String>): ByteArray {
    require(words.size == TOTAL_WORDS) { "exactly $TOTAL_WORDS words required" }

    val indices = IntArray(words.size)
    for (i in words.indices) {
        val word = words[i].trim().lowercase()
        val index = WORD_LIST.indexOf(word)
        require(index >= 0) { "unknown word at position ${i + 1}: \"${words[i]}\"" }
        indices[i] = index
    }

    var bitsBig = BigInteger.ZERO
    for (index in indices) {
        bitsBig = bitsBig.shiftLeft(WORD_BITS).or(BigInteger.valueOf(index.toLong()))
    }

    val entropy = ByteArray(16)
    for (i in 0 until 16) {
        val byteShift = 124 - i * 8
        entropy[i] = bitsBig.shiftRight(byteShift).and(BigInteger.valueOf(0xffL)).toByte()
    }

    val storedChecksum = bitsBig.and(BigInteger.valueOf(0xFL)).toInt()
    val entropyHash = MessageDigest.getInstance("SHA-256").digest(entropy)
    // See rootKeyToWords: mask to a byte before the unsigned shift (browser parity).
    val computedChecksum = (entropyHash[0].toInt() and 0xff) ushr 4
    require(storedChecksum == computedChecksum) { "invalid recovery words — checksum mismatch" }

    return ByteArray(32).also { it.fill(0); System.arraycopy(entropy, 0, it, 0, entropy.size) }
}

/**
 * Generate a random 32-byte root key and prove its words round-trip.
 *
 * The retry loop this used to carry existed only because the shipped word list
 * duplicated "abandon" at index 1, which made ~0.6% of keys undecodable. The
 * list is now the canonical 2048, so every key round-trips; the check stays as
 * an assertion because shipping a key whose words cannot be typed back is the
 * one failure this code must never allow.
 */
fun generateRecoverableRootKey(): ByteArray {
    val candidate = ByteArray(32).also(java.security.SecureRandom()::nextBytes)
    check(runCatching { wordsToRootKeyMaterial(rootKeyToWords(candidate)) }.isSuccess) {
        "generated a root key whose recovery words do not round-trip"
    }
    return candidate
}

/** Lowercase, unseparated hex of [bytes] — the exact ASCII the browser wraps in the escrow. */
fun rootKeyToHex(bytes: ByteArray): String = bytes.joinToString("") { "%02x".format(it) }

/** Inverse of [rootKeyToHex]; tolerant of upper/lower case. */
fun hexToRootKey(hex: String): ByteArray {
    require(hex.length % 2 == 0) { "odd-length hex" }
    return ByteArray(hex.length / 2) { i -> hex.substring(i * 2, i * 2 + 2).toInt(16).toByte() }
}

fun deriveRecoveryKey(wordMaterial: ByteArray, salt: ByteArray): SecretKeySpec {
    val generator = HKDFBytesGenerator(SHA256Digest())
    generator.init(HKDFParameters(wordMaterial, salt, "peach-recovery-v1".encodeToByteArray()))
    val out = ByteArray(32)
    generator.generateBytes(out, 0, out.size)
    return SecretKeySpec(out, "AES")
}

fun findInvalidWords(words: List<String>): List<Pair<Int, String>> {
    val invalid = ArrayList<Pair<Int, String>>()
    for (i in words.indices) {
        val w = words[i].trim().lowercase()
        if (w.isBlank() || WORD_LIST.indexOf(w) == -1) {
            invalid += i to words[i]
        }
    }
    return invalid
}

fun getWordSuggestions(prefix: String, maxResults: Int = 5): List<String> {
    val p = prefix.trim().lowercase()
    if (p.length < 2) return emptyList()
    return WORD_LIST.filter { it.startsWith(p) }.take(maxResults)
}

/**
 * BLAKE3 hash (hex) of the space-joined, trimmed, lowercased recovery words.
 * Stored in vault metadata so recovery-word entry can be verified without
 * persisting the words. Mirrors the browser's recovery-hash derivation.
 */
fun recoveryWordsHash(words: List<String>): String {
    val canonical = words.joinToString(" ") { it.trim().lowercase() }.encodeToByteArray()
    val digest = org.bouncycastle.crypto.digests.Blake3Digest(256)
    digest.update(canonical, 0, canonical.size)
    val out = ByteArray(32)
    digest.doFinal(out, 0)
    return out.joinToString("") { "%02x".format(it) }
}