package chat.vojo.proxy.core.socks

import android.util.Log
import chat.vojo.proxy.core.net.parseSocksAddress
import chat.vojo.proxy.core.net.readExact
import chat.vojo.proxy.core.net.readFully
import chat.vojo.proxy.core.net.readSocksAddress
import chat.vojo.proxy.core.net.u16be
import chat.vojo.proxy.core.upstream.UdpAssociation
import chat.vojo.proxy.core.upstream.Upstream
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.delay
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import java.io.BufferedInputStream
import java.io.Closeable
import java.io.InputStream
import java.io.OutputStream
import java.net.DatagramPacket
import java.net.DatagramSocket
import java.net.InetSocketAddress
import java.net.ServerSocket
import java.net.Socket
import java.net.SocketAddress
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicInteger

private const val TAG = "LocalSocks5"
private const val RELAY_BUFFER = 32 * 1024
private const val SOCKS_VERSION = 0x05
private const val CMD_CONNECT = 0x01
private const val CMD_UDP_ASSOCIATE = 0x03
private const val DNS_TCP_TIMEOUT_MS = 10_000

/**
 * A loopback SOCKS5 server the native tun2socks engine connects to. Each inbound
 * SOCKS5 connection is bridged to the configured [upstream] proxy. Supports TCP
 * CONNECT and UDP ASSOCIATE (with a DNS-over-TCP fallback for HTTP upstreams).
 */
class LocalSocks5Server(
    private val upstream: Upstream,
    private val onUpstreamError: (() -> Unit)? = null,
) {
    private val serverSocket = ServerSocket().apply {
        reuseAddress = true
        bind(InetSocketAddress("127.0.0.1", 0), 256)
    }

    /** Ephemeral loopback port the engine should target. */
    val port: Int get() = serverSocket.localPort

    // Elastic IO threads; blocking reads park their thread, so idle costs nothing.
    private val dispatcher = Dispatchers.IO.limitedParallelism(512)
    private val scope = CoroutineScope(SupervisorJob() + dispatcher)
    private val open = ConcurrentHashMap.newKeySet<Closeable>()

    fun start() {
        scope.launch { acceptLoop() }
    }

    fun stop() {
        runCatching { serverSocket.close() }
        open.forEach { runCatching { it.close() } }
        open.clear()
        scope.cancel()
    }

    private fun track(c: Closeable) { open.add(c) }
    private fun untrack(c: Closeable) { open.remove(c) }

    private suspend fun acceptLoop() {
        while (scope.isActive) {
            val client = try {
                serverSocket.accept()
            } catch (e: Exception) {
                // Transient failures (e.g. EMFILE) must not kill the listener;
                // only a closed server socket ends the loop.
                if (serverSocket.isClosed || !scope.isActive) break
                Log.w(TAG, "accept failed: ${e.message}")
                delay(100)
                continue
            }
            scope.launch { handleClient(client) }
        }
    }

    private suspend fun handleClient(client: Socket) {
        track(client)
        try {
            client.tcpNoDelay = true
            val cin = BufferedInputStream(client.getInputStream())
            val cout = client.getOutputStream()

            // Greeting: VER, NMETHODS, METHODS...
            if (cin.read() != SOCKS_VERSION) return
            val nMethods = cin.read()
            if (nMethods < 0) return
            cin.readFully(ByteArray(nMethods))
            cout.write(byteArrayOf(SOCKS_VERSION.toByte(), 0x00)) // no-auth
            cout.flush()

            // Request: VER, CMD, RSV, [ATYP+ADDR+PORT]
            val reqHead = cin.readExact(3)
            if (reqHead[0].toInt() != SOCKS_VERSION) return
            val cmd = reqHead[1].toInt() and 0xff
            val dest = try {
                readSocksAddress(cin)
            } catch (e: IllegalArgumentException) {
                runCatching { writeReply(cout, rep = 0x08) } // address type not supported
                return
            }

            when (cmd) {
                CMD_CONNECT -> handleConnect(client, cin, cout, dest)
                CMD_UDP_ASSOCIATE -> handleUdpAssociate(client, cin, cout)
                else -> writeReply(cout, rep = 0x07) // command not supported
            }
        } catch (e: Exception) {
            // fallthrough to cleanup
        } finally {
            closeQuietly(client)
            untrack(client)
        }
    }

    private suspend fun handleConnect(
        client: Socket,
        cin: InputStream,
        cout: OutputStream,
        dest: chat.vojo.proxy.core.net.Destination,
    ) {
        val conn = try {
            upstream.connectTcp(dest)
        } catch (e: Exception) {
            Log.w(TAG, "upstream connect failed for $dest: ${e.message}")
            onUpstreamError?.invoke()
            runCatching { writeReply(cout, rep = 0x05) } // connection refused
            return
        }
        track(conn)
        try {
            writeReply(cout, rep = 0x00) // success

            // Half-close aware relay: EOF in one direction only shuts down the
            // peer's write side (hev forwards SHUT_WR from apps); the connection
            // dies once both directions finish or anything errors out.
            val finished = AtomicInteger(0)
            val closeBoth = {
                closeQuietly(client)
                runCatching { conn.close() }
                Unit
            }
            val doneOne = { clean: Boolean, shutdownPeer: () -> Unit ->
                if (clean) {
                    // SSLSocket can't half-close — fall back to full close.
                    runCatching { shutdownPeer() }.onFailure { closeBoth() }
                    if (finished.incrementAndGet() >= 2) closeBoth()
                } else closeBoth()
            }
            val up = scope.launch {
                val clean = runCatching { pump(cin, conn.output) }.isSuccess
                doneOne(clean) { conn.shutdownOutput() }
            }
            val down = scope.launch {
                val clean = runCatching { pump(conn.input, cout) }.isSuccess
                doneOne(clean) { cout.flush(); client.shutdownOutput() }
            }
            up.join()
            down.join()
        } finally {
            runCatching { conn.close() }
            untrack(conn)
        }
    }

    private fun handleUdpAssociate(client: Socket, cin: InputStream, cout: OutputStream) {
        val relay = DatagramSocket(InetSocketAddress("127.0.0.1", 0))
        track(relay)
        val session = UdpSession(relay)
        try {
            // Tell the engine where to send UDP: 127.0.0.1:relayPort
            val rp = relay.localPort
            cout.write(
                byteArrayOf(
                    SOCKS_VERSION.toByte(), 0x00, 0x00, 0x01,
                    127, 0, 0, 1,
                    (rp ushr 8).toByte(), rp.toByte(),
                )
            )
            cout.flush()
            session.start()

            // Hold the association open until the engine closes the control connection.
            val drain = ByteArray(512)
            while (cin.read(drain) >= 0) { /* keep-alive */ }
        } catch (e: Exception) {
            // fallthrough to cleanup
        } finally {
            session.stop()
            closeQuietly(relay)
            untrack(relay)
        }
    }

    /** Bridges UDP datagrams between the engine's relay socket and the upstream proxy. */
    private inner class UdpSession(private val relay: DatagramSocket) {
        private val assoc: UdpAssociation? = runCatching { upstream.openUdpAssociation() }.getOrNull()
        private val jobs = mutableListOf<Job>()
        @Volatile private var engineSource: SocketAddress? = null
        @Volatile private var running = true

        fun start() {
            assoc?.let { track(it) }
            jobs += scope.launch { forwardLoop() }
            if (assoc != null) jobs += scope.launch { replyLoop() }
        }

        fun stop() {
            running = false
            assoc?.let { runCatching { it.close() }; untrack(it) }
            jobs.forEach { it.cancel() }
        }

        private fun forwardLoop() {
            val buf = ByteArray(64 * 1024)
            while (running) {
                val dp = DatagramPacket(buf, buf.size)
                try { relay.receive(dp) } catch (e: Exception) { break }
                if (engineSource == null) {
                    engineSource = dp.socketAddress
                    // Lock the relay to the engine's source so another local app
                    // can't inject datagrams or hijack the reply destination.
                    runCatching { relay.connect(dp.socketAddress) }
                } else {
                    engineSource = dp.socketAddress
                }
                try {
                    // [RSV(2)][FRAG(1)][ATYP+ADDR+PORT][DATA]
                    if (dp.length < 7 || buf[2].toInt() != 0) continue // FRAG unsupported
                    val (dest, headerLen) = parseSocksAddress(buf, 3)
                    val dataStart = 3 + headerLen
                    val dataLen = dp.length - dataStart
                    if (dataLen < 0) continue
                    if (assoc != null) {
                        assoc.send(dest, buf, dataStart, dataLen)
                    } else if (dest.port == 53) {
                        val q = buf.copyOfRange(dataStart, dp.length)
                        scope.launch { dnsOverTcp(dest, q) }
                    }
                } catch (e: Exception) {
                    // malformed datagram; skip
                }
            }
        }

        private fun replyLoop() {
            val a = assoc ?: return
            while (running) {
                val reply = try { a.receive() } catch (e: Exception) { break } ?: break
                sendToEngine(reply.from.raw, reply.data)
            }
        }

        /** HTTP-upstream DNS fallback: tunnel a single query as DNS-over-TCP (RFC 7766). */
        private fun dnsOverTcp(dest: chat.vojo.proxy.core.net.Destination, query: ByteArray) {
            val conn = try { upstream.connectTcp(dest) } catch (e: Exception) { return }
            track(conn)
            try {
                conn.socket.soTimeout = DNS_TCP_TIMEOUT_MS
                val out = conn.output
                out.write((query.size ushr 8) and 0xff)
                out.write(query.size and 0xff)
                out.write(query)
                out.flush()
                val lenHdr = conn.input.readExact(2)
                val rlen = u16be(lenHdr[0].toInt(), lenHdr[1].toInt())
                if (rlen <= 0 || rlen > 65535) return
                val resp = conn.input.readExact(rlen)
                sendToEngine(dest.raw, resp)
            } catch (e: Exception) {
                // drop on failure
            } finally {
                runCatching { conn.close() }
                untrack(conn)
            }
        }

        private fun sendToEngine(addrRaw: ByteArray, data: ByteArray) {
            val src = engineSource ?: return
            val pkt = ByteArray(3 + addrRaw.size + data.size)
            System.arraycopy(addrRaw, 0, pkt, 3, addrRaw.size)
            System.arraycopy(data, 0, pkt, 3 + addrRaw.size, data.size)
            runCatching { relay.send(DatagramPacket(pkt, pkt.size, src)) }
        }
    }

    private fun writeReply(out: OutputStream, rep: Int) {
        // VER, REP, RSV, ATYP=IPv4, BND.ADDR=0.0.0.0, BND.PORT=0
        out.write(byteArrayOf(SOCKS_VERSION.toByte(), rep.toByte(), 0x00, 0x01, 0, 0, 0, 0, 0, 0))
        out.flush()
    }

    private fun pump(input: InputStream, output: OutputStream) {
        val buf = ByteArray(RELAY_BUFFER)
        while (true) {
            val n = input.read(buf)
            if (n < 0) break
            if (n > 0) {
                output.write(buf, 0, n)
                output.flush()
            }
        }
    }

    private fun closeQuietly(c: Closeable) {
        runCatching { c.close() }
    }
}
