package com.peachpasswords.sync

import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.callbackFlow
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import org.json.JSONObject
import java.net.Inet4Address
import java.net.NetworkInterface
import java.net.ServerSocket
import java.net.Socket
import javax.inject.Inject
import javax.inject.Singleton

/**
 * Code-based pairing, phone side. The phone listens and displays a code; the
 * browser dials out and presents the token from it.
 *
 * Wire contract (peach-browser is the authority — `.agent/CONTINUITY.md`
 * 2026-07-24T17:36Z, implemented there in `shared/pairing-dial.ts`):
 *
 *   1. browser dials `ws://<ip>:<port>` taken from the code
 *   2. browser → phone, plaintext text frame: {type:"PEACH_HELLO", token:"<lowercase hex>"}
 *   3. phone verifies the token; on mismatch it closes the socket with NO error
 *      frame — that silence is the invalid-code signal
 *   4. phone → browser, plaintext: {type:"PEACH_HELLO_ACK", ephemeralPublicKeyHex:"<64 hex>"}
 *      (that key is the phone's Noise NK static `s`)
 *   5. Noise NK, phone as RESPONDER, over binary frames
 *   6. one encrypted PAIR_IDENTITY frame each way, BROWSER FIRST
 *   7. both derive the same 6-digit SAS from the handshake hash, and BOTH gate
 *      persistence on the user confirming it
 *
 * Step 7 is why [PairingSession] hands the SAS out and waits: the older
 * [MdnsPairingServer] wrote the peer into [PairedDeviceStore] as soon as the
 * handshake completed, before any human had compared anything, which makes the
 * SAS decorative. Nothing is persisted here until [PairingSession.confirm] is
 * called with true.
 */
sealed interface PairingEvent {
    /** The listener is up; show [code] to the user. */
    data class Listening(val code: String, val ip: String, val port: Int) : PairingEvent

    /** A peer completed the handshake. Show [sas]; nothing is stored yet. */
    data class AwaitingConfirmation(val sas: String, val deviceName: String, val session: PairingSession) : PairingEvent

    data class Paired(val deviceName: String) : PairingEvent

    /** Includes the silent-close case: a peer presented a bad token. */
    data class Rejected(val reason: String) : PairingEvent
}

class PairingSession internal constructor() {
    private val decision = CompletableDeferred<Boolean>()
    fun confirm(matches: Boolean) { decision.complete(matches) }
    internal suspend fun await(): Boolean = decision.await()
}

@Singleton
class CodePairingServer @Inject constructor(
    private val deviceStore: PairedDeviceStore,
) {
    private var serverSocket: ServerSocket? = null

    /**
     * Bind, publish a code, and accept pairing attempts until the flow is
     * cancelled. One token per call: after a successful pairing the caller
     * should stop this flow, which closes the pairing window while leaving any
     * long-lived sync listener alone.
     */
    fun start(
        scope: CoroutineScope,
        ourStaticPrivateKey: ByteArray,
        ourStaticPublicKey: ByteArray,
        localIdentity: LocalChirpIdentity,
    ): Flow<PairingEvent> = callbackFlow {
        val ip = localIpv4()
        if (ip == null) {
            trySend(PairingEvent.Rejected("no_local_network"))
            close()
            return@callbackFlow
        }

        val socket = withContext(Dispatchers.IO) { ServerSocket(0) }
        serverSocket = socket
        val token = PairingCode.newToken()
        val tokenHex = token.joinToString("") { "%02x".format(it) }

        trySend(PairingEvent.Listening(PairingCode.encode(ip, socket.localPort, token), ip, socket.localPort))

        val acceptJob = scope.launch(Dispatchers.IO) {
            while (isActive && !socket.isClosed) {
                val client = runCatching { socket.accept() }.getOrNull() ?: break
                launch {
                    trySend(
                        handleClient(client, tokenHex, ourStaticPrivateKey, ourStaticPublicKey, localIdentity),
                    )
                }
            }
        }

        awaitClose {
            acceptJob.cancel()
            stop()
        }
    }

    fun stop() {
        runCatching { serverSocket?.close() }
        serverSocket = null
    }

    private suspend fun handleClient(
        client: Socket,
        expectedTokenHex: String,
        ourStaticPrivateKey: ByteArray,
        ourStaticPublicKey: ByteArray,
        localIdentity: LocalChirpIdentity,
    ): PairingEvent = withContext(Dispatchers.IO) {
        val ws = WebSocketConnection(client, client.getInputStream(), client.getOutputStream())
        try {
            if (!ws.accept()) {
                ws.closeSilently()
                return@withContext PairingEvent.Rejected("not_a_websocket")
            }

            val hello = ws.receive() as? WsMessage.Text
                ?: return@withContext rejectSilently(ws, "no_hello")
            val helloJson = runCatching { JSONObject(hello.text) }.getOrNull()
                ?: return@withContext rejectSilently(ws, "bad_hello")
            if (helloJson.optString("type") != "PEACH_HELLO") {
                return@withContext rejectSilently(ws, "bad_hello")
            }
            if (!constantTimeEquals(helloJson.optString("token"), expectedTokenHex)) {
                return@withContext rejectSilently(ws, "bad_token")
            }

            ws.sendText(
                JSONObject()
                    .put("type", "PEACH_HELLO_ACK")
                    .put("ephemeralPublicKeyHex", ourStaticPublicKey.joinToString("") { "%02x".format(it) })
                    .toString(),
            )

            val session = noiseNKResponder(
                ourStaticPrivateKey,
                ourStaticPublicKey,
                { frame -> ws.send(frame) },
                { (ws.receive() as? WsMessage.Binary)?.bytes ?: throw IllegalStateException("peer closed") },
            )

            // Browser first, then us — matching the contract's ordering.
            val theirIdentity = parsePairIdentityFrame(
                session.decrypt((ws.receive() as? WsMessage.Binary)?.bytes ?: throw IllegalStateException("peer closed")),
            )
            ws.send(session.encrypt(encodePairIdentityFrame(localIdentity)))

            val sas = SASDerivation.derive(session.handshakeHash)
            val device = pairIdentityToPairedDevice(theirIdentity)
            val pending = PairingSession()
            val name = device.deviceName?.takeIf { it.isNotBlank() } ?: "Unknown device"

            // Hand the SAS to the UI and block here. Persisting before the user
            // has compared the digits would make the comparison meaningless.
            val confirmation = PairingEvent.AwaitingConfirmation(sas, name, pending)
            confirmationSink?.invoke(confirmation)
            if (!pending.await()) {
                ws.closeSilently()
                return@withContext PairingEvent.Rejected("sas_mismatch")
            }

            deviceStore.upsert(device)
            PairingEvent.Paired(name)
        } catch (e: Exception) {
            ws.closeSilently()
            PairingEvent.Rejected(e.message ?: "pairing_failed")
        } finally {
            runCatching { client.close() }
        }
    }

    /**
     * Set by the ViewModel so the SAS reaches the UI while [handleClient] is
     * still blocked inside the connection.
     */
    var confirmationSink: ((PairingEvent.AwaitingConfirmation) -> Unit)? = null

    private fun rejectSilently(ws: WebSocketConnection, reason: String): PairingEvent {
        ws.closeSilently()
        return PairingEvent.Rejected(reason)
    }

    /** The token is low-value, but comparing it in constant time costs nothing. */
    private fun constantTimeEquals(a: String, b: String): Boolean {
        if (a.length != b.length) return false
        var diff = 0
        for (i in a.indices) diff = diff or (a[i].code xor b[i].code)
        return diff == 0
    }

    /**
     * The address to put in the code. Picks a site-local IPv4 on an up,
     * non-loopback interface — a VPN or tethering interface would produce a
     * code the browser cannot dial.
     */
    private fun localIpv4(): String? =
        runCatching {
            NetworkInterface.getNetworkInterfaces().toList()
                .filter { it.isUp && !it.isLoopback && !it.isVirtual }
                .flatMap { it.inetAddresses.toList() }
                .filterIsInstance<Inet4Address>()
                .firstOrNull { it.isSiteLocalAddress }
                ?.hostAddress
        }.getOrNull()
}