package com.peachpasswords.sync

import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import okio.ByteString.Companion.toByteString
import org.json.JSONObject
import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.Test
import java.util.concurrent.CompletableFuture
import java.util.concurrent.TimeUnit

/**
 * Covers the two behaviours that are easy to get quietly wrong: a bad token
 * must be refused with no reply at all, and an abandoned code must stop working.
 * Both are invisible when broken — a leaky token check still pairs happily with
 * the right code, and a window that never closes looks identical to one that
 * does until months later.
 */
class CodePairingServerTest {

    private val identity = LocalChirpIdentity(
        deviceId = "test-device",
        deviceName = "Test Phone",
        deviceType = DeviceType.ANDROID,
        longTermPublicKeyHex = "00".repeat(32),
        signingPublicKeyHex = "00".repeat(32),
    )

    /** Enough of the registry for pairing; the real store needs Android crypto. */
    private class FakeRegistry : PairedDeviceRegistry {
        val devices = mutableListOf<PairedDevice>()
        override fun getAll(): List<PairedDevice> = devices
        override fun get(id: String): PairedDevice? = devices.firstOrNull { it.deviceId == id }
        override fun upsert(device: PairedDevice) { devices.add(device) }
    }

    private fun server() = CodePairingServer(FakeRegistry())

    private class Client {
        val opened = CompletableFuture<Boolean>()
        val closed = CompletableFuture<Boolean>()
        val frames = java.util.concurrent.LinkedBlockingQueue<String>()
        val binaryFrames = java.util.concurrent.LinkedBlockingQueue<ByteArray>()
        lateinit var ws: WebSocket

        fun connect(port: Int) {
            ws = OkHttpClient.Builder().readTimeout(10, TimeUnit.SECONDS).build().newWebSocket(
                Request.Builder().url("ws://127.0.0.1:$port").build(),
                object : WebSocketListener() {
                    override fun onOpen(webSocket: WebSocket, response: Response) { opened.complete(true) }
                    override fun onMessage(webSocket: WebSocket, text: String) { frames.add(text) }
                    override fun onMessage(webSocket: WebSocket, bytes: okio.ByteString) { binaryFrames.add(bytes.toByteArray()) }
                    override fun onFailure(webSocket: WebSocket, t: Throwable, r: Response?) { closed.complete(true) }
                    override fun onClosed(webSocket: WebSocket, code: Int, reason: String) { closed.complete(true) }
                },
            )
            opened.get(5, TimeUnit.SECONDS)
        }
    }

    /**
     * Collects into a queue on a job that stays alive. Terminating the flow with
     * `first {}` would trip awaitClose and shut the listener down before the
     * test could connect to it.
     */
    private class Harness(val scope: CoroutineScope) {
        val events = java.util.concurrent.LinkedBlockingQueue<PairingEvent>()
        inline fun <reified T : PairingEvent> await(timeoutMs: Long = 10_000): T {
            val deadline = System.currentTimeMillis() + timeoutMs
            val seen = mutableListOf<PairingEvent>()
            while (System.currentTimeMillis() < deadline) {
                val e = events.poll(200, TimeUnit.MILLISECONDS) ?: continue
                seen.add(e)
                if (e is T) return e
            }
            throw AssertionError("never saw ${T::class.simpleName}; saw $seen")
        }
    }

    private fun CoroutineScope.run(srv: CodePairingServer, windowMs: Long = CodePairingServer.PAIRING_WINDOW_MS): Harness {
        val h = Harness(this)
        launch { srv.start(this@run, ByteArray(32), ByteArray(32), identity, windowMs).collect { h.events.add(it) } }
        return h
    }

    @Test fun `a wrong token gets no reply and the socket closes`() = runBlocking {
        val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
        val srv = server()
        try {
            val listening = scope.run(srv).await<PairingEvent.Listening>()

            val c = Client()
            c.connect(listening.port)
            c.ws.send("""{"type":"PEACH_HELLO","token":"00000000"}""")

            assertTrue("socket should be closed on a bad token", c.closed.get(5, TimeUnit.SECONDS))
            assertEquals("nothing should be sent back — silence is the signal", 0, c.frames.size)
        } finally {
            srv.stop()
        }
    }

    /**
     * The browser's pairing socket is binary end to end (`openPairingSocket` in
     * peach-browser's pairing-dial.ts sets `binaryType = "arraybuffer"` and
     * sends `Uint8Array`). This listener required a TEXT hello and answered with
     * a TEXT ack, so every correct code was silently closed on and the browser
     * reported "invalid code" — pairing could not succeed at all. Neither half
     * fails loudly, which is exactly why both are pinned here.
     */
    @Test fun `a binary hello is accepted and answered with a binary ack`() = runBlocking {
        val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
        val srv = server()
        try {
            val listening = scope.run(srv).await<PairingEvent.Listening>()
            val token = PairingCode.parse(listening.code)!!.tokenHex

            val c = Client()
            c.connect(listening.port)
            c.ws.send("""{"type":"PEACH_HELLO","token":"$token"}""".toByteArray().toByteString())

            val ack = c.binaryFrames.poll(10, TimeUnit.SECONDS)
            assertTrue("the phone must answer a binary hello, not close on it", ack != null)
            assertTrue(
                "the ack must be a BINARY frame — a text one decodes to zero bytes in the browser",
                c.frames.isEmpty(),
            )
            assertTrue(String(ack!!), String(ack).contains("PEACH_HELLO_ACK"))
        } finally {
            srv.stop()
        }
    }

    @Test fun `LAN adoption dispatches before device-pair auth and keeps both SAS gates`() = runBlocking {
        val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
        val registry = FakeRegistry()
        val srv = CodePairingServer(registry)
        val root = ByteArray(32) { (it + 1).toByte() }
        val sourceKeys = deriveLocalSyncKeypair(root, "source-id")
        val receiverKeys = deriveLocalSyncKeypair(root, "receiver-id")
        val sourceIdentity = identity.copy(
            deviceId = "source-id",
            deviceName = "Source phone",
            longTermPublicKeyHex = Rendezvous.bytesToHex(sourceKeys.publicKey),
            relayUrl = "wss://relay.example.test",
        )
        var rootRequested = false
        val sourceDeps = object : AdoptionSourceDeps {
            override val sourceIdentity = sourceIdentity
            override suspend fun rootKeyHex(): String {
                rootRequested = true
                return Rendezvous.bytesToHex(root)
            }
            override suspend fun vaultId() = computeVaultId(root)
            override suspend fun onConsent(request: AdoptionFrame.Request) = true
            override suspend fun upsertPairedDevice(device: PairedDevice) {
                registry.upsert(device)
            }
        }
        val h = Harness(scope)
        srv.confirmationSink = { it.session.confirm(true) }
        scope.launch {
            srv.start(
                this,
                sourceKeys.secretKey,
                sourceKeys.publicKey,
                sourceIdentity,
                adoptionSourceDeps = sourceDeps,
            ).collect { h.events.add(it) }
        }

        try {
            val listening = h.await<PairingEvent.Listening>()
            val token = PairingCode.parse(listening.code)!!.tokenHex
            val client = Client()
            client.connect(listening.port)
            client.ws.send(
                JSONObject()
                    .put("type", "PEACH_HELLO")
                    .put("token", token)
                    .put("flow", "adoption")
                    .toString()
                    .toByteArray()
                    .toByteString(),
            )
            val ack = JSONObject(String(client.binaryFrames.poll(10, TimeUnit.SECONDS)))
            val noise = noiseNKInitiator(
                hexToPubkey(ack.getString("ephemeralPublicKeyHex")),
                { client.ws.send(it.toByteString()) },
                {
                    client.binaryFrames.poll(10, TimeUnit.SECONDS)
                        ?: throw IllegalStateException("timed out waiting for server frame")
                },
            )
            val channel = object : PairingTransport {
                override suspend fun send(msg: ByteArray) {
                    client.ws.send(noise.encrypt(msg).toByteString())
                }
                override suspend fun receive(): ByteArray {
                    val frame = client.binaryFrames.poll(10, TimeUnit.SECONDS)
                        ?: throw IllegalStateException("timed out waiting for encrypted frame")
                    return noise.decrypt(frame)
                }
                override suspend fun close() {
                    client.ws.close(1000, "done")
                }
            }
            val receiver = requestAdoption(
                channel,
                object : AdoptionReceiverDeps {
                    override val deviceName = "Browser"
                    override val deviceType = DeviceType.EXTENSION
                    override fun generateDeviceId() = "session-request"
                    override suspend fun confirmSas() = true
                    override suspend fun createVaultFromRootKey(rootKeyHex: String) =
                        computeVaultId(hexToPubkey(rootKeyHex))
                    override suspend fun deriveLocalIdentity() = LocalChirpIdentity(
                        deviceId = "receiver-id",
                        deviceName = "Browser",
                        deviceType = DeviceType.EXTENSION,
                        longTermPublicKeyHex = Rendezvous.bytesToHex(receiverKeys.publicKey),
                        signingPublicKeyHex = "ef".repeat(32),
                    )
                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit
                    override fun onError(error: String) = throw AssertionError(error)
                },
            )

            assertTrue(rootRequested)
            assertEquals("source-id", receiver?.deviceId)
            assertEquals("Browser", h.await<PairingEvent.Paired>().deviceName)
            assertEquals("receiver-id", registry.devices.single().deviceId)
            assertEquals(64, registry.devices.single().rendezvousSecretHex?.length)
            assertEquals("wss://relay.example.test", registry.devices.single().relayUrl)
        } finally {
            srv.stop()
        }
    }

    @Test fun `the pairing window expires and stops accepting`() = runBlocking {
        val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
        val srv = server()
        try {
            // 300ms rather than the real five minutes; the window is a parameter
            // precisely so this is testable at all.
            val h = scope.run(srv, windowMs = 300)
            val listening = h.await<PairingEvent.Listening>()
            h.await<PairingEvent.Expired>(timeoutMs = 5_000)

            var refused = false
            runCatching { java.net.Socket("127.0.0.1", listening.port).close() }
                .onFailure { refused = true }
            assertTrue("the port should be closed after the window", refused)
        } finally {
            srv.stop()
        }
    }

    @Test fun `the advertised address is never a tunnel interface`() {
        // Not asserting a specific address — CI and dev machines differ — only
        // that whatever is chosen is one a LAN peer could route back to.
        val ip = server().localIpv4()
        if (ip != null) {
            assertTrue("expected a site-local address, got $ip", java.net.InetAddress.getByName(ip).isSiteLocalAddress)
        }
    }
}