package chat.vojo.proxy.core.upstream

import chat.vojo.proxy.core.ProxyConfig
import chat.vojo.proxy.core.net.ATYP_DOMAIN
import chat.vojo.proxy.core.net.ATYP_IPV4
import chat.vojo.proxy.core.net.ATYP_IPV6
import chat.vojo.proxy.core.net.Destination
import chat.vojo.proxy.core.net.readExact
import chat.vojo.proxy.core.net.readFully
import java.io.IOException
import java.io.InputStream
import java.io.OutputStream
import java.net.DatagramPacket
import java.net.DatagramSocket
import java.net.InetAddress
import java.net.InetSocketAddress

/** Upstream that speaks standard SOCKS5 (RFC 1928 / 1929) to the proxy server. */
class Socks5Upstream(
    private val cfg: ProxyConfig,
    private val protector: SocketProtector,
) : Upstream {

    private val auth: Boolean get() = cfg.username.isNotEmpty()

    override fun connectTcp(dest: Destination): UpstreamConn {
        val socket = protector.connectedSocket(cfg.host, cfg.port)
        try {
            socket.soTimeout = HANDSHAKE_TIMEOUT_MS
            val out = socket.getOutputStream()
            val inp = socket.getInputStream()
            negotiate(inp, out)

            // CONNECT request: VER=5, CMD=1, RSV=0, then [ATYP][ADDR][PORT]
            val req = ByteArray(3 + dest.raw.size)
            req[0] = 0x05
            req[1] = 0x01
            req[2] = 0x00
            System.arraycopy(dest.raw, 0, req, 3, dest.raw.size)
            out.write(req)
            out.flush()
            readReply(inp)

            socket.soTimeout = 0
            return UpstreamConn(socket, inp, out)
        } catch (e: Throwable) {
            runCatching { socket.close() }
            throw e
        }
    }

    private fun negotiate(inp: InputStream, out: OutputStream) {
        if (auth) {
            out.write(byteArrayOf(0x05, 0x02, 0x00, 0x02))
        } else {
            out.write(byteArrayOf(0x05, 0x01, 0x00))
        }
        out.flush()
        val sel = inp.readExact(2)
        if (sel[0].toInt() != 0x05) throw IOException("bad socks5 version ${sel[0]}")
        when (sel[1].toInt() and 0xff) {
            0x00 -> {} // no auth
            0x02 -> userPassAuth(inp, out)
            0xFF -> throw IOException("socks5 proxy rejected auth methods")
            else -> throw IOException("unsupported socks5 method ${sel[1]}")
        }
    }

    private fun userPassAuth(inp: InputStream, out: OutputStream) {
        val u = cfg.username.toByteArray(Charsets.UTF_8)
        val p = cfg.password.toByteArray(Charsets.UTF_8)
        // RFC 1929 encodes ULEN/PLEN as a single byte each — refuse anything that
        // would overflow rather than silently truncate the length prefix.
        if (u.size > 255 || p.size > 255) throw IOException("socks5 credentials exceed 255 bytes")
        val msg = ByteArray(3 + u.size + p.size)
        msg[0] = 0x01
        msg[1] = u.size.toByte()
        System.arraycopy(u, 0, msg, 2, u.size)
        msg[2 + u.size] = p.size.toByte()
        System.arraycopy(p, 0, msg, 3 + u.size, p.size)
        out.write(msg)
        out.flush()
        val res = inp.readExact(2)
        if (res[1].toInt() != 0x00) throw IOException("socks5 auth failed")
    }

    private fun readReply(inp: InputStream) {
        val head = inp.readExact(4) // VER, REP, RSV, ATYP
        if (head[1].toInt() != 0x00) throw IOException("socks5 connect failed, rep=${head[1].toInt() and 0xff}")
        // consume BND.ADDR + BND.PORT
        when (head[3].toInt() and 0xff) {
            ATYP_IPV4 -> inp.readFully(ByteArray(4 + 2))
            ATYP_IPV6 -> inp.readFully(ByteArray(16 + 2))
            ATYP_DOMAIN -> {
                val len = inp.read()
                if (len < 0) throw IOException("eof reading bnd addr")
                inp.readFully(ByteArray(len + 2))
            }
            else -> throw IOException("bad atyp in socks5 reply")
        }
    }

    override fun openUdpAssociation(): UdpAssociation? = try {
        Socks5UdpAssociation(cfg, protector, ::negotiate)
    } catch (e: Throwable) {
        null
    }
}

/**
 * SOCKS5 UDP ASSOCIATE: a TCP control connection holds the association open while
 * datagrams flow over a side UDP socket to the relay endpoint the proxy returns.
 */
private class Socks5UdpAssociation(
    cfg: ProxyConfig,
    protector: SocketProtector,
    negotiate: (InputStream, OutputStream) -> Unit,
) : UdpAssociation {

    private val control = protector.connectedSocket(cfg.host, cfg.port)
    private val udp = DatagramSocket().also { protector.protect(it) }
    private val relay: InetSocketAddress

    init {
        try {
            control.soTimeout = HANDSHAKE_TIMEOUT_MS
            val out = control.getOutputStream()
            val inp = control.getInputStream()
            negotiate(inp, out)
            // UDP ASSOCIATE with wildcard endpoint
            out.write(byteArrayOf(0x05, 0x03, 0x00, ATYP_IPV4.toByte(), 0, 0, 0, 0, 0, 0))
            out.flush()
            val head = inp.readExact(4)
            if (head[1].toInt() != 0x00) throw IOException("udp associate failed rep=${head[1].toInt() and 0xff}")
            val addr: InetAddress
            when (head[3].toInt() and 0xff) {
                ATYP_IPV4 -> addr = InetAddress.getByAddress(inp.readExact(4))
                ATYP_IPV6 -> addr = InetAddress.getByAddress(inp.readExact(16))
                ATYP_DOMAIN -> {
                    val len = inp.read()
                    if (len < 0) throw IOException("eof reading bnd domain len")
                    addr = InetAddress.getByName(String(inp.readExact(len), Charsets.US_ASCII))
                }
                else -> throw IOException("bad atyp in udp associate reply")
            }
            val portBytes = inp.readExact(2)
            val port = ((portBytes[0].toInt() and 0xff) shl 8) or (portBytes[1].toInt() and 0xff)
            // Some proxies return 0.0.0.0 — fall back to the proxy host.
            val target = if (addr.isAnyLocalAddress) InetAddress.getByName(cfg.host) else addr
            relay = InetSocketAddress(target, port)
            control.soTimeout = 0
            watchControl()
        } catch (e: Throwable) {
            close()
            throw e
        }
    }

    /**
     * RFC 1928: the association lives only while the control connection does.
     * When the proxy drops it, closing the UDP socket unblocks receive() so the
     * session ends instead of silently black-holing datagrams.
     */
    private fun watchControl() {
        Thread({
            runCatching {
                val inp = control.getInputStream()
                val sink = ByteArray(64)
                while (inp.read(sink) >= 0) { /* drain */ }
            }
            close()
        }, "socks5-udp-control").apply { isDaemon = true; start() }
    }

    override fun send(dest: Destination, data: ByteArray, off: Int, len: Int) {
        // [RSV(2)=0][FRAG(1)=0][ATYP+ADDR+PORT][DATA]
        val pkt = ByteArray(3 + dest.raw.size + len)
        System.arraycopy(dest.raw, 0, pkt, 3, dest.raw.size)
        System.arraycopy(data, off, pkt, 3 + dest.raw.size, len)
        udp.send(DatagramPacket(pkt, pkt.size, relay))
    }

    override fun receive(): UdpReply? {
        val buf = ByteArray(64 * 1024)
        // Skip malformed/fragmented datagrams instead of letting one bad packet
        // kill the whole reply loop; only socket-level errors propagate.
        while (true) {
            val dp = DatagramPacket(buf, buf.size)
            udp.receive(dp)
            try {
                if (dp.length < 7 || buf[2].toInt() != 0) continue // FRAG unsupported
                val (dest, headerLen) = chat.vojo.proxy.core.net.parseSocksAddress(buf, 3)
                val dataStart = 3 + headerLen
                if (dataStart > dp.length) continue
                return UdpReply(dest, buf.copyOfRange(dataStart, dp.length))
            } catch (e: Exception) {
                continue
            }
        }
    }

    override fun close() {
        runCatching { udp.close() }
        runCatching { control.close() }
    }
}
