package com.durian.tssparty.util

import org.bouncycastle.jcajce.provider.digest.Keccak
import org.bouncycastle.jcajce.provider.digest.RIPEMD160
import org.bouncycastle.jcajce.provider.digest.SHA256
import java.security.MessageDigest

/**
 * Utility functions for address derivation
 */
object AddressUtils {

    /**
     * Derive Kava address from compressed public key
     * Kava uses Bech32 with "kava" prefix
     */
    fun deriveKavaAddress(compressedPubKey: ByteArray): String {
        // 1. Decompress public key if compressed (33 bytes -> 65 bytes)
        val uncompressedPubKey = if (compressedPubKey.size == 33) {
            decompressPublicKey(compressedPubKey)
        } else {
            compressedPubKey
        }

        // 2. For Cosmos/Kava: SHA256 -> RIPEMD160
        val sha256 = SHA256.Digest().digest(compressedPubKey)
        val ripemd160 = RIPEMD160.Digest().digest(sha256)

        // 3. Bech32 encode with "kava" prefix
        return Bech32.encode("kava", convertBits(ripemd160, 8, 5, true))
    }

    /**
     * Check if address is in EVM format (0x...)
     */
    fun isEvmAddress(address: String): Boolean {
        return address.startsWith("0x") && address.length == 42
    }

    /**
     * Get EVM address - either returns the address if already EVM format,
     * or derives it from the public key
     */
    fun getEvmAddress(address: String, publicKeyBase64: String): String {
        return if (isEvmAddress(address)) {
            address
        } else {
            // Derive EVM address from public key
            val publicKeyBytes = android.util.Base64.decode(publicKeyBase64, android.util.Base64.NO_WRAP)
            deriveEvmAddress(publicKeyBytes)
        }
    }

    /**
     * Derive EVM address from public key (for Kava EVM compatibility)
     */
    fun deriveEvmAddress(compressedPubKey: ByteArray): String {
        // 1. Decompress if needed
        val uncompressedPubKey = if (compressedPubKey.size == 33) {
            decompressPublicKey(compressedPubKey)
        } else {
            compressedPubKey
        }

        // 2. Take last 64 bytes (remove 0x04 prefix)
        val pubKeyNoPrefix = if (uncompressedPubKey.size == 65) {
            uncompressedPubKey.sliceArray(1..64)
        } else {
            uncompressedPubKey
        }

        // 3. Keccak256 hash
        val keccak = Keccak.Digest256().digest(pubKeyNoPrefix)

        // 4. Take last 20 bytes
        val addressBytes = keccak.sliceArray(12..31)

        // 5. Hex encode with 0x prefix
        return "0x" + addressBytes.toHexString()
    }

    /**
     * Decompress a compressed secp256k1 public key
     */
    private fun decompressPublicKey(compressed: ByteArray): ByteArray {
        require(compressed.size == 33) { "Invalid compressed public key size" }

        val prefix = compressed[0].toInt() and 0xFF
        require(prefix == 0x02 || prefix == 0x03) { "Invalid compression prefix" }

        val x = compressed.sliceArray(1..32)

        // secp256k1 curve parameters
        val p = "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F".toBigInteger(16)
        val xBigInt = x.toBigInteger()

        // y² = x³ + 7 (mod p)
        val ySquared = (xBigInt.pow(3) + 7.toBigInteger()).mod(p)

        // Calculate y using modular square root
        var y = ySquared.modPow((p + 1.toBigInteger()) / 4.toBigInteger(), p)

        // Check parity
        val isOdd = prefix == 0x03
        if (y.testBit(0) != isOdd) {
            y = p - y
        }

        // Build uncompressed key: 0x04 || x || y
        val result = ByteArray(65)
        result[0] = 0x04
        val xBytes = x
        val yBytes = y.toByteArray32()
        System.arraycopy(xBytes, 0, result, 1, 32)
        System.arraycopy(yBytes, 0, result, 33, 32)

        return result
    }

    /**
     * Convert between bit groups for Bech32
     */
    private fun convertBits(data: ByteArray, fromBits: Int, toBits: Int, pad: Boolean): ByteArray {
        var acc = 0
        var bits = 0
        val result = mutableListOf<Byte>()
        val maxv = (1 shl toBits) - 1

        for (value in data) {
            val v = value.toInt() and 0xFF
            acc = (acc shl fromBits) or v
            bits += fromBits
            while (bits >= toBits) {
                bits -= toBits
                result.add(((acc shr bits) and maxv).toByte())
            }
        }

        if (pad && bits > 0) {
            result.add(((acc shl (toBits - bits)) and maxv).toByte())
        }

        return result.toByteArray()
    }

    private fun ByteArray.toHexString(): String = joinToString("") { "%02x".format(it) }

    private fun ByteArray.toBigInteger(): java.math.BigInteger {
        return java.math.BigInteger(1, this)
    }

    private fun java.math.BigInteger.toByteArray32(): ByteArray {
        val bytes = this.toByteArray()
        return when {
            bytes.size == 32 -> bytes
            bytes.size > 32 -> bytes.sliceArray((bytes.size - 32) until bytes.size)
            else -> ByteArray(32 - bytes.size) + bytes
        }
    }
}

/**
 * Bech32 encoding utilities
 */
object Bech32 {
    private const val CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
    private val GENERATOR = intArrayOf(0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3)

    fun encode(hrp: String, data: ByteArray): String {
        val combined = data.map { it.toInt() and 0xFF }.toIntArray()
        val checksum = createChecksum(hrp, combined)
        val result = StringBuilder(hrp).append("1")
        for (d in combined) result.append(CHARSET[d])
        for (d in checksum) result.append(CHARSET[d])
        return result.toString()
    }

    private fun polymod(values: IntArray): Int {
        var chk = 1
        for (v in values) {
            val top = chk shr 25
            chk = ((chk and 0x1ffffff) shl 5) xor v
            for (i in 0..4) {
                if ((top shr i) and 1 == 1) {
                    chk = chk xor GENERATOR[i]
                }
            }
        }
        return chk
    }

    private fun hrpExpand(hrp: String): IntArray {
        val result = IntArray(hrp.length * 2 + 1)
        for (i in hrp.indices) {
            result[i] = hrp[i].code shr 5
            result[i + hrp.length + 1] = hrp[i].code and 31
        }
        result[hrp.length] = 0
        return result
    }

    private fun createChecksum(hrp: String, data: IntArray): IntArray {
        val values = hrpExpand(hrp) + data + intArrayOf(0, 0, 0, 0, 0, 0)
        val polymod = polymod(values) xor 1
        return IntArray(6) { (polymod shr (5 * (5 - it))) and 31 }
    }
}
