package wang.yaojia.webterm.session

import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Job
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.receiveAsFlow
import kotlinx.coroutines.launch
import kotlinx.coroutines.withTimeoutOrNull
import wang.yaojia.webterm.wire.ClientMessage
import wang.yaojia.webterm.wire.ConnectionPinger
import wang.yaojia.webterm.wire.HostEndpoint
import wang.yaojia.webterm.wire.MessageCodec
import wang.yaojia.webterm.wire.PingableTermTransport
import wang.yaojia.webterm.wire.ServerMessage
import wang.yaojia.webterm.wire.TermTransport
import wang.yaojia.webterm.wire.TimelineEvent
import wang.yaojia.webterm.wire.TransportConnection
import wang.yaojia.webterm.wire.Tunables

/**
 * The session lifecycle state machine (A14, plan §5 / §6.6). The Android analogue of the iOS
 * `SessionCore.SessionEngine`, and the composition root that drives the pure reducers
 * ([ReconnectMachine], [PingScheduler], [GateTracker], [AwayDigest]) against an injected
 * [TermTransport].
 *
 * PURE / TESTABLE (invariant #4 — the confinement contract): the engine imports NO OkHttp/Android
 * types. ALL engine state ([reconnect], [gateTracker], [currentConnection], [generation], …) is
 * touched ONLY on the injected [scope]'s dispatcher — a `Dispatchers.Default.limitedParallelism(1)`
 * in production, a `StandardTestDispatcher` under `runTest`. The only cross-boundary calls are the
 * event fan-out ([events], engine→UI) and [send]/[decideGate]/[notifyForegrounded]/[close]
 * (UI→engine), each of which hops onto [scope] before touching state. Because the dispatcher is
 * single-threaded the fields need no locks.
 *
 * The load-bearing invariants (plan §1):
 *  - **attach-first (#2):** [ClientMessage.Attach] is the FIRST frame sent on EVERY (re)connect,
 *    with the `sessionId` key always present (JSON `null` for a new session), sent BEFORE the
 *    connection is published so a queued UI send can never overtake it.
 *  - **adopt-server-id:** every `attached` frame adopts the server-issued id ([ReconnectMachine.withSessionId]),
 *    so the next reconnect re-attaches to the same PTY and the ring buffer replays.
 *  - **close == detach, NEVER kill (#2):** [close] closes the WS (server PTY survives) and stops
 *    reconnecting; it never sends a kill. A [close] that lands while a fresh connection is still
 *    inside the dial/attach window (before it is published to [currentConnection]) still detaches
 *    that connection — it is never leaked (see [runOneConnection]).
 *  - **generation-safe cancellation (defensive-only):** [generation] bumps on each (re)connect.
 *    Under the strictly-sequential supervisor no two connections ever run at once, so the
 *    `gen != generation` drop in [onFrame] is belt-and-braces — a superseded connection's late
 *    frame would be dropped rather than emitted onto a newer generation.
 *  - **oversized replay is terminal:** a frame past [maxFrameBytes] is the non-retryable
 *    [FailureReason.REPLAY_TOO_LARGE] — distinct from a retryable disconnect (never fed to backoff).
 */
public class SessionEngine(
    private val transport: TermTransport,
    private val endpoint: HostEndpoint,
    private val scope: CoroutineScope,
    initialSessionId: String? = null,
    private val initialCwd: String? = null,
    private val awayTimeline: suspend (sinceMs: Long) -> List<TimelineEvent> = { emptyList() },
    private val nowMs: () -> Long = { System.currentTimeMillis() },
    private val pingScheduler: PingScheduler = PingScheduler(),
    private val maxFrameBytes: Long = Tunables.MAX_WS_MESSAGE_BYTES,
    private val digestLimit: Int = DEFAULT_DIGEST_LIMIT,
) {
    /** A retryable end (drop/clean-server-close/ping-loss) vs a non-retryable [Terminal] one. */
    private sealed interface EndCause {
        /** The WS dropped while we were connected — reconnect with back-off. */
        data object Disconnected : EndCause

        /** `exit` or oversized-replay: the terminal event was already emitted; do NOT reconnect. */
        data object Terminal : EndCause
    }

    /** One live connection plus its optional ping capability (null = plain [TermTransport]). */
    private class OpenConn(val connection: TransportConnection, val pinger: ConnectionPinger?)

    private val eventsChannel = Channel<SessionEvent>(Channel.UNLIMITED)

    /** Engine→UI event stream (single consumer; the App's EventBus fans out per R10). */
    public val events: Flow<SessionEvent> = eventsChannel.receiveAsFlow()

    // ── engine-confined state (touched only on [scope]) ───────────────────────────────────
    private var reconnect: ReconnectMachine = ReconnectMachine.initial.withSessionId(initialSessionId)
    private var gateTracker: GateTracker = GateTracker.INITIAL
    private var currentConnection: TransportConnection? = null
    private var lastDims: Pair<Int, Int>? = null
    private var generation: Int = 0
    private var reconnectAttempt: Int = 0
    private var pendingDigest: Boolean = false
    private var disconnectedAtMs: Long? = null
    private var retryWakeup: CompletableDeferred<Unit>? = null
    private var started: Boolean = false
    private var closed: Boolean = false

    /** Begin connecting. Idempotent — a second call is a no-op. */
    public fun start() {
        scope.launch {
            // Idempotency flip happens ON the confined dispatcher (invariant #4): [started] is engine
            // state and must never be touched on the caller's thread.
            if (started) return@launch
            started = true
            try {
                supervisor()
            } finally {
                eventsChannel.close()
            }
        }
    }

    /**
     * Send a UI-originated frame (input / resize / approve / reject) on the live connection.
     * Dropped when disconnected (a `resize` is re-sent from [lastDims] on the next attach; input
     * while detached is meaningless). Encoding is byte-exact via [MessageCodec.encode].
     */
    public fun send(message: ClientMessage) {
        scope.launch { doSend(message) }
    }

    /**
     * Resolve a held gate ONLY if [epoch] still matches the current gate ([GateTracker.canDecide]);
     * a stale epoch is DROPPED so a slow tap can never approve the NEXT gate (two-line stale-guard).
     */
    public fun decideGate(epoch: Int, message: ClientMessage) {
        scope.launch { handleDecideGate(epoch, message) }
    }

    /**
     * The app returned to the foreground (plan §6.4/§6.6). If connected, re-send [lastDims] to
     * reclaim full-screen (latest-writer-wins). If in back-off, connect NOW without resetting the
     * ladder (no-reset-on-foreground, [ReconnectMachine]). [cols]/[rows] update the remembered dims.
     */
    public fun notifyForegrounded(cols: Int? = null, rows: Int? = null) {
        scope.launch { handleForegrounded(cols, rows) }
    }

    /**
     * Detach: close the WS (the server-side PTY keeps running) and stop reconnecting. NEVER a kill.
     *
     * Returns the launched [Job] so a caller can AWAIT the clean detach (the RFC-6455 `1000` close
     * frame handed to the transport writer) BEFORE it cancels the confinement [scope] (FIX 5 —
     * `RetainedSessionHolder.teardown()`). Fire-and-forget callers simply ignore the result.
     */
    public fun close(): Job = scope.launch { handleClose() }

    // ── supervisor: the connect / reconnect loop ──────────────────────────────────────────

    private suspend fun supervisor() {
        emit(Connection(ConnectionState.Connecting))
        while (!closed) {
            val gen = ++generation
            val cause = runOneConnection(gen)
            if (closed || cause == EndCause.Terminal) break
            scheduleAndWaitRetry()
        }
    }

    private suspend fun runOneConnection(gen: Int): EndCause {
        val open = dialOrNull() ?: return EndCause.Disconnected
        if (closed) return abandon(open) // close() landed during dial → detach, never publish
        if (!attachFirst(open)) return EndCause.Disconnected
        if (closed) return abandon(open) // close() landed during attach → detach, never publish
        currentConnection = open.connection
        onConnected()
        val cause = collectConnection(gen, open)
        currentConnection = null
        if (cause == EndCause.Disconnected) {
            // A live connection dropped → the next reattach owes an away-digest.
            pendingDigest = true
            disconnectedAtMs = nowMs()
        }
        return cause
    }

    /**
     * A [close] arrived while this connection was still opening (currentConnection was null, so
     * [handleClose] could not detach it). Detach it here and end terminally — emit NOTHING further
     * (no Connected/Adopted, no reconnect): [handleClose] already emitted [ConnectionState.Closed]
     * and [closed] short-circuits the supervisor. Upholds close == detach, NEVER kill (#2 / §6.6).
     */
    private suspend fun abandon(open: OpenConn): EndCause {
        ignoringNonCancel { open.connection.close() }
        return EndCause.Terminal
    }

    private suspend fun dialOrNull(): OpenConn? =
        try {
            if (transport is PingableTermTransport) {
                val pingable = transport.connectPingable(endpoint)
                OpenConn(pingable.connection, pingable.pinger)
            } else {
                OpenConn(transport.connect(endpoint), null)
            }
        } catch (e: CancellationException) {
            throw e
        } catch (_: Throwable) {
            null // connect-time failure → retryable
        }

    /** attach-first (+ replay [lastDims]) BEFORE publishing the connection. */
    private suspend fun attachFirst(open: OpenConn): Boolean =
        try {
            open.connection.send(MessageCodec.encode(ClientMessage.Attach(reconnect.sessionId, cwdForAttach())))
            lastDims?.let { (cols, rows) ->
                open.connection.send(MessageCodec.encode(ClientMessage.Resize(cols, rows)))
            }
            true
        } catch (e: CancellationException) {
            throw e
        } catch (_: Throwable) {
            ignoringNonCancel { open.connection.close() }
            false
        }

    /** `cwd` rides the attach only for a brand-new session (server appends it iff sessionId is null). */
    private fun cwdForAttach(): String? = if (reconnect.sessionId == null) initialCwd else null

    private fun onConnected() {
        reconnect = reconnect.reduce(ReconnectMachine.Input.Connected).first // reset ladder, keep id
        reconnectAttempt = 0
        emit(Connection(ConnectionState.Connected))
    }

    private suspend fun scheduleAndWaitRetry() {
        val (next, effect) = reconnect.reduce(ReconnectMachine.Input.Disconnected)
        reconnect = next
        reconnectAttempt += 1
        val after = (effect as ReconnectMachine.Effect.ScheduleRetry).after
        emit(Connection(ConnectionState.Reconnecting(reconnectAttempt, after)))
        val wake = CompletableDeferred<Unit>()
        retryWakeup = wake
        // Woken early by foreground/close → connect now; else the timer elapses → connect now.
        withTimeoutOrNull(after) { wake.await() }
        retryWakeup = null
    }

    // ── one connection's lifetime: frames + keep-alive ping ────────────────────────────────

    private suspend fun collectConnection(gen: Int, open: OpenConn): EndCause {
        var terminal: EndCause? = null
        return try {
            coroutineScope {
                val pingJob = open.pinger?.let { launchPing(open.connection, it) }
                try {
                    open.connection.frames.collect { frame ->
                        onFrame(gen, frame)?.let { cause ->
                            terminal = cause
                            ignoringNonCancel { open.connection.close() } // ends the frames flow → collect returns
                        }
                    }
                } finally {
                    pingJob?.cancel()
                }
                terminal ?: EndCause.Disconnected // clean finish with no exit frame = server-side close
            }
        } catch (e: CancellationException) {
            throw e
        } catch (_: Throwable) {
            EndCause.Disconnected // transport error → retryable
        }
    }

    private fun CoroutineScope.launchPing(conn: TransportConnection, pinger: ConnectionPinger): Job =
        launch {
            if (pingScheduler.run { pinger.ping() } == PingScheduler.Outcome.CONNECTION_LOST) {
                ignoringNonCancel { conn.close() } // declare dead → ends frames → collectConnection = Disconnected
            }
        }

    /** Handle one server frame; returns a terminal [EndCause] to stop the connection, else null. */
    private fun onFrame(gen: Int, frame: String): EndCause? {
        if (gen != generation) return null // defensive-only: the sequential supervisor never overlaps generations
        if (utf8Size(frame) > maxFrameBytes) {
            emit(Connection(ConnectionState.Failed(FailureReason.REPLAY_TOO_LARGE)))
            return EndCause.Terminal
        }
        return when (val msg = MessageCodec.decodeServer(frame)) {
            null -> null // undecodable → drop
            is ServerMessage.Attached -> { onAttached(msg.sessionId); null }
            is ServerMessage.Output -> { emit(Output(msg.data)); null }
            is ServerMessage.Status -> { onStatus(msg); null }
            is ServerMessage.Telemetry -> { emit(Telemetry(msg.telemetry)); null }
            is ServerMessage.Exit -> {
                emit(Exited(msg.code, msg.reason))
                EndCause.Terminal
            }
        }
    }

    private fun onAttached(sessionId: String) {
        reconnect = reconnect.withSessionId(sessionId) // always adopt the server-issued id
        emit(Adopted(sessionId))
        if (pendingDigest) {
            pendingDigest = false
            disconnectedAtMs?.let { since -> scope.launch { emitDigest(since) } }
        }
    }

    private fun onStatus(status: ServerMessage.Status) {
        val previous = gateTracker.current
        gateTracker = gateTracker.reduce(status.pending, status.gate, status.detail)
        val next = gateTracker.current
        if (next != previous) emit(Gate(next))
    }

    private suspend fun emitDigest(sinceMs: Long) {
        val timeline = try {
            awayTimeline(sinceMs)
        } catch (e: CancellationException) {
            throw e // never swallow cancellation (FIX 3)
        } catch (_: Throwable) {
            emptyList()
        }
        emit(Digest(AwayDigest.reduce(timeline, sinceMs, digestLimit))) // EMPTY still delivered; UI suppresses
    }

    // ── UI→engine command handlers (all on [scope]) ────────────────────────────────────────

    private suspend fun doSend(message: ClientMessage) {
        if (message is ClientMessage.Resize) lastDims = message.cols to message.rows
        currentConnection?.let { conn -> ignoringNonCancel { conn.send(MessageCodec.encode(message)) } }
    }

    private suspend fun handleDecideGate(epoch: Int, message: ClientMessage) {
        if (!gateTracker.canDecide(epoch)) return // stale epoch or no gate → drop (防误批新 gate)
        doSend(message)
        retireHeldGate() // a decision resolves the gate → a second same-epoch tap must NOT double-send
    }

    /**
     * Optimistically clear the held gate after a decision was sent (mirrors the web client setting
     * `pendingApprovalValue=false` inside `approve()`/`reject()`, public/terminal-session.ts). Folds a
     * synthetic falling edge through [GateTracker] so `canDecide(sameEpoch)` returns false afterwards;
     * [lastEpoch] is retained so the server's next rising edge still mints a fresh epoch. The server's
     * own next `pending=false` status is then a no-op (state already null → no duplicate [Gate] event).
     */
    private fun retireHeldGate() {
        if (gateTracker.current == null) return
        gateTracker = gateTracker.reduce(pending = false, gate = null, detail = null)
        emit(Gate(gateTracker.current)) // null → the UI hides the approve/reject affordances immediately
    }

    private suspend fun handleForegrounded(cols: Int?, rows: Int?) {
        if (closed) return
        if (cols != null && rows != null) lastDims = cols to rows
        val conn = currentConnection
        if (conn != null) {
            // Connected: reclaim full-screen by re-sending the remembered dims (latest-writer-wins).
            lastDims?.let { (c, r) -> ignoringNonCancel { conn.send(MessageCodec.encode(ClientMessage.Resize(c, r))) } }
        } else {
            // In back-off / connecting: connect NOW without touching the ladder (no-reset-on-foreground).
            retryWakeup?.complete(Unit)
        }
    }

    private suspend fun handleClose() {
        if (closed) return
        closed = true
        retryWakeup?.complete(Unit) // break any back-off wait so the supervisor exits
        val conn = currentConnection
        currentConnection = null
        conn?.let { ignoringNonCancel { it.close() } } // detach — the server PTY keeps running
        emit(Connection(ConnectionState.Closed))
    }

    private fun emit(event: SessionEvent) {
        eventsChannel.trySend(event) // UNLIMITED → never fails while open; no-op after close
    }

    /**
     * Run a best-effort suspend [block] (transport teardown / send), swallowing ordinary failures
     * but RETHROWING [CancellationException] so structured-concurrency cancellation is never
     * silently absorbed (FIX 3). Inline so the suspend calls run in the caller's coroutine.
     */
    private suspend inline fun ignoringNonCancel(block: () -> Unit) {
        try {
            block()
        } catch (e: CancellationException) {
            throw e
        } catch (_: Throwable) {
            // best-effort: a failed close/send is non-fatal to the engine's lifecycle
        }
    }

    /**
     * UTF-8 byte length of [text], short-circuiting the moment it passes [maxFrameBytes] so a
     * multi-MB oversized replay is not fully scanned. Cheaper than allocating a `ByteArray` for
     * every normal frame.
     */
    private fun utf8Size(text: String): Long {
        var bytes = 0L
        var i = 0
        while (i < text.length) {
            val code = text[i].code
            bytes += when {
                code < 0x80 -> 1
                code < 0x800 -> 2
                code in 0xD800..0xDBFF && i + 1 < text.length && text[i + 1].code in 0xDC00..0xDFFF -> {
                    i++ // consume the low surrogate: a full pair encodes to 4 bytes
                    4
                }
                else -> 3
            }
            if (bytes > maxFrameBytes) return bytes // over the cap — exact size no longer matters
            i++
        }
        return bytes
    }

    public companion object {
        /** Max [AwayDigest.recent] entries carried in the once-per-reconnect digest. */
        public const val DEFAULT_DIGEST_LIMIT: Int = 20
    }
}
