package com.peachpasswords.sync

import java.io.InputStream
import java.io.OutputStream
import java.net.Socket
import java.security.MessageDigest
import java.util.Base64

/**
 * A minimal RFC 6455 server endpoint — just enough of the protocol to talk to a
 * browser extension's `new WebSocket(...)`.
 *
 * Android already had a pairing listener, but it spoke
 * raw length-prefixed TCP frames. A browser extension cannot open a raw socket;
 * `ws://` is the only outbound connection it can make. Since the browser is
 * always the dialer and the phone always the listener, that asymmetry is the
 * whole reason this exists — so the phone has to be the one that speaks
 * WebSocket, and this is the smallest thing that does.
 *
 * Deliberately not implemented, because the peer is one known client and not
 * the open web: no extensions/compression negotiation, no subprotocols, no
 * fragmentation (the browser sends each frame whole), and no client mode.
 * Continuation frames are rejected rather than silently mishandled.
 */
private const val WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

sealed interface WsMessage {
    data class Text(val text: String) : WsMessage
    data class Binary(val bytes: ByteArray) : WsMessage
    data object Closed : WsMessage
}

class WebSocketConnection(
    private val socket: Socket,
    private val input: InputStream,
    private val output: OutputStream,
) {
    /**
     * Complete the HTTP upgrade. Returns false if this is not a WebSocket
     * handshake, in which case the caller should just drop the socket — an
     * unexpected plain-HTTP probe on the LAN is not worth answering.
     */
    fun accept(): Boolean {
        val request = readHttpHeaders() ?: return false
        val key = request.entries.firstOrNull { it.key.equals("sec-websocket-key", true) }?.value
            ?: return false
        val accept = Base64.getEncoder().encodeToString(
            MessageDigest.getInstance("SHA-1").digest((key + WS_GUID).toByteArray(Charsets.US_ASCII)),
        )
        output.write(
            (
                "HTTP/1.1 101 Switching Protocols\r\n" +
                    "Upgrade: websocket\r\n" +
                    "Connection: Upgrade\r\n" +
                    "Sec-WebSocket-Accept: $accept\r\n\r\n"
                ).toByteArray(Charsets.US_ASCII),
        )
        output.flush()
        return true
    }

    fun sendText(text: String) = writeFrame(0x1, text.toByteArray(Charsets.UTF_8))

    fun send(bytes: ByteArray) {
        try {
            requireSyncFrameSize(bytes.size)
        } catch (error: SyncFrameTooLargeException) {
            close(1009, "sync frame too large")
            throw error
        }
        writeFrame(0x2, bytes)
    }

    /** Blocks until a whole message arrives. [WsMessage.Closed] ends the conversation. */
    fun receive(): WsMessage {
        while (true) {
            val b0 = input.read()
            if (b0 < 0) return WsMessage.Closed
            val fin = (b0 and 0x80) != 0
            when (val opcode = b0 and 0x0f) {
                0x8 -> {
                    val payload = readPayload()
                    val code = if (payload.size >= 2) {
                        ((payload[0].toInt() and 0xff) shl 8) or (payload[1].toInt() and 0xff)
                    } else {
                        null
                    }
                    if (code == 1009) throw SyncFrameTooLargeException()
                    return WsMessage.Closed
                }
                0x9 -> { writeFrame(0xA, readPayload()); continue } // ping → pong
                0xA -> { readPayload(); continue }                  // pong, ignored
                0x1, 0x2 -> {
                    if (!fin) return WsMessage.Closed
                    val payload = readPayload()
                    return if (opcode == 0x1) WsMessage.Text(payload.toString(Charsets.UTF_8))
                    else WsMessage.Binary(payload)
                }
                else -> return WsMessage.Closed
            }
        }
    }

    /**
     * Close without a status frame. The wire contract makes this silence the
     * signal for a wrong or expired token — a phone that answered with an error
     * would tell an unrelated LAN host that it is a Peach device holding an open
     * pairing window.
     */
    fun closeSilently() {
        runCatching { socket.close() }
    }

    private fun readHttpHeaders(): Map<String, String>? {
        val head = StringBuilder()
        // Read to the blank line, bounded so a non-HTTP peer can't stream forever.
        while (head.length < 8192 && !head.endsWith("\r\n\r\n")) {
            val c = input.read()
            if (c < 0) return null
            head.append(c.toChar())
        }
        val lines = head.toString().split("\r\n")
        if (lines.firstOrNull()?.startsWith("GET ") != true) return null
        return lines.drop(1)
            .mapNotNull { line ->
                val i = line.indexOf(':')
                if (i <= 0) null else line.substring(0, i).trim() to line.substring(i + 1).trim()
            }
            .toMap()
    }

    private fun readPayload(): ByteArray {
        val b1 = input.read()
        if (b1 < 0) return ByteArray(0)
        val masked = (b1 and 0x80) != 0
        var len = (b1 and 0x7f).toLong()
        when (len) {
            126L -> len = ((readByte() shl 8) or readByte()).toLong()
            127L -> {
                len = 0
                repeat(8) { len = (len shl 8) or readByte().toLong() }
            }
        }
        // A vault sync frame is large but bounded; refuse an absurd length
        // rather than trying to allocate it.
        if (len !in 0..MAX_SYNC_FRAME_BYTES.toLong()) {
            close(1009, "sync frame too large")
            throw SyncFrameTooLargeException("sync frame exceeds $MAX_SYNC_FRAME_BYTES bytes")
        }

        val mask = if (masked) ByteArray(4) { readByte().toByte() } else null
        val payload = ByteArray(len.toInt())
        var read = 0
        while (read < payload.size) {
            val n = input.read(payload, read, payload.size - read)
            if (n < 0) break
            read += n
        }
        if (mask != null) {
            for (i in payload.indices) payload[i] = (payload[i].toInt() xor mask[i % 4].toInt()).toByte()
        }
        return payload
    }

    private fun readByte(): Int {
        val b = input.read()
        if (b < 0) throw java.io.EOFException("websocket closed mid-frame")
        return b
    }

    fun close(code: Int = 1000, reason: String = "") {
        val reasonBytes = reason.toByteArray(Charsets.UTF_8).take(123).toByteArray()
        val payload = byteArrayOf((code ushr 8).toByte(), code.toByte()) + reasonBytes
        runCatching { writeFrame(0x8, payload) }
        runCatching { socket.close() }
    }

    /** Server-to-client frames are never masked (RFC 6455 §5.1). */
    private fun writeFrame(opcode: Int, payload: ByteArray) {
        output.write(0x80 or opcode)
        when {
            payload.size < 126 -> output.write(payload.size)
            payload.size <= 0xffff -> {
                output.write(126)
                output.write((payload.size shr 8) and 0xff)
                output.write(payload.size and 0xff)
            }
            else -> {
                output.write(127)
                repeat(4) { output.write(0) }
                output.write((payload.size shr 24) and 0xff)
                output.write((payload.size shr 16) and 0xff)
                output.write((payload.size shr 8) and 0xff)
                output.write(payload.size and 0xff)
            }
        }
        output.write(payload)
        output.flush()
    }
}