package tech.xvanturing.freeproxy.vpn.proxy

import tech.xvanturing.freeproxy.data.model.ProxyProfile
import java.io.IOException
import java.io.InputStream
import java.net.DatagramSocket
import java.net.Inet6Address
import java.net.InetAddress
import java.net.InetSocketAddress
import java.net.Socket

/**
 * Minimal HTTP CONNECT upstream client used by DragonTCP Lite.
 *
 * The upstream proxy is always the local DragonTCP Go process on
 * 127.0.0.1:8080. DragonTCP then carries the stream over adaptive XOR-framed
 * TCP/53 to the remote server.
 */
class ProxyClient(
    private val profile: ProxyProfile,
    private val proxyAddress: InetSocketAddress,
    private val protector: SocketProtector,
) {
    @Throws(IOException::class)
    fun connectTcp(destination: InetAddress, destinationPort: Int): Socket {
        val socket = Socket()
        try {
            socket.bind(InetSocketAddress(0))
            if (!protector.protect(socket)) {
                throw IOException("Unable to protect local proxy socket from VPN")
            }
            socket.connect(proxyAddress, CONNECT_TIMEOUT_MS)
            socket.soTimeout = HANDSHAKE_TIMEOUT_MS
            socket.tcpNoDelay = true

            val literal = if (destination is Inet6Address) {
                "[${destination.hostAddress}]:$destinationPort"
            } else {
                "${destination.hostAddress}:$destinationPort"
            }
            val request = buildString {
                append("CONNECT ").append(literal).append(" HTTP/1.1\r\n")
                append("Host: ").append(literal).append("\r\n")
                append("Proxy-Connection: Keep-Alive\r\n")
                append("\r\n")
            }
            socket.getOutputStream().apply {
                write(request.toByteArray(Charsets.ISO_8859_1))
                flush()
            }

            val input = socket.getInputStream()
            val status = input.readLineCrLf()
            while (true) {
                val line = input.readLineCrLf()
                if (line.isEmpty()) break
            }
            val code = status.split(' ').getOrNull(1)?.toIntOrNull()
                ?: throw IOException("Local DragonTCP proxy returned invalid response: $status")
            if (code !in 200..299) {
                throw IOException("Local DragonTCP CONNECT failed: $status")
            }

            socket.soTimeout = 0
            return socket
        } catch (t: Throwable) {
            runCatching { socket.close() }
            throw t
        }
    }

    /** HTTP CONNECT does not support SOCKS5 UDP ASSOCIATE. */
    @Throws(IOException::class)
    fun openUdpAssociate(): UdpAssociation {
        throw IOException("UDP ASSOCIATE unavailable with local HTTP CONNECT proxy")
    }

    companion object {
        const val ATYP_IPV4 = 0x01
        const val ATYP_DOMAIN = 0x03
        const val ATYP_IPV6 = 0x04
        private const val CONNECT_TIMEOUT_MS = 10_000
        private const val HANDSHAKE_TIMEOUT_MS = 15_000
        private const val MAX_HEADER_LINE = 8192

        private fun InputStream.readLineCrLf(): String {
            val out = StringBuilder()
            while (true) {
                val b = read()
                if (b < 0) {
                    if (out.isEmpty()) throw IOException("Proxy closed connection during handshake")
                    break
                }
                if (b == '\n'.code) break
                if (b != '\r'.code) out.append(b.toChar())
                if (out.length > MAX_HEADER_LINE) throw IOException("Proxy response header too long")
            }
            return out.toString()
        }
    }
}

class UdpAssociation(
    private val controlSocket: Socket = Socket(),
    val relayAddress: InetSocketAddress = InetSocketAddress("127.0.0.1", 0),
) : AutoCloseable {
    val isAlive: Boolean get() = !controlSocket.isClosed && controlSocket.isConnected
    override fun close() { runCatching { controlSocket.close() } }
}

@Throws(IOException::class)
internal fun InputStream.readExactly(count: Int): ByteArray {
    val buffer = ByteArray(count)
    var offset = 0
    while (offset < count) {
        val n = read(buffer, offset, count - offset)
        if (n < 0) throw IOException("Connection closed with ${count - offset} bytes remaining")
        offset += n
    }
    return buffer
}
