package com.sshproxy.vpn.ssh

import java.io.InputStream
import java.io.OutputStream

data class SocksAuthCredentials(
    val username: String = "",
    val password: String = ""
) {
    val isConfigured: Boolean
        get() = username.isNotBlank() && password.isNotEmpty()

    companion object {
        val NONE = SocksAuthCredentials()
    }
}

internal object Socks5Protocol {
    const val VERSION: Int = 0x05
    const val AUTH_VERSION: Int = 0x01
    const val METHOD_NO_AUTH: Int = 0x00
    const val METHOD_USER: Int = 0x02
    const val METHOD_NO_ACCEPTABLE: Int = 0xFF
    const val CMD_FWD_UDP: Int = 0x05
    const val ATYP_IPV4: Int = 0x01
    const val ATYP_DOMAIN: Int = 0x03
    const val ATYP_IPV6: Int = 0x04
    const val REPLY_SUCCESS: Int = 0x00

    fun readFully(input: InputStream, buffer: ByteArray): Boolean {
        var offset = 0
        while (offset < buffer.size) {
            val read = input.read(buffer, offset, buffer.size - offset)
            if (read < 0) {
                return false
            }
            offset += read
        }
        return true
    }

    fun performAuthHandshake(
        input: InputStream,
        output: OutputStream,
        credentials: SocksAuthCredentials = SocksAuthCredentials.NONE
    ): AuthHandshakeResult {
        val greeting = if (credentials.isConfigured) {
            byteArrayOf(
                VERSION.toByte(),
                0x02,
                METHOD_NO_AUTH.toByte(),
                METHOD_USER.toByte()
            )
        } else {
            byteArrayOf(VERSION.toByte(), 0x01, METHOD_NO_AUTH.toByte())
        }
        output.write(greeting)
        output.flush()

        val reply = ByteArray(2)
        if (!readFully(input, reply)) {
            return AuthHandshakeResult(false, "no reply")
        }
        if (reply[0] != VERSION.toByte()) {
            return AuthHandshakeResult(false, "invalid version")
        }

        return when (reply[1].toInt() and 0xFF) {
            METHOD_NO_AUTH -> AuthHandshakeResult(true)
            METHOD_USER -> {
                if (!credentials.isConfigured) {
                    AuthHandshakeResult(false, "credentials required")
                } else {
                    performUsernamePasswordAuth(input, output, credentials)
                }
            }
            METHOD_NO_ACCEPTABLE -> AuthHandshakeResult(false, "credentials required")
            else -> AuthHandshakeResult(false, "unsupported method: 0x${reply[1].toInt() and 0xFF}")
        }
    }

    private fun performUsernamePasswordAuth(
        input: InputStream,
        output: OutputStream,
        credentials: SocksAuthCredentials
    ): AuthHandshakeResult {
        val userBytes = credentials.username.trim().toByteArray(Charsets.UTF_8)
        val passBytes = credentials.password.toByteArray(Charsets.UTF_8)
        if (userBytes.isEmpty() || userBytes.size > 255 || passBytes.size > 255) {
            return AuthHandshakeResult(false, "invalid credentials")
        }

        output.write(byteArrayOf(AUTH_VERSION.toByte(), userBytes.size.toByte()))
        output.write(userBytes)
        output.write(byteArrayOf(passBytes.size.toByte()))
        output.write(passBytes)
        output.flush()

        val reply = ByteArray(2)
        if (!readFully(input, reply)) {
            return AuthHandshakeResult(false, "no auth reply")
        }
        if (reply[0] != AUTH_VERSION.toByte() || reply[1] != 0x00.toByte()) {
            return AuthHandshakeResult(false, "username/password rejected")
        }
        return AuthHandshakeResult(true)
    }

    data class AuthHandshakeResult(
        val success: Boolean,
        val detail: String? = null
    )

    /** Запрос FWD_UDP как у hev-socks5-tunnel (IPv4 upstream 0.0.0.0:0). */
    fun writeFwdUdpRequest(output: OutputStream) {
        output.write(
            byteArrayOf(
                VERSION.toByte(),
                CMD_FWD_UDP.toByte(),
                0x00,
                ATYP_IPV4.toByte(),
                0x00,
                0x00,
                0x00,
                0x00,
                0x00,
                0x00
            )
        )
        output.flush()
    }

    fun readConnectResponse(input: InputStream): Boolean {
        val head = ByteArray(4)
        if (!readFully(input, head)) {
            return false
        }
        if (head[0] != VERSION.toByte() || head[1] != REPLY_SUCCESS.toByte()) {
            return false
        }
        val tailSize = when (head[3].toInt() and 0xFF) {
            ATYP_IPV4 -> 6
            ATYP_IPV6 -> 18
            ATYP_DOMAIN -> null
            else -> return false
        }
        if (tailSize != null) {
            return readFully(input, ByteArray(tailSize))
        }
        val lengthByte = ByteArray(1)
        if (!readFully(input, lengthByte)) {
            return false
        }
        val domainLen = lengthByte[0].toInt() and 0xFF
        return readFully(input, ByteArray(domainLen + 2))
    }

    fun writeUdpFrame(output: OutputStream, host: String, port: Int, payload: ByteArray) {
        val parts = encodeIpv4Address(host, port)
        val hdrlen = 3 + parts.size
        output.write(
            byteArrayOf(
                (payload.size shr 8).toByte(),
                (payload.size and 0xFF).toByte(),
                hdrlen.toByte()
            )
        )
        output.write(parts)
        output.write(payload)
        output.flush()
    }

    fun readUdpFrame(input: InputStream): ByteArray? {
        val head = ByteArray(3)
        if (!readFully(input, head)) {
            return null
        }
        val datlen = ((head[0].toInt() and 0xFF) shl 8) or (head[1].toInt() and 0xFF)
        val hdrlen = head[2].toInt() and 0xFF
        if (hdrlen < 5 || datlen <= 0 || datlen > 512) {
            return null
        }
        val addrLen = hdrlen - 3
        if (!readFully(input, ByteArray(addrLen))) {
            return null
        }
        val payload = ByteArray(datlen)
        return if (readFully(input, payload)) payload else null
    }

    /** Минимальный DNS A-запрос (один символ «a»). */
    fun minimalDnsQuery(): ByteArray = byteArrayOf(
        0x00, 0x01,
        0x01, 0x00,
        0x00, 0x01,
        0x00, 0x00,
        0x00, 0x00,
        0x00, 0x00,
        0x01, 0x61, 0x00,
        0x00, 0x01,
        0x00, 0x01
    )

    private fun encodeIpv4Address(host: String, port: Int): ByteArray {
        val octets = host.split('.').mapNotNull { it.toIntOrNull()?.coerceIn(0, 255) }
        require(octets.size == 4) { "Invalid IPv4: $host" }
        return byteArrayOf(
            ATYP_IPV4.toByte(),
            octets[0].toByte(),
            octets[1].toByte(),
            octets[2].toByte(),
            octets[3].toByte(),
            (port shr 8).toByte(),
            (port and 0xFF).toByte()
        )
    }
}
