package com.peachpasswords.sync

import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.delay
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 BINARY 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 BINARY: {type:"PEACH_HELLO_ACK", ephemeralPublicKeyHex:"<64 hex>"}
 *      (that key is the phone's Noise NK static `s`)
 *
 * Every frame in this exchange is binary, including the two JSON ones. This
 * document said "text" until 2026-07-24 and the code followed it; the browser's
 * socket is binary end to end, so pairing failed for every correct code. Verify
 * claims like this against `pairing-dial.ts`, not against this comment.
 *   5. Noise NK, phone as RESPONDER, over binary frames
 *   6. one encrypted PAIR_IDENTITY frame each way, BROWSER FIRST
 *   7. both derive the same 8-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

    /** The code's window elapsed without a pairing; the user must start again. */
    data object Expired : 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(
    // The interface, not the concrete store: the store needs Android's
    // EncryptedSharedPreferences, which would make this whole class untestable
    // off device — and the token/window logic is exactly what wants testing.
    private val deviceStore: PairedDeviceRegistry,
) {
    private var serverSocket: ServerSocket? = null

    companion object {
        /** Doc says 5–10 minutes; the low end is plenty to walk to a laptop. */
        const val PAIRING_WINDOW_MS = 5L * 60 * 1000

        /** Interfaces whose address a LAN peer cannot route back to. */
        private val TUNNEL_PREFIXES = listOf("tun", "ppp", "ipsec", "rmnet", "rndis", "ap")
    }

    /**
     * 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,
        windowMs: Long = PAIRING_WINDOW_MS,
        adoptionSourceDeps: AdoptionSourceDeps? = null,
    ): 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))

        // The window is not a race against an attacker — there isn't a
        // realistic one on your own LAN for a single-use code. It just stops an
        // abandoned code from staying live indefinitely on a phone in a pocket.
        val expiryJob = scope.launch {
            delay(windowMs)
            trySend(PairingEvent.Expired)
            stop()
            this@callbackFlow.close()
        }

        val acceptJob = scope.launch(Dispatchers.IO) {
            while (isActive && !socket.isClosed) {
                val client = runCatching { socket.accept() }.getOrNull() ?: break
                launch {
                    val event = handleClient(
                        client,
                        tokenHex,
                        ourStaticPrivateKey,
                        ourStaticPublicKey,
                        localIdentity,
                        adoptionSourceDeps,
                    )
                    trySend(event)
                    // One code, one device. A rejected or failed attempt leaves
                    // the window open so the user can retry without generating
                    // a new code; only a real pairing closes it.
                    if (event is PairingEvent.Paired) {
                        expiryJob.cancel()
                        stop()
                        this@callbackFlow.close()
                    }
                }
            }
        }

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

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

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

            // The browser sends this as a BINARY frame — `openPairingSocket` in
            // peach-browser's pairing-dial.ts is a binary seam end to end
            // (`ws.binaryType = "arraybuffer"`, `send(bytes: Uint8Array)`).
            // Requiring text here rejected every real dial and closed the socket
            // without a word, which the browser reports as "invalid code" — so a
            // perfectly good code always failed. Text is still accepted so any
            // other caller keeps working.
            val helloText = when (val frame = ws.receive()) {
                is WsMessage.Binary -> String(frame.bytes, Charsets.UTF_8)
                is WsMessage.Text -> frame.text
                else -> return@withContext rejectSilently(ws, "no_hello")
            }
            val helloJson = runCatching { JSONObject(helloText) }.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")
            }
            val flow = helloJson.optString("flow").ifBlank { null }
            if (flow != null && flow != "adoption") {
                return@withContext rejectSilently(ws, "unsupported_flow")
            }

            // Binary, for the same reason: the browser reads every frame as
            // `new Uint8Array(ev.data as ArrayBuffer)`. A text frame hands it a
            // string, which converts to a ZERO-LENGTH array rather than
            // throwing — and a zero-length ACK is also read as an invalid code.
            ws.send(
                JSONObject()
                    .put("type", "PEACH_HELLO_ACK")
                    .put("ephemeralPublicKeyHex", ourStaticPublicKey.joinToString("") { "%02x".format(it) })
                    .toString()
                    .toByteArray(Charsets.UTF_8),
            )

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

            if (flow == "adoption") {
                val deps = adoptionSourceDeps
                    ?: return@withContext rejectSilently(ws, "adoption_unavailable")
                return@withContext handleAdoption(
                    ws,
                    session,
                    localIdentity,
                    deps,
                )
            }

            // E1: the initiator proves it holds the private half of the key it
            // will advertise. Bind that proof to this NK transcript before any
            // identity frame is accepted or SAS consent is requested.
            val authenticatedInitiatorKey = performResponderPairingAuth(
                session = session,
                ourPrivateKey = ourStaticPrivateKey,
                ourPublicKey = ourStaticPublicKey,
                sendEncryptedJson = { plaintext -> ws.send(session.encrypt(plaintext)) },
                receiveEncryptedJson = {
                    session.decrypt(
                        (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")),
            )
            require(
                java.security.MessageDigest.isEqual(
                    authenticatedInitiatorKey,
                    hexToPubkey(theirIdentity.longTermPublicKeyHex),
                ),
            ) { "PAIR_IDENTITY key does not match proven initiator key" }
            ws.send(session.encrypt(encodePairIdentityFrame(localIdentity)))

            val sas = SASDerivation.derive(session.handshakeHash)
            val device = Rendezvous.bindPairState(
                device = pairIdentityToPairedDevice(theirIdentity),
                handshakeHash = session.handshakeHash,
                // The device displaying the code chooses the mesh relay.
                relayUrl = localIdentity.relayUrl,
            )
            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")
            }

            // A direct code pairing means a physically present user confirmed
            // the SAS on both screens — the only thing allowed to lift a
            // terminal revocation (see PairedDeviceStore.revoke).
            deviceStore.clearRevocation(device.deviceId)
            deviceStore.upsert(device)
            PairingEvent.Paired(name)
        } catch (e: Exception) {
            ws.closeSilently()
            PairingEvent.Rejected(e.message ?: "pairing_failed")
        } finally {
            noiseSession?.close()
            runCatching { client.close() }
        }
    }

    private suspend fun handleAdoption(
        ws: WebSocketConnection,
        session: NoiseNKSession,
        localIdentity: LocalChirpIdentity,
        sourceDeps: AdoptionSourceDeps,
    ): PairingEvent {
        val channel = object : PairingTransport {
            override suspend fun send(msg: ByteArray) = ws.send(session.encrypt(msg))
            override suspend fun receive(): ByteArray = session.decrypt(
                (ws.receive() as? WsMessage.Binary)?.bytes
                    ?: throw IllegalStateException("peer closed"),
            )
            override suspend fun close() = ws.closeSilently()
        }
        val request = receiveFrame(channel)
        require(isAdoptRequest(request)) { "expected ADOPT_REQUEST" }
        request as AdoptionFrame.Request

        val pending = PairingSession()
        val name = request.deviceName.ifBlank { "New device" }
        confirmationSink?.invoke(
            PairingEvent.AwaitingConfirmation(
                SASDerivation.derive(session.handshakeHash),
                name,
                pending,
            ),
        )
        if (!pending.await()) return PairingEvent.Rejected("sas_mismatch")

        val pairState = AdoptionPairState(
            rendezvousSecretHex = Rendezvous.bytesToHex(
                Rendezvous.deriveSecret(session.handshakeHash),
            ),
            relayUrl = localIdentity.relayUrl,
        )
        val boundDeps = object : AdoptionSourceDeps {
            override val sourceIdentity = localIdentity
            override val pairState = pairState
            override suspend fun rootKeyHex() = sourceDeps.rootKeyHex()
            override suspend fun vaultId() = sourceDeps.vaultId()
            override suspend fun license() = sourceDeps.license()
            override suspend fun onConsent(request: AdoptionFrame.Request) = true
            override suspend fun upsertPairedDevice(device: PairedDevice) {
                sourceDeps.upsertPairedDevice(device)
            }
        }
        return if (grantAdoption(channel, request, boundDeps) != null) {
            PairingEvent.Paired(name)
        } else {
            PairingEvent.Rejected("adoption_rejected")
        }
    }

    /**
     * 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 {
        return CryptoCompare.utf8(a, b)
    }

    /**
     * The address to put in the code — it has to be the one the browser can
     * actually dial, which is the Wi-Fi address.
     *
     * Interface order is not meaningful, so taking the first site-local IPv4
     * can hand back a VPN (tun/ppp) or tethering (ap/rndis) address. Both are
     * site-local and both produce a code that looks fine and then times out
     * with nothing to debug, so Wi-Fi-ish names are preferred explicitly and
     * tunnels are excluded rather than merely deprioritised.
     */
    internal fun localIpv4(): String? = runCatching {
        val candidates = NetworkInterface.getNetworkInterfaces().toList()
            .filter { it.isUp && !it.isLoopback && !it.isVirtual }
            .filterNot { iface -> TUNNEL_PREFIXES.any { iface.name.startsWith(it) } }

        fun addressOf(iface: NetworkInterface): String? = iface.inetAddresses.toList()
            .filterIsInstance<Inet4Address>()
            .firstOrNull { it.isSiteLocalAddress }
            ?.hostAddress

        candidates.firstOrNull { it.name.startsWith("wlan") }?.let(::addressOf)
            ?: candidates.firstNotNullOfOrNull(::addressOf)
    }.getOrNull()
}