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

    /** Discover then dial the first peer whose deviceId matches [deviceId]. */
    suspend fun syncWithKnownDevice(
        deviceId: String,
        localIdentity: LocalChirpIdentity,
        vaultFactory: () -> ChirpSyncVaultAdapter?,
        discoveryTimeoutMs: Long = 5_000L,
    ): PairingResult {
        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,
}