package com.rosetta.messenger.crypto

import android.util.Base64
import org.bouncycastle.crypto.engines.ChaCha7539Engine
import org.bouncycastle.crypto.macs.Poly1305
import org.bouncycastle.crypto.params.KeyParameter
import org.bouncycastle.crypto.params.ParametersWithIV
import org.bouncycastle.jce.ECNamedCurveTable
import org.bouncycastle.jce.provider.BouncyCastleProvider
import java.math.BigInteger
import java.security.MessageDigest
import java.security.SecureRandom
import java.security.Security
import javax.crypto.Cipher
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec

/**
 * Шифрование сообщений как в React Native версии
 * XChaCha20-Poly1305 для текста + ECDH + AES для ключа
 */
object MessageCrypto {
    
    private const val CHACHA_KEY_SIZE = 32
    private const val XCHACHA_NONCE_SIZE = 24
    private const val POLY1305_TAG_SIZE = 16
    
    init {
        if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
            Security.addProvider(BouncyCastleProvider())
        }
    }
    
    /**
     * Результат шифрования сообщения
     */
    data class EncryptedMessage(
        val ciphertext: String,   // Hex-encoded XChaCha20-Poly1305 ciphertext
        val key: String,          // Hex-encoded 32-byte key
        val nonce: String         // Hex-encoded 24-byte nonce
    )
    
    /**
     * XChaCha20-Poly1305 шифрование (совместимо с @noble/ciphers в RN)
     * 
     * XChaCha20-Poly1305 использует 24-byte nonce:
     * - Первые 16 байт nonce используются в HChaCha20 для derive subkey
     * - Последние 8 байт nonce используются как counter для ChaCha20
     */
    fun encryptMessage(plaintext: String): EncryptedMessage {
        val secureRandom = SecureRandom()
        
        // Генерируем случайный ключ (32 байта) и nonce (24 байта для XChaCha20)
        val key = ByteArray(CHACHA_KEY_SIZE)
        val nonce = ByteArray(XCHACHA_NONCE_SIZE)
        secureRandom.nextBytes(key)
        secureRandom.nextBytes(nonce)
        
        val ciphertext = xchacha20Poly1305Encrypt(plaintext.toByteArray(Charsets.UTF_8), key, nonce)
        
        return EncryptedMessage(
            ciphertext = ciphertext.toHex(),
            key = key.toHex(),
            nonce = nonce.toHex()
        )
    }
    
    /**
     * XChaCha20-Poly1305 encrypt implementation
     */
    private fun xchacha20Poly1305Encrypt(plaintext: ByteArray, key: ByteArray, nonce: ByteArray): ByteArray {
        // Step 1: Derive subkey using HChaCha20
        val subkey = hchacha20(key, nonce.copyOfRange(0, 16))
        
        // Step 2: Create ChaCha20 nonce (4 zeros + last 8 bytes of original nonce)
        val chacha20Nonce = ByteArray(12)
        // First 4 bytes are 0
        System.arraycopy(nonce, 16, chacha20Nonce, 4, 8)
        
        // Step 3: Encrypt with ChaCha20
        val engine = ChaCha7539Engine()
        engine.init(true, ParametersWithIV(KeyParameter(subkey), chacha20Nonce))
        
        val ciphertext = ByteArray(plaintext.size)
        engine.processBytes(plaintext, 0, plaintext.size, ciphertext, 0)
        
        // Step 4: Generate Poly1305 tag
        // For AEAD, we need to compute Poly1305 over AAD + ciphertext with proper padding
        val poly1305Key = ByteArray(32)
        val zeros = ByteArray(32)
        val polyKeyEngine = ChaCha7539Engine()
        polyKeyEngine.init(true, ParametersWithIV(KeyParameter(subkey), chacha20Nonce))
        polyKeyEngine.processBytes(zeros, 0, 32, poly1305Key, 0)
        
        val mac = Poly1305()
        mac.init(KeyParameter(poly1305Key))
        
        // No AAD in our case, just process ciphertext
        mac.update(ciphertext, 0, ciphertext.size)
        
        // Pad to 16 bytes
        val padding = (16 - (ciphertext.size % 16)) % 16
        if (padding > 0) {
            mac.update(ByteArray(padding), 0, padding)
        }
        
        // Length of AAD (0) as 64-bit LE
        mac.update(ByteArray(8), 0, 8)
        
        // Length of ciphertext as 64-bit LE
        val ctLen = longToLittleEndian(ciphertext.size.toLong())
        mac.update(ctLen, 0, 8)
        
        val tag = ByteArray(POLY1305_TAG_SIZE)
        mac.doFinal(tag, 0)
        
        // Return ciphertext + tag
        return ciphertext + tag
    }
    
    /**
     * HChaCha20 - derives a 256-bit subkey from a 256-bit key and 128-bit nonce
     */
    private fun hchacha20(key: ByteArray, nonce: ByteArray): ByteArray {
        val state = IntArray(16)
        
        // Constants "expand 32-byte k"
        state[0] = 0x61707865
        state[1] = 0x3320646e
        state[2] = 0x79622d32
        state[3] = 0x6b206574
        
        // Key
        for (i in 0..7) {
            state[4 + i] = littleEndianToInt(key, i * 4)
        }
        
        // Nonce
        for (i in 0..3) {
            state[12 + i] = littleEndianToInt(nonce, i * 4)
        }
        
        // 20 rounds
        for (i in 0 until 10) {
            // Column rounds
            quarterRound(state, 0, 4, 8, 12)
            quarterRound(state, 1, 5, 9, 13)
            quarterRound(state, 2, 6, 10, 14)
            quarterRound(state, 3, 7, 11, 15)
            // Diagonal rounds
            quarterRound(state, 0, 5, 10, 15)
            quarterRound(state, 1, 6, 11, 12)
            quarterRound(state, 2, 7, 8, 13)
            quarterRound(state, 3, 4, 9, 14)
        }
        
        // Output first 4 and last 4 words
        val subkey = ByteArray(32)
        intToLittleEndian(state[0], subkey, 0)
        intToLittleEndian(state[1], subkey, 4)
        intToLittleEndian(state[2], subkey, 8)
        intToLittleEndian(state[3], subkey, 12)
        intToLittleEndian(state[12], subkey, 16)
        intToLittleEndian(state[13], subkey, 20)
        intToLittleEndian(state[14], subkey, 24)
        intToLittleEndian(state[15], subkey, 28)
        
        return subkey
    }
    
    private fun quarterRound(state: IntArray, a: Int, b: Int, c: Int, d: Int) {
        state[a] += state[b]; state[d] = rotl(state[d] xor state[a], 16)
        state[c] += state[d]; state[b] = rotl(state[b] xor state[c], 12)
        state[a] += state[b]; state[d] = rotl(state[d] xor state[a], 8)
        state[c] += state[d]; state[b] = rotl(state[b] xor state[c], 7)
    }
    
    private fun rotl(v: Int, c: Int): Int = (v shl c) or (v ushr (32 - c))
    
    private fun littleEndianToInt(bs: ByteArray, off: Int): Int {
        return (bs[off].toInt() and 0xff) or
               ((bs[off + 1].toInt() and 0xff) shl 8) or
               ((bs[off + 2].toInt() and 0xff) shl 16) or
               ((bs[off + 3].toInt() and 0xff) shl 24)
    }
    
    private fun intToLittleEndian(n: Int, bs: ByteArray, off: Int) {
        bs[off] = n.toByte()
        bs[off + 1] = (n ushr 8).toByte()
        bs[off + 2] = (n ushr 16).toByte()
        bs[off + 3] = (n ushr 24).toByte()
    }
    
    private fun longToLittleEndian(n: Long): ByteArray {
        val bs = ByteArray(8)
        bs[0] = n.toByte()
        bs[1] = (n ushr 8).toByte()
        bs[2] = (n ushr 16).toByte()
        bs[3] = (n ushr 24).toByte()
        bs[4] = (n ushr 32).toByte()
        bs[5] = (n ushr 40).toByte()
        bs[6] = (n ushr 48).toByte()
        bs[7] = (n ushr 56).toByte()
        return bs
    }
    
    /**
     * Расшифровка текста сообщения (XChaCha20-Poly1305)
     */
    fun decryptMessage(ciphertext: String, keyHex: String, nonceHex: String): String {
        val key = keyHex.hexToBytes()
        val nonce = nonceHex.hexToBytes()
        val ciphertextWithTag = ciphertext.hexToBytes()
        
        val plaintext = xchacha20Poly1305Decrypt(ciphertextWithTag, key, nonce)
        return String(plaintext, Charsets.UTF_8)
    }
    
    /**
     * XChaCha20-Poly1305 decrypt implementation
     */
    private fun xchacha20Poly1305Decrypt(ciphertextWithTag: ByteArray, key: ByteArray, nonce: ByteArray): ByteArray {
        if (ciphertextWithTag.size < POLY1305_TAG_SIZE) {
            throw IllegalArgumentException("Ciphertext too short")
        }
        
        val ciphertext = ciphertextWithTag.copyOfRange(0, ciphertextWithTag.size - POLY1305_TAG_SIZE)
        val tag = ciphertextWithTag.copyOfRange(ciphertextWithTag.size - POLY1305_TAG_SIZE, ciphertextWithTag.size)
        
        // Step 1: Derive subkey using HChaCha20
        val subkey = hchacha20(key, nonce.copyOfRange(0, 16))
        
        // Step 2: Create ChaCha20 nonce
        val chacha20Nonce = ByteArray(12)
        System.arraycopy(nonce, 16, chacha20Nonce, 4, 8)
        
        // Step 3: Verify Poly1305 tag
        val poly1305Key = ByteArray(32)
        val zeros = ByteArray(32)
        val polyKeyEngine = ChaCha7539Engine()
        polyKeyEngine.init(true, ParametersWithIV(KeyParameter(subkey), chacha20Nonce))
        polyKeyEngine.processBytes(zeros, 0, 32, poly1305Key, 0)
        
        val mac = Poly1305()
        mac.init(KeyParameter(poly1305Key))
        mac.update(ciphertext, 0, ciphertext.size)
        
        val padding = (16 - (ciphertext.size % 16)) % 16
        if (padding > 0) {
            mac.update(ByteArray(padding), 0, padding)
        }
        
        mac.update(ByteArray(8), 0, 8)
        val ctLen = longToLittleEndian(ciphertext.size.toLong())
        mac.update(ctLen, 0, 8)
        
        val computedTag = ByteArray(POLY1305_TAG_SIZE)
        mac.doFinal(computedTag, 0)
        
        if (!tag.contentEquals(computedTag)) {
            throw SecurityException("Authentication failed")
        }
        
        // Step 4: Decrypt with ChaCha20
        val engine = ChaCha7539Engine()
        engine.init(false, ParametersWithIV(KeyParameter(subkey), chacha20Nonce))
        
        val plaintext = ByteArray(ciphertext.size)
        engine.processBytes(ciphertext, 0, ciphertext.size, plaintext, 0)
        
        return plaintext
    }
    
    /**
     * ECDH шифрование ключа для получателя
     * Использует secp256k1 + AES как в RN версии
     */
    fun encryptKeyForRecipient(keyAndNonce: ByteArray, recipientPublicKeyHex: String): String {
        val secureRandom = SecureRandom()
        val ecSpec = ECNamedCurveTable.getParameterSpec("secp256k1")
        
        // Генерируем эфемерный приватный ключ
        val ephemeralPrivateKeyBytes = ByteArray(32)
        secureRandom.nextBytes(ephemeralPrivateKeyBytes)
        val ephemeralPrivateKey = BigInteger(1, ephemeralPrivateKeyBytes)
        
        // Получаем эфемерный публичный ключ
        val ephemeralPublicKey = ecSpec.g.multiply(ephemeralPrivateKey)
        val ephemeralPublicKeyHex = ephemeralPublicKey.getEncoded(false).toHex()
        
        // Парсим публичный ключ получателя
        val recipientPublicKeyBytes = recipientPublicKeyHex.hexToBytes()
        val recipientPublicKey = ecSpec.curve.decodePoint(recipientPublicKeyBytes)
        
        // ECDH: получаем общий секрет
        val sharedPoint = recipientPublicKey.multiply(ephemeralPrivateKey)
        val sharedSecret = sharedPoint.normalize().xCoord.encoded
        
        // Derive AES key from shared secret
        val aesKey = MessageDigest.getInstance("SHA-256").digest(sharedSecret)
        
        // Генерируем IV для AES
        val iv = ByteArray(16)
        secureRandom.nextBytes(iv)
        
        // Шифруем keyAndNonce с AES-CBC
        val secretKey = SecretKeySpec(aesKey, "AES")
        val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding")
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, IvParameterSpec(iv))
        val encryptedKey = cipher.doFinal(keyAndNonce)
        
        // Формат: iv:ciphertext:ephemeralPublicKey (Base64)
        val result = ByteArray(iv.size + encryptedKey.size + recipientPublicKeyBytes.size)
        System.arraycopy(iv, 0, result, 0, iv.size)
        System.arraycopy(encryptedKey, 0, result, iv.size, encryptedKey.size)
        
        // Возвращаем как Base64 с ephemeral public key в конце
        return Base64.encodeToString(iv, Base64.NO_WRAP) + ":" +
               Base64.encodeToString(encryptedKey, Base64.NO_WRAP) + ":" +
               ephemeralPublicKeyHex
    }
    
    /**
     * ECDH расшифровка ключа
     */
    fun decryptKeyFromSender(encryptedKeyBase64: String, myPrivateKeyHex: String): ByteArray {
        val parts = encryptedKeyBase64.split(":")
        if (parts.size != 3) throw IllegalArgumentException("Invalid encrypted key format")
        
        val iv = Base64.decode(parts[0], Base64.NO_WRAP)
        val encryptedKey = Base64.decode(parts[1], Base64.NO_WRAP)
        val ephemeralPublicKeyHex = parts[2]
        
        val ecSpec = ECNamedCurveTable.getParameterSpec("secp256k1")
        
        // Парсим мой приватный ключ
        val myPrivateKey = BigInteger(myPrivateKeyHex, 16)
        
        // Парсим эфемерный публичный ключ отправителя
        val ephemeralPublicKeyBytes = ephemeralPublicKeyHex.hexToBytes()
        val ephemeralPublicKey = ecSpec.curve.decodePoint(ephemeralPublicKeyBytes)
        
        // ECDH: получаем общий секрет
        val sharedPoint = ephemeralPublicKey.multiply(myPrivateKey)
        val sharedSecret = sharedPoint.normalize().xCoord.encoded
        
        // Derive AES key from shared secret
        val aesKey = MessageDigest.getInstance("SHA-256").digest(sharedSecret)
        
        // Расшифровываем
        val secretKey = SecretKeySpec(aesKey, "AES")
        val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding")
        cipher.init(Cipher.DECRYPT_MODE, secretKey, IvParameterSpec(iv))
        
        return cipher.doFinal(encryptedKey)
    }
    
    /**
     * Полное шифрование сообщения для отправки
     */
    fun encryptForSending(plaintext: String, recipientPublicKey: String): Pair<String, String> {
        // 1. Шифруем текст
        val encrypted = encryptMessage(plaintext)
        
        // 2. Собираем key + nonce
        val keyAndNonce = encrypted.key.hexToBytes() + encrypted.nonce.hexToBytes()
        
        // 3. Шифруем ключ для получателя
        val encryptedKey = encryptKeyForRecipient(keyAndNonce, recipientPublicKey)
        
        return Pair(encrypted.ciphertext, encryptedKey)
    }
    
    /**
     * Полная расшифровка входящего сообщения
     */
    fun decryptIncoming(
        ciphertext: String,
        encryptedKey: String,
        myPrivateKey: String
    ): String {
        // 1. Расшифровываем ключ
        val keyAndNonce = decryptKeyFromSender(encryptedKey, myPrivateKey)
        
        // 2. Разделяем key и nonce
        val key = keyAndNonce.slice(0 until 32).toByteArray()
        val nonce = keyAndNonce.slice(32 until keyAndNonce.size).toByteArray()
        
        // 3. Расшифровываем сообщение
        return decryptMessage(ciphertext, key.toHex(), nonce.toHex())
    }
}

// Extension functions для конвертации
private fun ByteArray.toHex(): String = joinToString("") { "%02x".format(it) }

private fun String.hexToBytes(): ByteArray {
    check(length % 2 == 0) { "Hex string must have even length" }
    return chunked(2).map { it.toInt(16).toByte() }.toByteArray()
}