package com.sshproxy.vpn.ssh

import android.net.VpnService
import com.jcraft.jsch.JSch
import com.jcraft.jsch.Session
import com.sshproxy.vpn.data.SshAuthMethod
import com.sshproxy.vpn.data.TunnelConfig
import java.io.File
import java.net.InetSocketAddress
import java.net.Socket
import java.util.Properties

/**
 * SSH-сессия + локальный SOCKS5 (TCP через direct-tcpip, UDP через relay на сервер).
 */
class SshTunnelManager {
    private var jsch: JSch? = null
    private var session: Session? = null
    private var sshSocket: Socket? = null
    private var socksProxy: LocalSocks5Proxy? = null
    private var localSocksPort: Int = 0
    private var remoteSocksCredentials: SocksAuthCredentials = SocksAuthCredentials.NONE

    val isConnected: Boolean
        get() = session?.isConnected == true

    fun checkHealth(): Boolean {
        val activeSession = session ?: return false
        if (!activeSession.isConnected) {
            return false
        }
        return runCatching {
            activeSession.sendKeepAliveMsg()
            true
        }.getOrElse { false }
    }

    fun connect(config: TunnelConfig): Int {
        disconnect()

        val remoteHost = config.remoteSocksHost.trim().ifBlank { TunnelConfig.DEFAULT_REMOTE_SOCKS_HOST }
        remoteSocksCredentials = SocksAuthCredentials(
            username = config.remoteSocksUsername,
            password = config.remoteSocksPassword
        )

        val localJsch = JSch()
        when (config.authMethod) {
            SshAuthMethod.PUBLIC_KEY -> {
                val keyFile = File(config.privateKeyPath)
                require(keyFile.exists()) { "Файл ключа не найден: ${config.privateKeyPath}" }
                if (config.keyPassphrase.isNotBlank()) {
                    localJsch.addIdentity(
                        keyFile.absolutePath,
                        config.keyPassphrase.toByteArray(Charsets.UTF_8)
                    )
                } else {
                    localJsch.addIdentity(keyFile.absolutePath)
                }
            }
            SshAuthMethod.PASSWORD -> {
                require(config.password.isNotBlank()) { "Укажите пароль SSH" }
            }
        }

        val newSession = localJsch.getSession(config.username, config.host, config.port)
        if (config.authMethod == SshAuthMethod.PASSWORD) {
            newSession.setPassword(config.password)
        }

        val props = Properties()
        props["StrictHostKeyChecking"] = if (config.acceptUnknownHostKey) "no" else "yes"
        props["PreferredAuthentications"] = when (config.authMethod) {
            SshAuthMethod.PUBLIC_KEY -> "publickey"
            SshAuthMethod.PASSWORD -> "password"
        }
        newSession.setConfig(props)
        val keepAliveMs = config.keepAliveIntervalSec.coerceIn(
            TunnelConfig.MIN_KEEP_ALIVE_INTERVAL_SEC,
            TunnelConfig.MAX_KEEP_ALIVE_INTERVAL_SEC
        ) * 1000
        newSession.setServerAliveInterval(keepAliveMs)
        newSession.setServerAliveCountMax(SERVER_ALIVE_COUNT_MAX)

        newSession.setSocketFactory(object : com.jcraft.jsch.SocketFactory {
            override fun createSocket(host: String, port: Int): Socket {
                val socket = Socket()
                socket.bind(InetSocketAddress(0))
                socket.connect(InetSocketAddress(host, port), CONNECT_TIMEOUT_MS)
                sshSocket = socket
                return socket
            }

            override fun getInputStream(socket: Socket) = socket.getInputStream()
            override fun getOutputStream(socket: Socket) = socket.getOutputStream()
        })

        newSession.connect(30_000)

        val proxy = LocalSocks5Proxy(
            session = newSession,
            bindPort = config.socksPort,
            remoteUdpRelayHost = remoteHost,
            remoteUdpRelayPort = config.remoteSocksPort,
            remoteSocksCredentials = remoteSocksCredentials
        )
        val assignedPort = proxy.start()
        if (assignedPort <= 0) {
            proxy.stop()
            newSession.disconnect()
            throw IllegalStateException("Не удалось поднять локальный SOCKS на порту ${config.socksPort}")
        }

        if (!SocksHealthCheck.verifyLocalSocks(assignedPort)) {
            proxy.stop()
            newSession.disconnect()
            throw IllegalStateException("Локальный SOCKS не отвечает на порту $assignedPort")
        }

        jsch = localJsch
        session = newSession
        socksProxy = proxy
        localSocksPort = assignedPort
        return assignedPort
    }

    fun protectSshSocket(vpnService: VpnService): Boolean {
        val socket = sshSocket ?: return false
        if (!socket.isConnected) {
            return false
        }
        return runCatching { vpnService.protect(socket) }.getOrDefault(false)
    }

    fun disconnect() {
        runCatching { socksProxy?.stop() }
        runCatching { session?.disconnect() }
        session = null
        jsch = null
        sshSocket = null
        socksProxy = null
        localSocksPort = 0
        remoteSocksCredentials = SocksAuthCredentials.NONE
    }

    fun getLocalSocksPort(): Int = localSocksPort

    fun isUdpRelayAvailable(remoteHost: String, remotePort: Int): Boolean {
        return checkUdpRelay(remoteHost, remotePort).available
    }

    fun checkUdpRelay(
        remoteHost: String,
        remotePort: Int,
        credentials: SocksAuthCredentials = remoteSocksCredentials
    ): UdpRelayHealthCheck.Result {
        if (session == null || session?.isConnected != true) {
            return UdpRelayHealthCheck.Result(
                available = false,
                stage = UdpRelayHealthCheck.Stage.TCP_CONNECT,
                detail = "SSH disconnected"
            )
        }
        if (localSocksPort > 0) {
            val local = UdpRelayHealthCheck.verifyThroughLocalSocks(localSocksPort)
            if (local.available) {
                return local
            }
            if (local.stage != UdpRelayHealthCheck.Stage.TCP_CONNECT) {
                return local
            }
        }
        val host = remoteHost.trim().ifBlank { TunnelConfig.DEFAULT_REMOTE_SOCKS_HOST }
        return UdpRelayHealthCheck.verifyFull(session!!, host, remotePort, credentials)
    }

    fun checkUdpRelayDetailed(
        remoteHost: String,
        remotePort: Int,
        credentials: SocksAuthCredentials = remoteSocksCredentials
    ): UdpRelayHealthCheck.Result {
        return checkUdpRelay(remoteHost, remotePort, credentials)
    }

    companion object {
        private const val CONNECT_TIMEOUT_MS = 30_000
        private const val SERVER_ALIVE_COUNT_MAX = 3
    }
}
