package wang.yaojia.webterm.session

import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.receiveAsFlow
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.TestScope
import kotlinx.coroutines.test.runTest
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertNotNull
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import wang.yaojia.webterm.testsupport.FakePingableTermTransport
import wang.yaojia.webterm.testsupport.FakeTermTransport
import wang.yaojia.webterm.testsupport.FakeTimeSource
import wang.yaojia.webterm.wire.ClientMessage
import wang.yaojia.webterm.wire.HostEndpoint
import wang.yaojia.webterm.wire.MessageCodec
import wang.yaojia.webterm.wire.TermTransport
import wang.yaojia.webterm.wire.TimelineEvent
import wang.yaojia.webterm.wire.TransportConnection
import wang.yaojia.webterm.wire.Tunables
import kotlin.time.ExperimentalTime
import kotlin.time.Duration.Companion.seconds

/**
 * A14 verification (plan §5 "A14 Verify"): the [SessionEngine] state machine driven purely through
 * the [FakeTermTransport]/[FakePingableTermTransport] doubles and the coroutine virtual clock under
 * `runTest`. No OkHttp, no Android, no wall-clock waits.
 */
class SessionEngineTest {

    private companion object {
        /** A valid lower-case UUID-v4 the server can issue on `attached`. */
        const val SERVER_ID = "11111111-1111-4111-8111-111111111111"
        const val ATTACH_NEW = """{"type":"attach","sessionId":null}"""
        const val APPROVE = """{"type":"approve"}"""
    }

    private fun endpoint(): HostEndpoint = HostEndpoint.fromBaseUrl("http://localhost:3000")!!

    private fun attachWith(id: String): String = MessageCodec.encode(ClientMessage.Attach(id))

    private fun attachedFrame(id: String): String = """{"type":"attached","sessionId":"$id"}"""

    private fun statusFrame(pending: Boolean, gate: String? = null): String {
        val gatePart = if (gate != null) ""","gate":"$gate"""" else ""
        return """{"type":"status","status":"working","pending":$pending$gatePart}"""
    }

    private fun TestScope.newEngine(
        transport: TermTransport,
        initialSessionId: String? = null,
        initialCwd: String? = null,
        maxFrameBytes: Long = Tunables.MAX_WS_MESSAGE_BYTES,
        awayTimeline: suspend (Long) -> List<TimelineEvent> = { emptyList() },
        nowMs: () -> Long = { 0L },
    ): SessionEngine = SessionEngine(
        transport = transport,
        endpoint = endpoint(),
        scope = backgroundScope,
        initialSessionId = initialSessionId,
        initialCwd = initialCwd,
        awayTimeline = awayTimeline,
        nowMs = nowMs,
        maxFrameBytes = maxFrameBytes,
    )

    /** Live-updating list of everything the engine emits (single consumer, buffered channel). */
    private fun TestScope.collectEvents(engine: SessionEngine): List<SessionEvent> {
        val out = mutableListOf<SessionEvent>()
        backgroundScope.launch { engine.events.collect { out += it } }
        return out
    }

    private fun List<SessionEvent>.connectionStates(): List<ConnectionState> =
        filterIsInstance<Connection>().map { it.state }

    // ── attach-first ordering ──────────────────────────────────────────────────────────────

    @Test
    fun `attach is the first frame on connect with an explicit null sessionId for a new session`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport)
        collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()

        assertEquals(listOf(ATTACH_NEW), transport.sentFramesByConnection[0])
    }

    @Test
    fun `on reconnect attach is re-sent first carrying the adopted server id`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport)
        collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()
        transport.emit(attachedFrame(SERVER_ID)) // server confirms + issues the id
        testScheduler.runCurrent()

        transport.emitError(RuntimeException("ws dropped")) // conn 0 drops → back-off 1s
        testScheduler.runCurrent()
        testScheduler.advanceTimeBy(1.seconds)
        testScheduler.runCurrent()

        assertEquals(attachWith(SERVER_ID), transport.sentFramesByConnection[1].first())
    }

    // ── adopt-server-id ─────────────────────────────────────────────────────────────────────

    @Test
    fun `every attached frame is adopted and surfaced as an Adopted event`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()
        transport.emit(attachedFrame(SERVER_ID))
        testScheduler.runCurrent()

        assertTrue(events.contains(Adopted(SERVER_ID)), "expected an Adopted($SERVER_ID) event, got $events")
    }

    // ── connect-now-on-foreground (and no-reset-on-foreground) ──────────────────────────────

    @Test
    fun `foregrounding connects immediately without resetting the back-off ladder`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent() // conn 0 connected
        transport.emitError(RuntimeException("drop")) // → Reconnecting(1, 1s), waiting
        testScheduler.runCurrent()

        transport.scriptConnectFailure() // the foreground connect will fail so we can read the ladder
        engine.notifyForegrounded()
        testScheduler.runCurrent() // wake → connect NOW (no 1s advance)

        assertEquals(2, transport.connectAttempts.size) // connected immediately, not after 1s
        val reconnecting = events.connectionStates().filterIsInstance<ConnectionState.Reconnecting>()
        assertEquals(2.seconds, reconnecting.last().next) // ladder advanced to 2s, NOT reset to 1s
    }

    // ── close == detach, NEVER kill ─────────────────────────────────────────────────────────

    @Test
    fun `close detaches the connection without a kill and never reconnects`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()
        engine.close()
        testScheduler.runCurrent()

        assertEquals(1, transport.closeCallCount) // exactly one detach
        assertEquals(listOf(ATTACH_NEW), transport.sentFrames) // only attach was ever sent — no kill frame
        assertTrue(events.contains(Connection(ConnectionState.Closed)))

        testScheduler.advanceTimeBy(60.seconds) // no back-off reconnect after a client close
        testScheduler.runCurrent()
        assertEquals(1, transport.connectAttempts.size)
    }

    // ── oversized-replay is a terminal, non-retryable failure ───────────────────────────────

    @Test
    fun `an oversized frame yields replayTooLarge and never reconnects`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport, maxFrameBytes = 32)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()
        transport.emit("x".repeat(100)) // 100 bytes > 32-byte cap
        testScheduler.runCurrent()

        assertTrue(
            events.contains(Connection(ConnectionState.Failed(FailureReason.REPLAY_TOO_LARGE))),
            "expected a REPLAY_TOO_LARGE failure, got $events",
        )
        // Distinct from a retryable disconnect: no Reconnecting, no second connect.
        assertFalse(events.connectionStates().any { it is ConnectionState.Reconnecting })
        testScheduler.advanceTimeBy(60.seconds)
        testScheduler.runCurrent()
        assertEquals(1, transport.connectAttempts.size)
    }

    // ── gate-decision epoch drop (two-line stale-guard) ─────────────────────────────────────

    @Test
    fun `a decision against a stale gate epoch is dropped while a live epoch sends`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()

        transport.emit(statusFrame(pending = true, gate = "tool")) // rising edge → epoch 1
        testScheduler.runCurrent()
        val epoch1 = events.filterIsInstance<Gate>().last().gate!!.epoch
        assertEquals(1, epoch1)

        engine.decideGate(epoch1, ClientMessage.Approve(null)) // live epoch → sends
        testScheduler.runCurrent()

        transport.emit(statusFrame(pending = false)) // gate lifted
        testScheduler.runCurrent()
        transport.emit(statusFrame(pending = true, gate = "tool")) // rising edge → epoch 2
        testScheduler.runCurrent()

        engine.decideGate(epoch1, ClientMessage.Approve(null)) // stale epoch → dropped
        testScheduler.runCurrent()

        assertEquals(1, transport.sentFrames.count { it == APPROVE }) // exactly one approve reached the wire
    }

    @Test
    fun `a second decision against the same live gate epoch does not double-send`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()
        transport.emit(statusFrame(pending = true, gate = "tool")) // rising edge → epoch 1, gate held
        testScheduler.runCurrent()
        val epoch1 = events.filterIsInstance<Gate>().last().gate!!.epoch

        engine.decideGate(epoch1, ClientMessage.Approve(null)) // first tap → sends, retires the gate
        testScheduler.runCurrent()
        engine.decideGate(epoch1, ClientMessage.Approve(null)) // second same-epoch tap BEFORE any status update
        testScheduler.runCurrent()

        assertEquals(1, transport.sentFrames.count { it == APPROVE }) // exactly ONE approve reached the wire
    }

    // ── close during the dial/attach window still detaches (close == detach, #2 / §6.6) ─────

    @Test
    fun `close during the dial window detaches the freshly-opened connection and emits nothing further`() = runTest {
        val gate = CompletableDeferred<Unit>()
        val transport = GatedConnectTransport(connectGate = gate)
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent() // supervisor runs; connect() suspends on the gate (currentConnection is null)

        engine.close() // close lands DURING the dial window
        testScheduler.runCurrent()

        gate.complete(Unit) // the connection now opens, after close() already fired
        testScheduler.runCurrent()

        assertTrue(transport.connection.isClosed, "the connection opened during close() must be detached")
        val states = events.connectionStates()
        assertFalse(states.contains(ConnectionState.Connected), "must not emit Connected after close, got $states")
        assertFalse(events.any { it is Adopted }, "must not adopt a server id after close, got $events")
        assertEquals(1, transport.connectCount) // no reconnect scheduled after a client close
    }

    @Test
    fun `close during the attach window detaches the freshly-opened connection and emits nothing further`() = runTest {
        val sendGate = CompletableDeferred<Unit>()
        val transport = GatedConnectTransport(sendGate = sendGate) // connect succeeds; the first send suspends
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent() // connect returns; attachFirst's send() suspends on the gate

        engine.close() // close lands DURING the attach window
        testScheduler.runCurrent()

        sendGate.complete(Unit) // the attach send now completes, after close() already fired
        testScheduler.runCurrent()

        assertTrue(transport.connection.isClosed, "the connection opened during close() must be detached")
        val states = events.connectionStates()
        assertFalse(states.contains(ConnectionState.Connected), "must not emit Connected after close, got $states")
        assertEquals(1, transport.connectCount) // no reconnect scheduled after a client close
    }

    // ── back-off ladder 1 → 2 → 4 → 8 → 16 → 30 ─────────────────────────────────────────────

    @Test
    fun `the reconnect back-off ladder climbs 1 2 4 8 16 30`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent() // conn 0 connected
        repeat(6) { transport.scriptConnectFailure() } // every reconnect attempt fails → ladder climbs
        transport.emitError(RuntimeException("drop")) // → Reconnecting(1, 1s)
        testScheduler.runCurrent()

        for (rung in listOf(1, 2, 4, 8, 16, 30)) {
            testScheduler.advanceTimeBy(rung.seconds)
            testScheduler.runCurrent()
        }

        val delays = events.connectionStates()
            .filterIsInstance<ConnectionState.Reconnecting>()
            .map { it.next }
        assertEquals(listOf(1, 2, 4, 8, 16, 30).map { it.seconds }, delays.take(6))
    }

    // ── ping 25s / 2-miss (only through a PingableTermTransport) ─────────────────────────────

    @Test
    fun `two consecutive missed pongs on a pingable transport trigger a reconnect`() = runTest {
        val transport = FakePingableTermTransport()
        transport.pinger.scriptPongs(false, false) // both pings miss
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()
        assertEquals(1, transport.connectPingableCount) // took the pingable path
        assertEquals(0, transport.pinger.pingCallCount)

        testScheduler.advanceTimeBy(Tunables.PING_INTERVAL) // ping #1 (miss)
        testScheduler.runCurrent()
        assertEquals(1, transport.pinger.pingCallCount)

        testScheduler.advanceTimeBy(Tunables.PING_INTERVAL) // ping #2 (miss) → CONNECTION_LOST
        testScheduler.runCurrent()
        assertEquals(2, transport.pinger.pingCallCount)
        assertTrue(events.connectionStates().any { it is ConnectionState.Reconnecting })

        engine.close()
        testScheduler.runCurrent()
    }

    @Test
    fun `a plain non-pingable transport is never pinged`() = runTest {
        val transport = FakeTermTransport() // does NOT implement PingableTermTransport
        val engine = newEngine(transport)
        collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()
        // Advance well past several ping intervals — a plain transport must not disconnect on its own.
        testScheduler.advanceTimeBy(Tunables.PING_INTERVAL * 5)
        testScheduler.runCurrent()

        assertEquals(1, transport.connectAttempts.size) // still the one live connection, no ping-driven drop
    }

    // ── away-digest once-per-reconnect (all-zero suppressed by the UI, still delivered here) ──

    @OptIn(ExperimentalTime::class)
    @Test
    fun `a completed reconnect emits exactly one away digest reduced since the disconnect`() = runTest {
        val transport = FakeTermTransport()
        val clock = FakeTimeSource() // epochMillis starts at 0
        val timeline = listOf(
            TimelineEvent(at = 500, eventClass = "tool", label = "ran Bash"),
            TimelineEvent(at = 600, eventClass = "waiting", label = "needs approval"),
        )
        val engine = newEngine(transport, awayTimeline = { timeline }, nowMs = clock::epochMillis)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()
        transport.emit(attachedFrame(SERVER_ID)) // first attach → NO digest
        testScheduler.runCurrent()

        transport.emitError(RuntimeException("drop")) // disconnect at clock=0 → digest owed
        testScheduler.runCurrent()
        testScheduler.advanceTimeBy(1.seconds)
        testScheduler.runCurrent()
        transport.emit(attachedFrame(SERVER_ID)) // reconnect adopt → digest fires once
        testScheduler.runCurrent() // deliver the frame → onAttached launches + runs emitDigest

        val digests = events.filterIsInstance<Digest>()
        assertEquals(1, digests.size)
        val digest = digests.single().digest
        assertEquals(1, digest.toolRuns)
        assertEquals(1, digest.waitingCount)
        assertNotNull(digest)
    }

    // ── output + exit pass-through ──────────────────────────────────────────────────────────

    @Test
    fun `output frames pass through verbatim and an exit is terminal`() = runTest {
        val transport = FakeTermTransport()
        val engine = newEngine(transport)
        val events = collectEvents(engine)

        engine.start()
        testScheduler.runCurrent()
        transport.emit("""{"type":"output","data":"hello"}""")
        testScheduler.runCurrent()
        transport.emit("""{"type":"exit","code":-1,"reason":"spawn failed"}""")
        testScheduler.runCurrent()

        assertTrue(events.contains(Output("hello")))
        assertTrue(events.contains(Exited(-1, "spawn failed")))
        // Terminal: no reconnect after an exit.
        testScheduler.advanceTimeBy(60.seconds)
        testScheduler.runCurrent()
        assertEquals(1, transport.connectAttempts.size)
    }
}

/**
 * A [TransportConnection] that records its [close]/[send] and can suspend the FIRST send on a gate,
 * so a test can wedge `close()` into the attach window (before the connection is published).
 */
private class RecordingConnection(private val sendGate: CompletableDeferred<Unit>?) : TransportConnection {
    private val framesChannel = Channel<String>(Channel.UNLIMITED)
    override val frames: Flow<String> = framesChannel.receiveAsFlow()
    val sent: MutableList<String> = mutableListOf()

    @Volatile
    var isClosed: Boolean = false
        private set

    private var firstSend = true

    override suspend fun send(frame: String) {
        if (firstSend) {
            firstSend = false
            sendGate?.await() // suspend the attach send → the dispatcher is free to run close()
        }
        sent += frame
    }

    override suspend fun close() {
        isClosed = true
        framesChannel.close()
    }
}

/**
 * A [TermTransport] whose [connect] (and/or the connection's first [send]) suspends on a gate, so a
 * test can fire `close()` while the engine is still inside the dial/attach window (FIX 1 regression).
 */
private class GatedConnectTransport(
    private val connectGate: CompletableDeferred<Unit>? = null,
    sendGate: CompletableDeferred<Unit>? = null,
) : TermTransport {
    val connection: RecordingConnection = RecordingConnection(sendGate)

    @Volatile
    var connectCount: Int = 0
        private set

    override suspend fun connect(endpoint: HostEndpoint): TransportConnection {
        connectCount += 1
        connectGate?.await() // suspend the dial → the dispatcher is free to run close()
        return connection
    }
}
