package chat.vojo.proxy.core.upstream

import chat.vojo.proxy.core.ProxyConfig
import chat.vojo.proxy.core.ProxyType
import chat.vojo.proxy.core.net.ATYP_IPV6
import chat.vojo.proxy.core.net.Destination
import java.io.BufferedInputStream
import java.io.IOException
import java.io.InputStream
import java.net.Socket
import java.util.Base64
import javax.net.ssl.SSLSocket
import javax.net.ssl.SSLSocketFactory

/** Upstream that tunnels through an HTTP(S) proxy using the CONNECT method. */
class HttpUpstream(
    private val cfg: ProxyConfig,
    private val protector: SocketProtector,
) : Upstream {

    private val tls: Boolean get() = cfg.type == ProxyType.HTTPS

    override fun connectTcp(dest: Destination): UpstreamConn {
        var socket = protector.connectedSocket(cfg.host, cfg.port)
        try {
            socket.soTimeout = HANDSHAKE_TIMEOUT_MS
            if (tls) {
                val factory = SSLSocketFactory.getDefault() as SSLSocketFactory
                val ssl = factory.createSocket(socket, cfg.host, cfg.port, true) as SSLSocket
                ssl.soTimeout = HANDSHAKE_TIMEOUT_MS
                // Raw SSLSocket validates the chain but NOT the hostname; without
                // this a valid cert for any domain could impersonate the proxy.
                val params = ssl.sslParameters
                params.endpointIdentificationAlgorithm = "HTTPS"
                ssl.sslParameters = params
                ssl.startHandshake()
                socket = ssl
            }
            val out = socket.getOutputStream()
            val inp = BufferedInputStream(socket.getInputStream())

            val authority = if (dest.atyp == ATYP_IPV6) "[${dest.host}]:${dest.port}" else "${dest.host}:${dest.port}"
            val sb = StringBuilder()
            sb.append("CONNECT ").append(authority).append(" HTTP/1.1\r\n")
            sb.append("Host: ").append(authority).append("\r\n")
            if (cfg.username.isNotEmpty()) {
                val token = Base64.getEncoder()
                    .encodeToString("${cfg.username}:${cfg.password}".toByteArray(Charsets.UTF_8))
                sb.append("Proxy-Authorization: Basic ").append(token).append("\r\n")
            }
            sb.append("Proxy-Connection: keep-alive\r\n")
            sb.append("\r\n")
            out.write(sb.toString().toByteArray(Charsets.US_ASCII))
            out.flush()

            val status = readStatusLine(inp)
            if (status !in 200..299) throw IOException("http proxy CONNECT failed: $status")

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

    /** Reads the response head up to the blank line and returns the status code. */
    private fun readStatusLine(inp: InputStream): Int {
        val head = StringBuilder()
        var last3 = 0 // rolling match for "\r\n\r\n"
        while (true) {
            val c = inp.read()
            if (c < 0) throw IOException("eof during http proxy response")
            head.append(c.toChar())
            last3 = when {
                c == '\r'.code -> if (last3 == 2) 3 else 1
                c == '\n'.code -> if (last3 == 1) 2 else if (last3 == 3) 4 else 0
                else -> 0
            }
            if (last3 == 4) break // saw \r\n\r\n
            if (head.length > 16 * 1024) throw IOException("http proxy response head too large")
        }
        val firstLine = head.lineSequence().firstOrNull()?.trim().orEmpty()
        val parts = firstLine.split(' ')
        return parts.getOrNull(1)?.toIntOrNull() ?: throw IOException("malformed status line: $firstLine")
    }

    /** HTTP proxies can't carry UDP; the local server falls back to DNS-over-TCP. */
    override fun openUdpAssociation(): UdpAssociation? = null
}
