ssh-mobile-client/app/src/main/java/com/kosbarokah/sshclient/ssh/SshEngine.kt

package com.kosbarokah.sshclient.ssh

import com.kosbarokah.sshclient.data.AuthMethod
import com.kosbarokah.sshclient.data.ServerProfile
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.delay
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import net.schmizz.sshj.Config
import net.schmizz.sshj.DefaultConfig
import net.schmizz.sshj.SSHClient
import net.schmizz.sshj.connection.channel.direct.Session
import net.schmizz.sshj.connection.channel.direct.SessionChannel
import net.schmizz.sshj.transport.verification.PromiscuousVerifier
import net.schmizz.sshj.userauth.keyprovider.OpenSSHKeyFile
import java.io.IOException

/**
 * High-level SSH session engine built on sshj.
 *
 * - Opens a PTY ("xterm-256color") interactive shell.
 * - Streams remote bytes to [onData] callback and forwards local typed bytes
 *   via [write].
 * - Supports password, keyboard-interactive and private-key auth.
 * - Reconnects with backoff and sends keep-alive pings.
 */
class SshEngine(
    private val profile: ServerProfile,
    private val password: String?,
    private val privateKey: String?,
    private val passphrase: String?,
    private val passwordPolicy: PasswordPolicy = PasswordPolicy.DENY_INTERACTIVE_PROMPT
) {

    enum class PasswordPolicy {
        /** Do not auto-answer interactive password/kbf prompts. */
        DENY_INTERACTIVE_PROMPT,

        /** Auto-answer keyboard-interactive password prompts from the same password. */
        ALLOW_INTERACTIVE_PROMPT
    }

    interface Listener {
        fun onConnected()
        fun onDisconnected(reason: String)
        fun onError(message: String)
    }

    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
    private val maxReconnectAttempts = 5
    private var reconnectJob: Job? = null

    @Volatile
    private var client: SSHClient? = null

    @Volatile
    private var session: Session? = null

    @Volatile
    private var running = false

    private var varListener: Listener? = null
    private var varOnData: ((ByteArray) -> Unit)? = null

    fun setListener(listener: Listener) { varListener = listener }
    fun setOnData(cb: (ByteArray) -> Unit) { varOnData = cb }

    fun start() {
        if (running) return
        running = true
        reconnectJob = scope.launch {
            var attempt = 0
            while (isActive && running) {
                try {
                    attempt++.also { if (attempt > 1) varListener?.onDisconnected("Koneksi putus. Mendapatkan ulang (percobaan $attempt)…") }
                    connectOnce()
                    attempt = 0
                    // connected — this loop continues only on disconnection
                } catch (e: Exception) {
                    val msg = e.message ?: "koneksi gagal"
                    if (!running) break
                    val willRetry = attempt < maxReconnectAttempts
                    if (willRetry) {
                        varListener?.onError("Koneksi gagal: $msg")
                        delay((attempt.coerceAtMost(4) * 1000L)) // backoff 1s,2s,3s,4s,4s
                    } else {
                        varListener?.onError("Gagal setelah $attempt percobaan: $msg")
                        running = false
                        break
                    }
                }
            }
        }
    }

    private suspend fun connectOnce() {
        withContext(Dispatchers.IO) {
            val ssl = newClient()
            client = ssl

            // auth
            authenticate(ssl)

            // start interactive shell with PTY
            session = ssl.startSession().also {
                it.allocatePTY("xterm-256color", 80, 24, 0, 0, HashMap())
                it.startShell()
            }

            varListener?.onConnected()
            running = true

            // read loop (blocks; runs on IO)
            val srIn = session!!.inputStream
            val buf = ByteArray(8192)
            while (running) {
                val n = try {
                    srIn.read(buf)
                } catch (e: IOException) {
                    if (!running) 0 else throw e
                }
                if (n < 0) throw IOException("SSH koneksi ditutup oleh server")
                if (n > 0) {
                    val chunk = ByteArray(n)
                    System.arraycopy(buf, 0, chunk, 0, n)
                    varOnData?.invoke(chunk)
                }
            }
        }
    }

    private fun newClient(): SSHClient {
        val config: Config = DefaultConfig()
        val c = SSHClient(config)
        c.useCompression()
        c.loadKnownHosts()
        // Validate host keys: in this v1 we accept any host key but warn.
        // TODO: real known-host pinning (store fingerprint per profile).
        c.addHostKeyVerifier(PromiscuousVerifier())
        val p = if (profile.port > 0) profile.port else 22
        c.connect(profile.host, p)
        c.timeout = 20000
        // Enable SSH keep-alive (client-side ping every 60s) so NAT/lazy routers
        // don't drop the session during idle.
        try {
            c.connection.keepAlive.setKeepAliveInterval(60)
        } catch (_: Exception) {
        }
        return c
    }

    private fun authenticate(c: SSHClient) {
        val user = profile.username
        when (profile.authMethod) {
            AuthMethod.PASSWORD -> {
                c.authPassword(user, password ?: "")
            }
            AuthMethod.PRIVATE_KEY -> {
                val kf = OpenSSHKeyFile()
                kf.init(privateKey, passphrase)
                c.authPublickey(user, kf)
            }
        }
    }

    /**
     * Send bytes to the remote (user typed keys). Called on the UI/main thread.
     */
    fun write(data: ByteArray) {
        val sess = session ?: return
        try {
            sess.outputStream.write(data)
            sess.outputStream.flush()
        } catch (_: IOException) {
            // stream closed; reconnect loop will handle
        }
    }

    fun writeUtf8(text: String) = write(text.toByteArray(Charsets.UTF_8))

    fun resize(cols: Int, rows: Int) {
        val s = session
        if (s != null && s.isOpen) {
            try {
                (s as? SessionChannel)?.changeWindowDimensions(cols, rows, 0, 0)
            } catch (_: Exception) {
            }
        }
    }

    fun disconnect(reason: String = "terputus") {
        running = false
        reconnectJob?.cancel()
        scope.cancel()
        try { session?.close() } catch (_: Exception) {}
        try { client?.disconnect() } catch (_: Exception) {}
        client = null
        session = null
        varListener?.onDisconnected(reason)
    }
}