package chat.vojo.proxy.core.net

import java.io.InputStream
import java.net.InetAddress

const val ATYP_IPV4 = 0x01
const val ATYP_DOMAIN = 0x03
const val ATYP_IPV6 = 0x04

/**
 * A connection target in SOCKS5 address form. [raw] is the exact wire encoding
 * `[ATYP][ADDR][PORT]`, reusable verbatim as a Shadowsocks address header or an
 * upstream SOCKS5 request body.
 */
class Destination(
    val host: String,
    val port: Int,
    val atyp: Int,
    val raw: ByteArray,
) {
    override fun toString(): String = if (atyp == ATYP_IPV6) "[$host]:$port" else "$host:$port"
}

/** Parses `[ATYP][ADDR][PORT]` from [input] (already positioned at ATYP). */
fun readSocksAddress(input: InputStream): Destination {
    val atyp = input.read()
    if (atyp < 0) throw java.io.EOFException("eof reading atyp")
    val host: String
    val addrBytes: ByteArray
    when (atyp) {
        ATYP_IPV4 -> {
            addrBytes = input.readExact(4)
            host = InetAddress.getByAddress(addrBytes).hostAddress!!
        }
        ATYP_DOMAIN -> {
            val len = input.read()
            if (len < 0) throw java.io.EOFException("eof reading domain len")
            addrBytes = ByteArray(1 + len)
            addrBytes[0] = len.toByte()
            input.readFully(addrBytes, 1, len)
            host = String(addrBytes, 1, len, Charsets.US_ASCII)
        }
        ATYP_IPV6 -> {
            addrBytes = input.readExact(16)
            host = InetAddress.getByAddress(addrBytes).hostAddress!!
        }
        else -> throw IllegalArgumentException("unknown atyp $atyp")
    }
    val portBytes = input.readExact(2)
    val port = u16be(portBytes[0].toInt(), portBytes[1].toInt())

    val raw = ByteArray(1 + addrBytes.size + 2)
    raw[0] = atyp.toByte()
    System.arraycopy(addrBytes, 0, raw, 1, addrBytes.size)
    raw[raw.size - 2] = portBytes[0]
    raw[raw.size - 1] = portBytes[1]
    return Destination(host, port, atyp, raw)
}

/** Parses a SOCKS5 address from a raw datagram payload starting at [off]; returns dest + header length. */
fun parseSocksAddress(data: ByteArray, off: Int): Pair<Destination, Int> {
    var p = off
    val atyp = data[p++].toInt() and 0xff
    val host: String
    when (atyp) {
        ATYP_IPV4 -> {
            host = InetAddress.getByAddress(data.copyOfRange(p, p + 4)).hostAddress!!
            p += 4
        }
        ATYP_DOMAIN -> {
            val len = data[p++].toInt() and 0xff
            host = String(data, p, len, Charsets.US_ASCII)
            p += len
        }
        ATYP_IPV6 -> {
            host = InetAddress.getByAddress(data.copyOfRange(p, p + 16)).hostAddress!!
            p += 16
        }
        else -> throw IllegalArgumentException("unknown atyp $atyp")
    }
    val port = u16be(data[p].toInt(), data[p + 1].toInt())
    p += 2
    val headerLen = p - off
    val raw = data.copyOfRange(off, off + headerLen)
    return Destination(host, port, atyp, raw) to headerLen
}
