package com.peachpasswords.sync

import org.junit.Assert.assertArrayEquals
import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.Test
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import okio.ByteString
import okio.ByteString.Companion.toByteString
import java.net.ServerSocket
import java.util.concurrent.CompletableFuture
import java.util.concurrent.LinkedBlockingQueue
import java.util.concurrent.TimeUnit
import kotlin.concurrent.thread

/**
 * Exercises the server against OkHttp's RFC 6455 client rather than a
 * hand-rolled one. The whole point of [WebSocketConnection] is that a browser
 * extension can talk to it, so testing it with an independent implementation of
 * the same spec is the closest thing to that available off device — a client
 * that shared my own misreading of the framing rules would prove nothing.
 */
class PeachWebSocketTest {

    private class Received {
        val text = LinkedBlockingQueue<String>()
        val binary = LinkedBlockingQueue<ByteArray>()
        val opened = CompletableFuture<Boolean>()
    }

    private fun listener(into: Received) = object : WebSocketListener() {
        override fun onMessage(webSocket: WebSocket, text: String) { into.text.add(text) }
        override fun onMessage(webSocket: WebSocket, bytes: ByteString) { into.binary.add(bytes.toByteArray()) }
        override fun onOpen(webSocket: WebSocket, response: Response) { into.opened.complete(true) }
    }

    /** Server side: accept one client, upgrade, then echo whatever arrives. */
    private fun echoServer(port: ServerSocket, onAccepted: (Boolean) -> Unit = {}) = thread {
        val client = port.accept()
        val ws = WebSocketConnection(client, client.getInputStream(), client.getOutputStream())
        val upgraded = ws.accept()
        onAccepted(upgraded)
        if (!upgraded) return@thread
        while (true) {
            when (val msg = ws.receive()) {
                is WsMessage.Text -> ws.sendText(msg.text)
                is WsMessage.Binary -> ws.send(msg.bytes)
                WsMessage.Closed -> return@thread
            }
        }
    }

    private fun connect(port: Int, into: Received): WebSocket {
        val ws = OkHttpClient.Builder()
            .readTimeout(10, TimeUnit.SECONDS)
            .build()
            .newWebSocket(Request.Builder().url("ws://127.0.0.1:$port").build(), listener(into))
        into.opened.get(5, TimeUnit.SECONDS)
        return ws
    }

    @Test fun `completes the upgrade handshake with a real client`() {
        ServerSocket(0).use { server ->
            val accepted = CompletableFuture<Boolean>()
            echoServer(server) { accepted.complete(it) }
            val ws = connect(server.localPort, Received())
            assertTrue("server should accept the upgrade", accepted.get(5, TimeUnit.SECONDS))
            ws.cancel()
        }
    }

    @Test fun `round-trips the plaintext hello frame`() {
        ServerSocket(0).use { server ->
            echoServer(server)
            val got = Received()
            val ws = connect(server.localPort, got)
            val hello = """{"type":"PEACH_HELLO","token":"deadbeef01234567"}"""
            ws.send(hello)
            assertEquals(hello, got.text.poll(5, TimeUnit.SECONDS))
            ws.cancel()
        }
    }

    @Test fun `round-trips binary frames across the length-encoding boundaries`() {
        // 125 is the last 7-bit length, 126 starts the 16-bit header, and 70000
        // crosses into the 64-bit one — a Noise-wrapped vault frame lands there.
        ServerSocket(0).use { server ->
            echoServer(server)
            val got = Received()
            val ws = connect(server.localPort, got)
            for (size in intArrayOf(0, 1, 125, 126, 8192, 70_000)) {
                val payload = ByteArray(size) { (it % 251).toByte() }
                ws.send(payload.toByteString())
                val echoed = got.binary.poll(10, TimeUnit.SECONDS)
                assertArrayEquals("size $size did not survive the round trip", payload, echoed)
            }
            ws.cancel()
        }
    }

    @Test fun `a plain HTTP request is refused rather than upgraded`() {
        ServerSocket(0).use { server ->
            val accepted = CompletableFuture<Boolean>()
            echoServer(server) { accepted.complete(it) }
            java.net.Socket("127.0.0.1", server.localPort).use { s ->
                s.getOutputStream().write("GET / HTTP/1.1\r\nHost: x\r\n\r\n".toByteArray())
                s.getOutputStream().flush()
                assertTrue("no Sec-WebSocket-Key means no upgrade", !accepted.get(5, TimeUnit.SECONDS))
            }
        }
    }

    @Test fun `close code 1009 surfaces a distinct frame too large error`() {
        ServerSocket(0).use { server ->
            val surfaced = CompletableFuture<Boolean>()
            thread {
                val client = server.accept()
                val ws = WebSocketConnection(client, client.getInputStream(), client.getOutputStream())
                if (!ws.accept()) {
                    surfaced.complete(false)
                    return@thread
                }
                surfaced.complete(
                    runCatching { ws.receive() }.exceptionOrNull() is SyncFrameTooLargeException,
                )
            }

            val ws = connect(server.localPort, Received())
            assertTrue(ws.close(1009, "sync frame too large"))
            assertTrue(surfaced.get(5, TimeUnit.SECONDS))
        }
    }

    @Test fun `sync frame cap is 24 MiB`() {
        assertEquals(24 * 1024 * 1024, MAX_SYNC_FRAME_BYTES)
        assertTrue(runCatching { requireSyncFrameSize(MAX_SYNC_FRAME_BYTES + 1) }.isFailure)
    }
}