package com.peachpasswords.sync

import android.content.Context
import android.net.nsd.NsdManager
import android.net.nsd.NsdServiceInfo
import android.net.wifi.WifiManager
import dagger.hilt.android.qualifiers.ApplicationContext
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.callbackFlow
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeoutOrNull
import org.json.JSONObject
import java.net.InetSocketAddress
import java.net.Socket
import java.net.SocketTimeoutException
import javax.inject.Inject
import javax.inject.Singleton

/**
 * C7 TCP + mDNS reconnect CLIENT (initiator side). Discovers `_peach._tcp`
 * peers, dials one, and runs the REAL initiator pipeline:
 *   1. An opaque, rotating PEACH_RECONNECT capability lets the responder choose
 *      an already-paired local record, then Noise KK authenticates both stored
 *      static keys.
 *   2. PAIR_IDENTITY exchange (initiator writes first, then reads), matching
 *      ChirpPairing.runPairingInitiator ordering.
 *   3. SAS via the REAL [SASDerivation.derive] over the Noise handshake hash.
 *   4. Encrypted vault sync via the REAL [runChirpSyncInitiator] over a
 *      [ChirpSyncSession].
 *
 * ── Reconciliation with the brief ──────────────────────────────────────────
 *  - Reconnect is strictly KK: both local static keys are inputs and a failure
 *    never negotiates or retries NK. The returned [NoiseNKSession] interface is
 *    transport-neutral legacy naming, not the active handshake pattern.
 *  - `syncAdapter.buildRequest/applyResponse` do NOT exist. Real round driver is
 *    [runChirpSyncInitiator] against a [ChirpSyncVaultAdapter].
 *  - `peer.serviceInfo.host/.port`: on API < 34, [NsdServiceInfo.getHost] is
 *    deprecated; on the supported floor it still returns the resolved
 *    InetAddress, which we use for the dial. The resolved port is authoritative.
 */
@Singleton
class PairingDiscoveryClient @Inject constructor(
    @ApplicationContext private val ctx: Context,
    private val deviceStore: PairedDeviceStore,
) {
    /**
     * Discover nearby Peach instances via mDNS. Each found service is resolved,
     * its TXT records parsed into a [DiscoveredPeer], and emitted. The flow runs
     * until the collector cancels (e.g. via kotlinx `.timeout(...)` or `first`).
     */
    fun discoverPeers(): Flow<DiscoveredPeer> = callbackFlow {
        val manager = ctx.getSystemService(NsdManager::class.java)
            ?: run { close(); return@callbackFlow }

        // mDNS/NSD relies on Wi-Fi multicast, which Android drops by default to save
        // power. Hold a MulticastLock for the lifetime of discovery so service
        // announcements actually reach us (requires CHANGE_WIFI_MULTICAST_STATE).
        val multicastLock = (ctx.getSystemService(Context.WIFI_SERVICE) as? WifiManager)
            ?.createMulticastLock("peach-mdns")
            ?.apply { setReferenceCounted(false); runCatching { acquire() } }

        val resolveListener = object : NsdManager.ResolveListener {
            override fun onServiceResolved(resolved: NsdServiceInfo) {
                parsePeer(resolved)?.let { trySend(it) }
            }
            override fun onResolveFailed(info: NsdServiceInfo, errorCode: Int) {}
        }

        val discoveryListener = object : NsdManager.DiscoveryListener {
            override fun onServiceFound(info: NsdServiceInfo) {
                if (info.serviceType.trimEnd('.') == PEACH_SERVICE_TYPE) {
                    runCatching { manager.resolveService(info, resolveListener) }
                }
            }
            override fun onServiceLost(info: NsdServiceInfo) {}
            override fun onDiscoveryStarted(serviceType: String) {}
            override fun onDiscoveryStopped(serviceType: String) {}
            override fun onStartDiscoveryFailed(serviceType: String, errorCode: Int) { close() }
            override fun onStopDiscoveryFailed(serviceType: String, errorCode: Int) {}
        }

        manager.discoverServices(PEACH_SERVICE_TYPE, NsdManager.PROTOCOL_DNS_SD, discoveryListener)
        awaitClose {
            runCatching { manager.stopServiceDiscovery(discoveryListener) }
            multicastLock?.let { if (it.isHeld) runCatching { it.release() } }
        }
    }

    /**
     * Dial [peer], run the Noise KK initiator handshake, exchange identities and
     * derive the SAS, then run one initiator sync round.
     *
     * Caller supplies the local static identity and private key required by KK,
     * plus a vault adapter factory. The pre-handshake selector is an opaque
     * rendezvous capability; it never discloses either stored public key.
     */
    private suspend fun connectAndSync(
        peer: DiscoveredPeer,
        localIdentity: LocalChirpIdentity,
        localStaticPrivateKey: ByteArray,
        rendezvousId: String,
        vaultFactory: () -> ChirpSyncVaultAdapter?,
        connectTimeoutMs: Int = 15_000,
    ): PairingResult = withContext(Dispatchers.IO) {
        val socket = Socket()
        var noiseSession: NoiseNKSession? = null
        try {
            socket.connect(InetSocketAddress(peer.host, peer.port), connectTimeoutMs)
            socket.soTimeout = connectTimeoutMs

            val transport = SocketFrameIO(
                input = socket.getInputStream(),
                output = socket.getOutputStream(),
                onClose = { withContext(Dispatchers.IO) { runCatching { socket.close() } } },
            )

            val localPublicKey = hexToPubkey(localIdentity.longTermPublicKeyHex)
            transport.send(
                JSONObject()
                    .put("type", "PEACH_RECONNECT")
                    .put("rendezvous_id", rendezvousId)
                    .toString()
                    .toByteArray(Charsets.UTF_8),
            )
            val session = try {
                noiseKKInitiator(
                    localStaticPrivateKey,
                    localPublicKey,
                    peer.remoteStaticPublicKey,
                    transport::send,
                    transport::receive,
                )
            } catch (e: Throwable) {
                transport.close()
                return@withContext PairingResult.Error(PairingError.HANDSHAKE_FAILED)
            }
            noiseSession = session

            // PAIR_IDENTITY: initiator writes first, then reads (mirrors
            // ChirpPairing.runPairingInitiator).
            transport.send(session.encrypt(SyncPadding.pad(encodePairIdentityFrame(localIdentity))))
            val theirIdentity = parsePairIdentityFrame(
                SyncPadding.unpad(session.decrypt(transport.receive())),
            )
            require(theirIdentity.deviceId == peer.deviceId) { "peer device id changed after KK" }
            require(
                CryptoCompare.hex(
                    theirIdentity.longTermPublicKeyHex,
                    pubkeyToHex(peer.remoteStaticPublicKey),
                ),
            ) { "peer identity does not match authenticated KK key" }

            val sas = SASDerivation.derive(session.handshakeHash)
            val incoming = pairIdentityToPairedDevice(theirIdentity)
            val existing = deviceStore.getByPubkey(incoming.longTermPublicKeyHex.orEmpty())
                ?: deviceStore.get(incoming.deviceId)
            val device = incoming.copy(
                rendezvousSecretHex = existing?.rendezvousSecretHex,
                relayUrl = existing?.relayUrl ?: incoming.relayUrl,
            )
            deviceStore.upsert(device)

            // Encrypted vault sync — real initiator round.
            val vault = vaultFactory()
            if (vault != null) {
                vault.bindAuthenticatedPeer(device)
                val syncSession = chirpSessionOver(session, transport, peer.remoteStaticPublicKey)
                val outcome = runChirpSyncInitiator(syncSession, vault)
                if (outcome.status == SyncOutcomeStatus.COMPLETED) {
                    deviceStore.updateLastSync(device.deviceId)
                }
            }

            transport.close()
            PairingResult.Success(device, sas)
        } catch (e: SocketTimeoutException) {
            runCatching { socket.close() }
            PairingResult.Error(PairingError.TIMEOUT)
        } catch (e: Throwable) {
            runCatching { socket.close() }
            PairingResult.Error(PairingError.IO_ERROR)
        } finally {
            noiseSession?.close()
        }
    }

    /**
     * Dial a known peer: the cached address first, discovery only as a fallback.
     *
     * [cached] is the peer's last-known endpoint, refreshed by the listener on
     * every inbound reconnect. Trying it first is the "silent reconnect" the
     * sync design asks for (§6) — this used to go straight to mDNS and report
     * NOT_FOUND after a 5s timeout, so a peer that was perfectly reachable at
     * its cached address failed to sync whenever discovery was unavailable,
     * which on a home network is often.
     *
     * The cached attempt gets a short timeout: a stale address is the normal
     * case after a network change, and it must not delay the fallback.
     */
    suspend fun syncWithKnownDevice(
        deviceId: String,
        localIdentity: LocalChirpIdentity,
        localStaticPrivateKey: ByteArray,
        vaultFactory: () -> ChirpSyncVaultAdapter?,
        discoveryTimeoutMs: Long = 5_000L,
        cached: DiscoveredPeer? = null,
        cachedTimeoutMs: Int = 2_000,
    ): PairingResult {
        val paired = deviceStore.get(deviceId)
            ?: return PairingResult.Error(PairingError.NOT_FOUND)
        val secret = paired.rendezvousSecretHex
            ?.let { runCatching { Rendezvous.secretFromHex(it) }.getOrNull() }
            ?: return PairingResult.Error(PairingError.HANDSHAKE_FAILED)
        val rendezvousIds = Rendezvous.dialingIds(
            secret,
            System.currentTimeMillis() / 1000L,
        )
        if (cached != null) {
            val pinnedCached = pinDiscoveredPeerToPairedState(paired, cached)
            for (rendezvousId in rendezvousIds) {
                val viaCache = connectAndSync(
                    pinnedCached,
                    localIdentity,
                    localStaticPrivateKey,
                    rendezvousId,
                    vaultFactory,
                    connectTimeoutMs = cachedTimeoutMs,
                )
                if (viaCache is PairingResult.Success) return viaCache
            }
        }
        val discovered = withTimeoutOrNull(discoveryTimeoutMs) {
            discoverPeers().first { it.deviceId == deviceId }
        } ?: return PairingResult.Error(PairingError.NOT_FOUND)
        val peer = pinDiscoveredPeerToPairedState(paired, discovered)
        var last: PairingResult = PairingResult.Error(PairingError.HANDSHAKE_FAILED)
        for (rendezvousId in rendezvousIds) {
            last = connectAndSync(
                peer,
                localIdentity,
                localStaticPrivateKey,
                rendezvousId,
                vaultFactory,
            )
            if (last is PairingResult.Success) return last
        }
        return last
    }

    private fun parsePeer(info: NsdServiceInfo): DiscoveredPeer? {
        val attrs = info.attributes ?: return null
        val pkHex = attrs["pk"]?.let { String(it, Charsets.UTF_8) } ?: return null
        if (!pkHex.matches(Regex("^[0-9a-f]{64}$"))) return null
        val deviceId = attrs["id"]?.let { String(it, Charsets.UTF_8) } ?: return null
        @Suppress("DEPRECATION")
        val host = info.host?.hostAddress ?: return null
        return DiscoveredPeer(
            deviceId = deviceId,
            deviceName = info.serviceName,
            host = host,
            port = info.port,
            remoteStaticPublicKey = hexToPubkey(pkHex),
        )
    }

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

data class DiscoveredPeer(
    val deviceId: String,
    val deviceName: String,
    val host: String,
    val port: Int,
    /** Paired responder static public key, parsed from TXT "pk" and pinned in KK. */
    val remoteStaticPublicKey: ByteArray,
)

/**
 * Discovery supplies reachability only. KK's responder key always comes from
 * paired local state; a spoofed mDNS TXT record can redirect or deny service,
 * but cannot choose the key the initiator authenticates.
 */
internal fun pinDiscoveredPeerToPairedState(
    paired: PairedDevice,
    discovered: DiscoveredPeer,
): DiscoveredPeer {
    require(discovered.deviceId == paired.deviceId) { "discovered peer id does not match paired state" }
    val expectedKey = paired.longTermPublicKeyHex
        ?: throw IllegalArgumentException("paired peer missing static key")
    return discovered.copy(remoteStaticPublicKey = hexToPubkey(expectedKey))
}

sealed class PairingResult {
    data class Success(val device: PairedDevice, val sasCode: String) : PairingResult()
    data class Error(val reason: PairingError) : PairingResult()
}

enum class PairingError {
    NOT_FOUND,
    HANDSHAKE_FAILED,
    VERSION_MISMATCH,
    ALREADY_PAIRED,
    TIMEOUT,
    IO_ERROR,
}