package chat.vojo.proxy.core.upstream

import chat.vojo.proxy.core.ProxyConfig
import chat.vojo.proxy.core.ProxyType
import chat.vojo.proxy.core.net.Destination
import java.io.Closeable
import java.io.InputStream
import java.io.OutputStream
import java.net.DatagramSocket
import java.net.Socket

/** Excludes app-owned sockets from the VPN so upstream traffic doesn't loop back into the tun. */
interface SocketProtector {
    fun protect(socket: Socket): Boolean
    fun protect(socket: DatagramSocket): Boolean
}

/** A live TCP tunnel to a target, with the proxy framing already negotiated. */
class UpstreamConn(
    val socket: Socket,
    val input: InputStream,
    val output: OutputStream,
) : Closeable {
    /**
     * Half-closes the write side so the peer sees EOF while reads continue.
     * SSLSocket does not support shutdownOutput — throws, caller falls back to close().
     */
    fun shutdownOutput() {
        output.flush()
        socket.shutdownOutput()
    }

    override fun close() {
        runCatching { socket.close() }
    }
}

class UdpReply(val from: Destination, val data: ByteArray)

/** A UDP relay through the proxy that multiplexes all destinations over one path. */
interface UdpAssociation : Closeable {
    fun send(dest: Destination, data: ByteArray, off: Int, len: Int)

    /** Blocks for the next inbound datagram; returns null when the association closes. */
    fun receive(): UdpReply?
}

/** One configured proxy endpoint, able to open TCP tunnels (and maybe UDP) for targets. */
interface Upstream {
    /** Connects to [dest] through the proxy. Throws on handshake failure. */
    fun connectTcp(dest: Destination): UpstreamConn

    /** A UDP relay, or null if this proxy type can't carry UDP (HTTP/HTTPS). */
    fun openUdpAssociation(): UdpAssociation?

    companion object {
        fun create(cfg: ProxyConfig, protector: SocketProtector): Upstream = when (cfg.type) {
            ProxyType.SHADOWSOCKS -> ShadowsocksUpstream(cfg, protector)
            ProxyType.SOCKS5 -> Socks5Upstream(cfg, protector)
            ProxyType.HTTP, ProxyType.HTTPS -> HttpUpstream(cfg, protector)
        }
    }
}

const val CONNECT_TIMEOUT_MS = 10_000
const val HANDSHAKE_TIMEOUT_MS = 15_000

/** Creates an unconnected, VPN-protected TCP socket and connects it to the proxy server. */
internal fun SocketProtector.connectedSocket(host: String, port: Int): Socket {
    val socket = Socket()
    protect(socket)
    socket.tcpNoDelay = true
    socket.connect(java.net.InetSocketAddress(host, port), CONNECT_TIMEOUT_MS)
    return socket
}
