package chat.vojo.proxy.core.crypto

import chat.vojo.proxy.core.net.readExact
import chat.vojo.proxy.core.net.u16be
import java.io.ByteArrayOutputStream
import java.io.EOFException
import java.io.IOException
import java.io.InputStream
import java.io.OutputStream
import java.security.SecureRandom

/** Max plaintext bytes per Shadowsocks AEAD chunk (14-bit length field). */
private const val MAX_PAYLOAD = 0x3FFF

private val rng = SecureRandom()

/**
 * Encrypts a byte stream into the Shadowsocks AEAD TCP framing:
 * `[salt][len-chunk][payload-chunk]...`, prepending the salt on the first write.
 */
class SsOutputStream(
    private val out: OutputStream,
    private val aead: Aead,
    masterKey: ByteArray,
) : OutputStream() {

    private val salt = ByteArray(aead.saltSize).also { rng.nextBytes(it) }
    private val subkey = ShadowsocksCrypto.sessionSubkey(masterKey, salt, aead.keySize)
    private val nonce = ByteArray(12)
    private var saltSent = false

    override fun write(b: Int) = write(byteArrayOf(b.toByte()), 0, 1)

    @Synchronized
    override fun write(b: ByteArray, off: Int, len: Int) {
        if (len <= 0) return
        val buf = ByteArrayOutputStream(len + 64)
        if (!saltSent) {
            buf.write(salt)
            saltSent = true
        }
        var p = off
        var rem = len
        while (rem > 0) {
            val n = minOf(rem, MAX_PAYLOAD)
            val lenBytes = byteArrayOf((n ushr 8).toByte(), n.toByte())
            buf.write(aead.seal(subkey, nonce, lenBytes))
            ShadowsocksCrypto.incrementNonce(nonce)
            buf.write(aead.seal(subkey, nonce, b, p, n))
            ShadowsocksCrypto.incrementNonce(nonce)
            p += n
            rem -= n
        }
        out.write(buf.toByteArray())
    }

    override fun flush() = out.flush()
    override fun close() = out.close()
}

/**
 * Decrypts a Shadowsocks AEAD TCP stream back into plaintext. Reads the salt
 * lazily on first access. A clean upstream close surfaces as normal EOF (-1).
 */
class SsInputStream(
    private val inp: InputStream,
    private val aead: Aead,
    private val masterKey: ByteArray,
) : InputStream() {

    private var subkey: ByteArray? = null
    private val nonce = ByteArray(12)
    private var plain: ByteArray = ByteArray(0)
    private var plainPos = 0

    private fun fill(): Boolean {
        try {
            if (subkey == null) {
                val salt = inp.readExact(aead.saltSize)
                subkey = ShadowsocksCrypto.sessionSubkey(masterKey, salt, aead.keySize)
            }
            val key = subkey!!
            val encLen = inp.readExact(2 + aead.tagSize)
            val lenPlain = aead.open(key, nonce, encLen)
            ShadowsocksCrypto.incrementNonce(nonce)
            val payloadLen = u16be(lenPlain[0].toInt(), lenPlain[1].toInt())
            if (payloadLen <= 0 || payloadLen > MAX_PAYLOAD) {
                throw IOException("invalid shadowsocks chunk length $payloadLen")
            }
            val encPayload = inp.readExact(payloadLen + aead.tagSize)
            plain = aead.open(key, nonce, encPayload)
            ShadowsocksCrypto.incrementNonce(nonce)
            plainPos = 0
            return true
        } catch (e: EOFException) {
            return false
        }
    }

    @Synchronized
    override fun read(b: ByteArray, off: Int, len: Int): Int {
        if (len <= 0) return 0
        while (plainPos >= plain.size) {
            if (!fill()) return -1
        }
        val n = minOf(len, plain.size - plainPos)
        System.arraycopy(plain, plainPos, b, off, n)
        plainPos += n
        return n
    }

    override fun read(): Int {
        val one = ByteArray(1)
        val n = read(one, 0, 1)
        return if (n <= 0) -1 else one[0].toInt() and 0xff
    }

    override fun close() = inp.close()
}

/** Seals a single Shadowsocks UDP datagram: `[salt][AEAD(payload, nonce=0)]`. */
fun ssUdpSeal(aead: Aead, masterKey: ByteArray, plaintext: ByteArray, off: Int, len: Int): ByteArray {
    val salt = ByteArray(aead.saltSize).also { rng.nextBytes(it) }
    val subkey = ShadowsocksCrypto.sessionSubkey(masterKey, salt, aead.keySize)
    val ct = aead.seal(subkey, ByteArray(12), plaintext, off, len)
    return salt + ct
}

/** Opens a single Shadowsocks UDP datagram; returns the decrypted `[ADDR][PAYLOAD]`. */
fun ssUdpOpen(aead: Aead, masterKey: ByteArray, packet: ByteArray, off: Int, len: Int): ByteArray {
    require(len >= aead.saltSize + aead.tagSize) { "short ss udp packet" }
    val salt = packet.copyOfRange(off, off + aead.saltSize)
    val subkey = ShadowsocksCrypto.sessionSubkey(masterKey, salt, aead.keySize)
    return aead.open(subkey, ByteArray(12), packet, off + aead.saltSize, len - aead.saltSize)
}
