package wang.yaojia.webterm.transport

import app.cash.turbine.test
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineStart
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withTimeout
import okhttp3.OkHttpClient
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import okhttp3.mockwebserver.MockResponse
import okhttp3.mockwebserver.MockWebServer
import okhttp3.mockwebserver.SocketPolicy
import org.junit.jupiter.api.AfterEach
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test
import wang.yaojia.webterm.wire.HostEndpoint
import java.io.IOException

/**
 * A7 WS contract over MockWebServer's `withWebSocketUpgrade`:
 *  - the `Origin` header on the upgrade byte-equals `endpoint.originHeader` (THE CSWSH defence);
 *  - frames arrive in order; a clean server close COMPLETES the flow normally; a server abort
 *    SURFACES as a flow error (the two must stay distinguishable for A14);
 *  - `send` writes the exact frame; `close` detaches (server sees a 1000 close).
 */
class OkHttpTermTransportTest {

    private lateinit var server: MockWebServer
    private val client: OkHttpClient = OkHttpClientFactory.create()

    @BeforeEach
    fun setUp() {
        server = MockWebServer()
        server.start()
    }

    @AfterEach
    fun tearDown() {
        // Force-close any live WS connection so MockWebServer isn't left waiting on an open socket.
        client.dispatcher.cancelAll()
        client.connectionPool.evictAll()
        // MockWebServer's WS shutdown occasionally times out draining its own task queue even after
        // the socket is closed — a teardown artifact, not a product concern (assertions already ran).
        runCatching { server.shutdown() }
        client.dispatcher.executorService.shutdown()
    }

    private fun endpoint(): HostEndpoint =
        requireNotNull(HostEndpoint.fromBaseUrl("http://${server.hostName}:${server.port}"))

    private fun transport(): OkHttpTermTransport = OkHttpTermTransport(client)

    @Test
    fun stampsOriginHeaderOnTheWsUpgrade() = runBlocking {
        // Arrange
        val serverWs = RecordingServerWebSocket()
        server.enqueue(MockResponse().withWebSocketUpgrade(serverWs))
        val endpoint = endpoint()

        // Act
        val connection = withTimeout(TIMEOUT_MS) { transport().connect(endpoint) }

        // Assert: the upgrade request carried the byte-equal Origin.
        val upgrade = server.takeRequest()
        assertEquals(endpoint.originHeader, upgrade.getHeader("Origin"))
        connection.close()
    }

    @Test
    fun deliversFramesInArrivalOrderThenCompletesOnCleanClose() = runBlocking {
        // Arrange
        val serverWs = RecordingServerWebSocket()
        server.enqueue(MockResponse().withWebSocketUpgrade(serverWs))
        val connection = withTimeout(TIMEOUT_MS) { transport().connect(endpoint()) }
        val socket = withTimeout(TIMEOUT_MS) { serverWs.opened.await() }

        // Act + Assert
        connection.frames.test {
            socket.send("frame-1")
            socket.send("frame-2")
            assertEquals("frame-1", awaitItem())
            assertEquals("frame-2", awaitItem())
            socket.close(NORMAL_CLOSURE_CODE, "done")
            awaitComplete() // clean close → NORMAL flow completion
        }
    }

    @Test
    fun serverAbortSurfacesAsAFlowError() = runBlocking {
        // Arrange
        val serverWs = RecordingServerWebSocket()
        server.enqueue(MockResponse().withWebSocketUpgrade(serverWs))
        val connection = withTimeout(TIMEOUT_MS) { transport().connect(endpoint()) }
        withTimeout(TIMEOUT_MS) { serverWs.opened.await() }

        // Act + Assert
        connection.frames.test {
            server.shutdown() // abrupt TCP teardown (no WS close frame), not a clean close
            val error = awaitError()
            assertTrue(error is IOException, "expected a transport IOException, got $error")
        }
    }

    @Test
    fun sendWritesTheExactFrameToTheServer() = runBlocking {
        // Arrange
        val serverWs = RecordingServerWebSocket()
        server.enqueue(MockResponse().withWebSocketUpgrade(serverWs))
        val connection = withTimeout(TIMEOUT_MS) { transport().connect(endpoint()) }
        withTimeout(TIMEOUT_MS) { serverWs.opened.await() }

        // Act
        connection.send("""{"type":"attach","sessionId":null}""")

        // Assert
        val received = withTimeout(TIMEOUT_MS) { serverWs.messages.receive() }
        assertEquals("""{"type":"attach","sessionId":null}""", received)
        connection.close()
    }

    @Test
    fun closeDetachesWithANormalClosure() = runBlocking {
        // Arrange
        val serverWs = RecordingServerWebSocket()
        server.enqueue(MockResponse().withWebSocketUpgrade(serverWs))
        val connection = withTimeout(TIMEOUT_MS) { transport().connect(endpoint()) }
        withTimeout(TIMEOUT_MS) { serverWs.opened.await() }

        // Act: client detach (server PTY keeps running — this is never a kill).
        connection.close()

        // Assert: the server observed a graceful 1000 close.
        val code = withTimeout(TIMEOUT_MS) { serverWs.closingCode.await() }
        assertEquals(NORMAL_CLOSURE_CODE, code)
    }

    @Test
    fun connectPingableReportsLivenessAcrossClose() = runBlocking {
        // Arrange
        val serverWs = RecordingServerWebSocket()
        server.enqueue(MockResponse().withWebSocketUpgrade(serverWs))

        // Act
        val pingable = withTimeout(TIMEOUT_MS) { transport().connectPingable(endpoint()) }
        withTimeout(TIMEOUT_MS) { serverWs.opened.await() }

        // Assert: alive before close, dead after a client detach.
        assertTrue(pingable.pinger.ping())
        pingable.connection.close()
        assertFalse(pingable.pinger.ping())
    }

    @Test
    fun cancellingAnInFlightDialTearsDownTheSocketWithNoLeak() = runBlocking {
        // Arrange: the server accepts the TCP connection but never completes the WS upgrade, so onOpen
        // never fires — connect() stays suspended in the handshake (a genuine in-flight dial).
        server.enqueue(MockResponse().apply { socketPolicy = SocketPolicy.NO_RESPONSE })

        // Act: dial in a child job, wait until OkHttp has the call in flight, then cancel mid-handshake.
        val dialing = launch(start = CoroutineStart.UNDISPATCHED) {
            transport().connect(endpoint())
        }
        withTimeout(TIMEOUT_MS) {
            while (client.dispatcher.runningCallsCount() == 0) delay(POLL_MS)
        }
        dialing.cancelAndJoin()

        // Assert: the just-created WebSocket was torn down — no leaked running call survives the cancel.
        // (Without the fix the call sits blocked on readTimeout(0) forever, so this poll never settles.)
        withTimeout(TIMEOUT_MS) {
            while (client.dispatcher.runningCallsCount() != 0) delay(POLL_MS)
        }
        assertEquals(0, client.dispatcher.runningCallsCount())
    }

    @Test
    fun cancellingFrameCollectionPropagatesCancellationCleanly() = runBlocking {
        // Arrange
        val serverWs = RecordingServerWebSocket()
        server.enqueue(MockResponse().withWebSocketUpgrade(serverWs))
        val connection = withTimeout(TIMEOUT_MS) { transport().connect(endpoint()) }
        withTimeout(TIMEOUT_MS) { serverWs.opened.await() }

        // Act: park a collector inside the flow (no frames sent → the pump suspends draining the mailbox).
        val collector = launch(start = CoroutineStart.UNDISPATCHED) {
            connection.frames.collect { }
        }

        // Assert: cancelling the collection lands as a cancellation (the pump must NOT swallow the
        // CancellationException into a normal close) and join returns promptly — no deadlock.
        withTimeout(TIMEOUT_MS) { collector.cancelAndJoin() }
        assertTrue(collector.isCancelled, "collector must end cancelled, not completed normally")
    }

    private companion object {
        const val TIMEOUT_MS = 5_000L
        const val POLL_MS = 10L
    }
}

/** Server side of the wire: records the upgrade socket, inbound frames, and the client's close code. */
private class RecordingServerWebSocket : WebSocketListener() {
    val opened = CompletableDeferred<WebSocket>()
    val messages = Channel<String>(Channel.UNLIMITED)
    val closingCode = CompletableDeferred<Int>()

    override fun onOpen(webSocket: WebSocket, response: Response) {
        opened.complete(webSocket)
    }

    override fun onMessage(webSocket: WebSocket, text: String) {
        messages.trySend(text)
    }

    override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
        closingCode.complete(code)
    }
}
