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 java.net.InetSocketAddress
import java.net.Socket
import java.net.SocketTimeoutException
import javax.inject.Inject
import javax.inject.Singleton

/**
 * C7 TCP + mDNS pairing/sync CLIENT (initiator side). Discovers `_peach._tcp`
 * peers, dials one, and runs the REAL initiator pipeline:
 *   1. Noise NK initiator handshake ([noiseNKInitiator]) — pins the responder's
 *      static public key from the mDNS TXT "pk" record.
 *   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 ──────────────────────────────────────────
 *  - `noiseNK.initiateHandshake(stream, out, remotePublicKey)` does NOT exist.
 *    Real: `noiseNKInitiator(responderStaticPublicKey, send, receive)` →
 *    [NoiseNKSession]. One bidirectional cipher; no `.sendCipher/.recvCipher`.
 *  - `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 NK initiator handshake, exchange identities and
     * derive the SAS, then run one initiator sync round.
     *
     * Caller supplies the local static identity (Noise initiator does not have a
     * static key in NK, but our PAIR_IDENTITY + the sync layer still carry our
     * long-term identity) and a vault adapter factory.
     */
    suspend fun connectAndSync(
        peer: DiscoveredPeer,
        localIdentity: LocalChirpIdentity,
        vaultFactory: () -> ChirpSyncVaultAdapter?,
        connectTimeoutMs: Int = 15_000,
    ): PairingResult = withContext(Dispatchers.IO) {
        val socket = Socket()
        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 session = try {
                noiseNKInitiator(peer.remoteStaticPublicKey, transport::send, transport::receive)
            } catch (e: Throwable) {
                transport.close()
                return@withContext PairingResult.Error(PairingError.HANDSHAKE_FAILED)
            }

            // PAIR_IDENTITY: initiator writes first, then reads (mirrors
            // ChirpPairing.runPairingInitiator).
            transport.send(session.encrypt(encodePairIdentityFrame(localIdentity)))
            val theirIdentity = parsePairIdentityFrame(session.decrypt(transport.receive()))

            val sas = SASDerivation.derive(session.handshakeHash)
            val device = pairIdentityToPairedDevice(theirIdentity)
            deviceStore.upsert(device)

            // Encrypted vault sync — real initiator round.
            val vault = vaultFactory()
            if (vault != null) {
                val syncSession = chirpSessionOver(session, transport)
                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)
        }
    }

    /**
     * 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,
        vaultFactory: () -> ChirpSyncVaultAdapter?,
        discoveryTimeoutMs: Long = 5_000L,
        cached: DiscoveredPeer? = null,
        cachedTimeoutMs: Int = 2_000,
    ): PairingResult {
        if (cached != null) {
            val viaCache = connectAndSync(cached, localIdentity, vaultFactory, connectTimeoutMs = cachedTimeoutMs)
            if (viaCache is PairingResult.Success) return viaCache
        }
        val peer = withTimeoutOrNull(discoveryTimeoutMs) {
            discoverPeers().first { it.deviceId == deviceId }
        } ?: return PairingResult.Error(PairingError.NOT_FOUND)
        return connectAndSync(peer, localIdentity, vaultFactory)
    }

    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,
    /** Noise NK responder static public key, parsed from TXT "pk". */
    val remoteStaticPublicKey: ByteArray,
)

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,
}