package chat.vojo.proxy.core.crypto

import chat.vojo.proxy.core.SsCipher
import java.security.MessageDigest
import javax.crypto.Cipher
import javax.crypto.Mac
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec

/**
 * Shadowsocks AEAD primitives (SIP004).
 *
 * Master key derivation, per-session subkey derivation (HKDF-SHA1), and the
 * AEAD seal/open operations. Implemented entirely on platform JCA so we pull in
 * no extra crypto dependency. ChaCha20-Poly1305 requires API 28+ (our minSdk).
 */
object ShadowsocksCrypto {

    private val SUBKEY_INFO = "ss-subkey".toByteArray(Charsets.US_ASCII)

    /**
     * OpenSSL EVP_BytesToKey(MD5) — the legacy KDF Shadowsocks uses to turn a
     * password into the master key. No salt, single iteration per block.
     */
    fun deriveMasterKey(password: String, keyLen: Int): ByteArray {
        val pw = password.toByteArray(Charsets.UTF_8)
        val md5 = MessageDigest.getInstance("MD5")
        val out = ByteArray(keyLen)
        var generated = 0
        var prev = ByteArray(0)
        while (generated < keyLen) {
            md5.reset()
            md5.update(prev)
            md5.update(pw)
            prev = md5.digest()
            val n = minOf(prev.size, keyLen - generated)
            System.arraycopy(prev, 0, out, generated, n)
            generated += n
        }
        return out
    }

    /** HKDF-SHA1 expand+extract producing [outLen] bytes of session subkey. */
    fun hkdfSha1(key: ByteArray, salt: ByteArray, info: ByteArray, outLen: Int): ByteArray {
        val mac = Mac.getInstance("HmacSHA1")

        // Extract
        mac.init(SecretKeySpec(salt, "HmacSHA1"))
        val prk = mac.doFinal(key)

        // Expand
        val out = ByteArray(outLen)
        var t = ByteArray(0)
        var generated = 0
        var counter = 1
        while (generated < outLen) {
            mac.reset()
            mac.init(SecretKeySpec(prk, "HmacSHA1"))
            mac.update(t)
            mac.update(info)
            mac.update(counter.toByte())
            t = mac.doFinal()
            val n = minOf(t.size, outLen - generated)
            System.arraycopy(t, 0, out, generated, n)
            generated += n
            counter++
        }
        return out
    }

    fun sessionSubkey(masterKey: ByteArray, salt: ByteArray, keyLen: Int): ByteArray =
        hkdfSha1(masterKey, salt, SUBKEY_INFO, keyLen)

    /** Little-endian increment of the 12-byte AEAD nonce, in place. */
    fun incrementNonce(nonce: ByteArray) {
        var i = 0
        while (i < nonce.size) {
            val v = (nonce[i].toInt() and 0xff) + 1
            nonce[i] = v.toByte()
            if (v <= 0xff) return
            i++
        }
    }
}

/** Abstracts the two AEAD families behind one seal/open surface. */
sealed class Aead(val cipher: SsCipher) {
    val keySize: Int get() = cipher.keySize
    val saltSize: Int get() = cipher.saltSize
    val tagSize: Int get() = cipher.tagSize

    abstract fun seal(subkey: ByteArray, nonce: ByteArray, plaintext: ByteArray, off: Int, len: Int): ByteArray
    abstract fun open(subkey: ByteArray, nonce: ByteArray, ciphertext: ByteArray, off: Int, len: Int): ByteArray

    fun seal(subkey: ByteArray, nonce: ByteArray, plaintext: ByteArray): ByteArray =
        seal(subkey, nonce, plaintext, 0, plaintext.size)

    fun open(subkey: ByteArray, nonce: ByteArray, ciphertext: ByteArray): ByteArray =
        open(subkey, nonce, ciphertext, 0, ciphertext.size)

    companion object {
        fun of(cipher: SsCipher): Aead = when (cipher) {
            SsCipher.AES_128_GCM, SsCipher.AES_256_GCM -> AesGcmAead(cipher)
            SsCipher.CHACHA20_IETF_POLY1305 -> ChaChaAead(cipher)
        }
    }
}

private class AesGcmAead(cipher: SsCipher) : Aead(cipher) {
    // Cipher.getInstance is a provider lookup + allocation; the relay hot path
    // runs four AEAD ops per 16 KB of traffic, so cache one instance per thread.
    private val local = ThreadLocal.withInitial { Cipher.getInstance("AES/GCM/NoPadding") }

    override fun seal(subkey: ByteArray, nonce: ByteArray, plaintext: ByteArray, off: Int, len: Int): ByteArray {
        val c = local.get()!!
        c.init(Cipher.ENCRYPT_MODE, SecretKeySpec(subkey, "AES"), GCMParameterSpec(tagSize * 8, nonce))
        return c.doFinal(plaintext, off, len)
    }

    override fun open(subkey: ByteArray, nonce: ByteArray, ciphertext: ByteArray, off: Int, len: Int): ByteArray {
        val c = local.get()!!
        c.init(Cipher.DECRYPT_MODE, SecretKeySpec(subkey, "AES"), GCMParameterSpec(tagSize * 8, nonce))
        return c.doFinal(ciphertext, off, len)
    }
}

private class ChaChaAead(cipher: SsCipher) : Aead(cipher) {
    // "ChaCha20/Poly1305/NoPadding" is registered by Conscrypt since API 28;
    // the hyphenated alias only exists in newer Mainline Conscrypt builds.
    private val local = ThreadLocal.withInitial { Cipher.getInstance("ChaCha20/Poly1305/NoPadding") }

    override fun seal(subkey: ByteArray, nonce: ByteArray, plaintext: ByteArray, off: Int, len: Int): ByteArray {
        val c = local.get()!!
        c.init(Cipher.ENCRYPT_MODE, SecretKeySpec(subkey, "ChaCha20"), IvParameterSpec(nonce))
        return c.doFinal(plaintext, off, len)
    }

    override fun open(subkey: ByteArray, nonce: ByteArray, ciphertext: ByteArray, off: Int, len: Int): ByteArray {
        val c = local.get()!!
        c.init(Cipher.DECRYPT_MODE, SecretKeySpec(subkey, "ChaCha20"), IvParameterSpec(nonce))
        return c.doFinal(ciphertext, off, len)
    }
}
