package com.sshproxy.vpn.ssh



import com.jcraft.jsch.ChannelDirectTCPIP

import com.jcraft.jsch.Session

import java.net.InetSocketAddress

import java.net.Socket



/**

 * Полная проверка UDP для звонков:

 * 1) TCP до прокси

 * 2) SOCKS5 auth

 * 3) CMD FWD_UDP (hev-socks5-server)

 * 4) UDP-in-TCP probe — DNS к 1.1.1.1:53 и ответ

 */

object UdpRelayHealthCheck {



    private const val DEFAULT_TIMEOUT_MS = 8_000

    private const val DNS_PROBE_HOST = "1.1.1.1"

    private const val DNS_PROBE_PORT = 53



    enum class Stage {

        TCP_CONNECT,

        SOCKS_AUTH,

        FWD_UDP,

        UDP_PROBE

    }



    data class Result(

        val available: Boolean,

        val stage: Stage,

        val detail: String? = null

    )



    /** Проверка через локальный SOCKS (тот же путь, что у tun2socks при звонке). */

    fun verifyThroughLocalSocks(
        localPort: Int,
        timeoutMs: Int = DEFAULT_TIMEOUT_MS
    ): Result {
        if (localPort !in 1..65535) {
            return Result(false, Stage.TCP_CONNECT, "invalid local port")
        }
        return runCatching {
            Socket().use { socket ->
                socket.connect(InetSocketAddress("127.0.0.1", localPort), timeoutMs)
                socket.soTimeout = timeoutMs
                runSocksUdpProbe(
                    socket.getInputStream(),
                    socket.getOutputStream(),
                    SocksAuthCredentials.NONE
                )
            }
        }.getOrElse { error ->
            Result(false, Stage.TCP_CONNECT, error.message)
        }
    }



    /** Проверка напрямую до UDP-прокси на SSH-сервере (direct-tcpip). */

    fun verifyFull(

        session: Session,

        host: String,

        port: Int,

        credentials: SocksAuthCredentials = SocksAuthCredentials.NONE,

        timeoutMs: Int = DEFAULT_TIMEOUT_MS

    ): Result {

        if (host.isBlank() || port !in 1..65535) {

            return Result(false, Stage.TCP_CONNECT, "invalid host or port")

        }



        return runCatching {

            val channel = session.openChannel("direct-tcpip") as ChannelDirectTCPIP

            channel.setHost(host)

            channel.setPort(port)

            try {

                channel.connect(timeoutMs)

                runSocksUdpProbe(channel.inputStream, channel.outputStream, credentials)

            } finally {

                runCatching { channel.disconnect() }

            }

        }.getOrElse { error ->

            Result(false, Stage.TCP_CONNECT, error.message)

        }

    }



    private fun runSocksUdpProbe(

        input: java.io.InputStream,

        output: java.io.OutputStream,

        credentials: SocksAuthCredentials

    ): Result {

        val auth = Socks5Protocol.performAuthHandshake(input, output, credentials)

        if (!auth.success) {

            return Result(false, Stage.SOCKS_AUTH, auth.detail ?: "SOCKS5 auth failed")

        }



        Socks5Protocol.writeFwdUdpRequest(output)

        if (!Socks5Protocol.readConnectResponse(input)) {

            return Result(false, Stage.FWD_UDP, "FWD_UDP rejected or not hev-socks5")

        }



        Socks5Protocol.writeUdpFrame(

            output = output,

            host = DNS_PROBE_HOST,

            port = DNS_PROBE_PORT,

            payload = Socks5Protocol.minimalDnsQuery()

        )



        val response = Socks5Protocol.readUdpFrame(input)

        if (response == null || response.isEmpty()) {

            return Result(false, Stage.UDP_PROBE, "no UDP response from server")

        }



        return Result(true, Stage.UDP_PROBE)

    }



    /** Быстрая проверка — только TCP до порта. */

    fun isTcpPortOpen(

        session: Session,

        host: String,

        port: Int,

        timeoutMs: Int = 5_000

    ): Boolean {

        if (host.isBlank() || port !in 1..65535) {

            return false

        }

        return runCatching {

            val channel = session.openChannel("direct-tcpip") as ChannelDirectTCPIP

            channel.setHost(host)

            channel.setPort(port)

            try {

                channel.connect(timeoutMs)

                true

            } finally {

                runCatching { channel.disconnect() }

            }

        }.getOrDefault(false)

    }

}

