package com.peachpasswords.sync

import android.content.Context
import android.net.nsd.NsdManager
import android.net.nsd.NsdServiceInfo
import android.os.Build
import dagger.hilt.android.qualifiers.ApplicationContext
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.ServerSocket
import java.net.Socket
import java.nio.charset.StandardCharsets
import javax.inject.Inject
import javax.inject.Singleton

/**
 * C7 TCP + mDNS reconnect SERVER (responder side).
 *
 * Listens on an OS-assigned TCP port, advertises itself over mDNS as
 * `_peach._tcp`, and for each inbound connection runs the REAL responder
 * pipeline:
 *   1. Resolve PEACH_RECONNECT against paired local state, then run Noise KK
 *      over length-prefixed frames ([SocketFrameIO]).
 *   2. PAIR_IDENTITY frame exchange — encoded/decoded with the SAME private
 *      helpers ChirpPairing uses; since those are file-private there, this draft
 *      re-derives the identity frame via [LocalChirpIdentity] and a minimal
 *      JSON, and parses the peer's via [parsePairIdentityFrame]. See
 *      C7-INTEGRATION.md: the cleanest integration promotes ChirpPairing's
 *      encode/decodePairingFrame to internal and deletes these shims.
 *   3. Encrypted vault sync via the REAL [runChirpSyncResponder] over a
 *      [ChirpSyncSession] backed by the Noise session.
 *
 * Reconnect is strictly KK. The opaque selector capability is resolved against
 * local paired state before the handshake, both static keys are authenticated
 * by the named pattern, and failure never falls back to first-pair NK.
 *  - The brief calls `syncAdapter.buildResponse/applyRequest`. NO SUCH METHODS.
 *    The real per-round protocol is [runChirpSyncResponder]; the adapter is a
 *    [ChirpSyncVaultAdapter] ([VaultSyncAdapter]) with getEncryptedEntries/
 *    applyEntries/etc. We drive the whole round, not single request/response.
 */
@Singleton
class MdnsPairingServer @Inject constructor(
    private val deviceStore: PairedDeviceStore,
    @ApplicationContext private val ctx: Context,
) {
    private var serverSocket: ServerSocket? = null
    private var nsdManager: NsdManager? = null
    private var registrationListener: NsdManager.RegistrationListener? = null

    /**
     * Start listening + advertise + accept. Caller supplies:
     *  - [ourStaticPrivateKey]/[ourStaticPublicKey]: the device's long-term sync
     *    identity (from [deriveSyncIdentityKeypair]) — the KK responder static
     *    key, also published in mDNS TXT "pk".
     *  - [vaultFactory]: builds a fresh [ChirpSyncVaultAdapter] per connection
     *    (the vault may lock/unlock between connections).
     *  - [localIdentity]: our PAIR_IDENTITY payload.
     */
    fun start(
        scope: CoroutineScope,
        ourStaticPrivateKey: ByteArray,
        ourStaticPublicKey: ByteArray,
        localIdentity: LocalChirpIdentity,
        vaultFactory: () -> ChirpSyncVaultAdapter?,
    ): Flow<PairingServerEvent> = callbackFlow {
        val socket = withContext(Dispatchers.IO) { ServerSocket(0) }
        serverSocket = socket
        val port = socket.localPort

        registerMdns(port, ourStaticPublicKey, localIdentity.deviceId)
        trySend(PairingServerEvent.Started(port))

        val acceptJob = scope.launch(Dispatchers.IO) {
            while (isActive && !socket.isClosed) {
                val client = try {
                    socket.accept()
                } catch (_: Throwable) {
                    break
                }
                launch {
                    val event = handleClient(client, ourStaticPrivateKey, ourStaticPublicKey, localIdentity, vaultFactory)
                    trySend(event)
                }
            }
        }

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

    fun stop() {
        registrationListener?.let { l -> runCatching { nsdManager?.unregisterService(l) } }
        registrationListener = null
        runCatching { serverSocket?.close() }
        serverSocket = null
    }

    private suspend fun handleClient(
        socket: Socket,
        ourStaticPrivateKey: ByteArray,
        ourStaticPublicKey: ByteArray,
        localIdentity: LocalChirpIdentity,
        vaultFactory: () -> ChirpSyncVaultAdapter?,
    ): PairingServerEvent {
        val transport = SocketFrameIO(
            input = socket.getInputStream(),
            output = socket.getOutputStream(),
            onClose = { withContext(Dispatchers.IO) { runCatching { socket.close() } } },
        )
        var noiseSession: NoiseNKSession? = null
        return try {
            val selector = JSONObject(String(transport.receive(), Charsets.UTF_8))
            require(selector.optString("type") == "PEACH_RECONNECT") { "expected PEACH_RECONNECT" }
            val rendezvousId = selector.optString("rendezvous_id")
            val known = Rendezvous.resolvePeer(deviceStore.getAll(), rendezvousId)
                ?: throw IllegalArgumentException("unknown reconnect peer")
            val initiatorKeyHex = known.longTermPublicKeyHex
                ?: throw IllegalArgumentException("known peer missing static key")
            val initiatorKey = hexToPubkey(initiatorKeyHex)

            val session = noiseKKResponder(
                ourStaticPrivateKey,
                ourStaticPublicKey,
                initiatorKey,
                transport::send,
                transport::receive,
            )
            noiseSession = session

            // 2. PAIR_IDENTITY exchange (responder reads first, then answers) —
            //    mirrors ChirpPairing.runPairingResponder ordering.
            val theirIdentity = parsePairIdentityFrame(
                SyncPadding.unpad(session.decrypt(transport.receive())),
            )
            require(theirIdentity.deviceId == known.deviceId) { "peer device id changed after KK" }
            require(
                CryptoCompare.hex(theirIdentity.longTermPublicKeyHex, initiatorKeyHex),
            ) { "peer identity does not match authenticated KK key" }
            transport.send(session.encrypt(SyncPadding.pad(encodePairIdentityFrame(localIdentity))))

            val sas = SASDerivation.derive(session.handshakeHash)
            // 3. Encrypted vault sync — real responder round over a ChirpSyncSession.
            val vault = vaultFactory()
            if (vault != null) {
                vault.bindAuthenticatedPeer(known)
                val syncSession = chirpSessionOver(session, transport, initiatorKey)
                val outcome = runChirpSyncResponder(syncSession, vault)
                if (outcome.status == SyncOutcomeStatus.COMPLETED) {
                    deviceStore.updateLastSync(known.deviceId)
                }
            }
            transport.close()
            PairingServerEvent.SyncComplete(known.deviceId, sas)
        } catch (e: Throwable) {
            transport.close()
            PairingServerEvent.Error(e.message ?: "pairing_server_error")
        } finally {
            noiseSession?.close()
        }
    }

    private fun registerMdns(port: Int, staticPub: ByteArray, deviceId: String) {
        val manager = ctx.getSystemService(NsdManager::class.java) ?: return
        nsdManager = manager
        val info = NsdServiceInfo().apply {
            serviceName = buildServiceName(deviceId)
            serviceType = PEACH_SERVICE_TYPE
            this.port = port
            // NsdServiceInfo.setAttribute is API 21+; minSdk is 26.
            setAttribute("v", PROTOCOL_VERSION)
            setAttribute("pk", pubkeyToHex(staticPub))
            setAttribute("id", deviceId)
        }
        val listener = object : NsdManager.RegistrationListener {
            override fun onServiceRegistered(info: NsdServiceInfo) {}
            override fun onRegistrationFailed(info: NsdServiceInfo, errorCode: Int) {}
            override fun onServiceUnregistered(info: NsdServiceInfo) {}
            override fun onUnregistrationFailed(info: NsdServiceInfo, errorCode: Int) {}
        }
        registrationListener = listener
        manager.registerService(info, NsdManager.PROTOCOL_DNS_SD, listener)
    }

    private fun buildServiceName(deviceId: String): String {
        val shortId = deviceId.take(8)
        val model = Build.MODEL.replace(Regex("[^A-Za-z0-9-]"), "")
        // mDNS service-instance names are capped at 63 bytes; reserve room for
        // the "Peach-" prefix and the "-${shortId}" suffix.
        val maxModel = 63 - "Peach-".length - 1 - shortId.length
        val trimmedModel = if (model.length > maxModel && maxModel > 0) model.take(maxModel) else model
        return "Peach-$trimmedModel-$shortId".take(63)
    }

    private companion object {
        const val PEACH_SERVICE_TYPE = "_peach._tcp"
        const val PROTOCOL_VERSION = "1"
    }
}

sealed class PairingServerEvent {
    data class Started(val port: Int) : PairingServerEvent()
    data class SyncComplete(val deviceId: String, val sas: String) : PairingServerEvent()
    data class Error(val message: String) : PairingServerEvent()
    data object Stopped : PairingServerEvent()
}

/* ── PAIR_IDENTITY shims (see C7-INTEGRATION.md item: promote ChirpPairing's
   private encode/decodePairingFrame to internal and drop these) ───────────── */

internal data class PairIdentity(
    val deviceId: String,
    val deviceName: String,
    val deviceType: DeviceType,
    val longTermPublicKeyHex: String,
    val signingPublicKeyHex: String?,
    val wsPort: Int?,
    val relayUrl: String?,
)

internal fun encodePairIdentityFrame(id: LocalChirpIdentity): ByteArray {
    val sb = StringBuilder()
    sb.append('{')
    sb.append("\"type\":\"PAIR_IDENTITY\"")
    sb.append(",\"device_id\":\"").append(jsonEscape(id.deviceId)).append('"')
    sb.append(",\"device_name\":\"").append(jsonEscape(id.deviceName)).append('"')
    sb.append(",\"device_type\":\"").append(id.deviceType.wireValue).append('"')
    sb.append(",\"long_term_public_key_hex\":\"").append(id.longTermPublicKeyHex).append('"')
    sb.append(",\"signing_public_key_hex\":\"").append(id.signingPublicKeyHex).append('"')
    id.wsPort?.let { sb.append(",\"ws_port\":").append(it) }
    sb.append(",\"relay_url\":\"").append(jsonEscape(id.relayUrl)).append('"')
    sb.append('}')
    return sb.toString().toByteArray(StandardCharsets.UTF_8)
}

internal fun parsePairIdentityFrame(bytes: ByteArray): PairIdentity {
    val map = miniJsonObject(String(bytes, StandardCharsets.UTF_8))
    require(map["type"] == "PAIR_IDENTITY") { "expected PAIR_IDENTITY frame" }
    val deviceType = when (map["device_type"]) {
        "extension" -> DeviceType.EXTENSION
        "android" -> DeviceType.ANDROID
        "ios" -> DeviceType.IOS
        else -> error("device_type missing")
    }
    val ltpk = map["long_term_public_key_hex"] ?: error("long_term_public_key_hex missing")
    require(ltpk.matches(Regex("^[0-9a-f]{64}$"))) { "long_term_public_key_hex not 32-byte hex" }
    return PairIdentity(
        deviceId = map["device_id"] ?: error("device_id missing"),
        deviceName = map["device_name"] ?: error("device_name missing"),
        deviceType = deviceType,
        longTermPublicKeyHex = ltpk,
        signingPublicKeyHex = map["signing_public_key_hex"]?.takeIf { it.isNotBlank() },
        wsPort = map["ws_port"]?.toIntOrNull(),
        relayUrl = map["relay_url"]?.let(Rendezvous::normalizeRelayUrl),
    )
}

internal fun pairIdentityToPairedDevice(id: PairIdentity): PairedDevice = PairedDevice(
    deviceId = id.deviceId,
    deviceName = id.deviceName,
    deviceType = id.deviceType.wireValue,
    longTermPublicKeyHex = id.longTermPublicKeyHex,
    signingPublicKeyHex = id.signingPublicKeyHex,
    trustSource = "direct",
    wsPort = id.wsPort,
    relayUrl = id.relayUrl,
    pairedAt = System.currentTimeMillis(),
    lastSyncedAt = System.currentTimeMillis(),
)

/**
 * Build a [ChirpSyncSession] over an established Noise session + frame
 * transport, so the REAL [runChirpSyncInitiator]/[runChirpSyncResponder] can run
 * the encrypted vault exchange. Send/receive splice the Noise cipher onto the
 * length-prefixed frame layer; `remoteStaticPublicKey` is informational here.
 */
internal fun chirpSessionOver(
    session: NoiseNKSession,
    transport: SocketFrameIO,
    remoteStaticPublicKey: ByteArray = ByteArray(0),
): ChirpSyncSession =
    ChirpSyncSession(
        remoteStaticPublicKey = remoteStaticPublicKey.copyOf(),
        // ChirpSyncSession.send is non-suspending;
        // SocketFrameIO.sendBlocking writes a frame synchronously.
        send = { plaintext -> transport.sendBlocking(session.encrypt(SyncPadding.pad(plaintext))) },
        receive = { SyncPadding.unpad(session.decrypt(transport.receive())) },
        close = {
            try {
                session.close()
            } finally {
                transport.close()
            }
        },
    )

private fun jsonEscape(value: String): String = buildString(value.length) {
    for (ch in value) when (ch) {
        '\\' -> append("\\\\"); '"' -> append("\\\"")
        '\n' -> append("\\n"); '\r' -> append("\\r"); '\t' -> append("\\t")
        else -> append(ch)
    }
}

private fun miniJsonObject(text: String): Map<String, String> {
    val trimmed = text.trim()
    require(trimmed.startsWith("{") && trimmed.endsWith("}")) { "expected JSON object" }
    val body = trimmed.substring(1, trimmed.length - 1)
    val out = linkedMapOf<String, String>()
    var i = 0
    while (i < body.length) {
        while (i < body.length && body[i].isWhitespace()) i++
        if (i >= body.length || body[i] != '"') break
        i++
        val keyStart = i
        while (i < body.length && body[i] != '"') i++
        val key = body.substring(keyStart, i); i++
        while (i < body.length && (body[i].isWhitespace() || body[i] == ':')) i++
        val value: String
        if (i < body.length && body[i] == '"') {
            i++
            val vStart = i
            while (i < body.length && body[i] != '"') i++
            value = body.substring(vStart, i); i++
        } else {
            val vStart = i
            while (i < body.length && body[i] != ',') i++
            value = body.substring(vStart, i).trim()
        }
        out[key] = value
        while (i < body.length && body[i].isWhitespace()) i++
        if (i < body.length && body[i] == ',') i++
    }
    return out
}