package com.peachpasswords.sync

import kotlinx.coroutines.async
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.delay
import kotlinx.coroutines.isActive
import kotlinx.coroutines.withTimeoutOrNull
import java.security.SecureRandom
import java.util.concurrent.atomic.AtomicBoolean
import org.bouncycastle.crypto.params.X25519PrivateKeyParameters

const val CHIRP_PAIRING_PORT: Int = 49142

interface PairingTransport {
    suspend fun send(msg: ByteArray)
    suspend fun receive(): ByteArray
    suspend fun close()
}

data class PairingEndpoint(
    val ip: String,
    val port: Int,
)

data class PairingServer(
    val transport: PairingTransport,
    val peerIp: String,
)

data class PairingHandshakeResult(
    val device: PairedDevice,
    val sas: String,
    /**
     * The peer's signed Pro license blob, if it advertised one. Activated
     * (signature-verified) ONLY after the SAS is confirmed. Never persisted.
     */
    val peerLicense: String? = null,
)

enum class PairingRole {
    INITIATOR,
    RESPONDER,
}

interface ChirpPairingDeps {
    suspend fun fetchLocalIdentity(): LocalChirpIdentity
    suspend fun completePairing(paired: PairedDevice)
    suspend fun confirmSAS(sas: String): Boolean
    suspend fun cacheLastKnownIp(pubkeyHex: String, ip: String)
    suspend fun getLocalSubnet(): String?
    suspend fun openPairingTransport(endpoint: PairingEndpoint): PairingTransport
    suspend fun openPairingServer(): PairingServer
    suspend fun openAudio(): AudioIO
    val chirpTimeoutMs: Long
    fun onProgress(progress: PairingProgress)

    /**
     * Activate a Pro license the (now SAS-confirmed) peer advertised. The
     * implementation verifies the signature and ignores it if already Pro.
     * Default no-op so existing callers/tests need no change.
     */
    suspend fun activatePeerLicense(license: String) {}
}

class ChirpPairingFailure(
    val phase: PairingPhase,
    message: String,
) : Error(message)

interface AudioIO {
    val sampleRate: Int
    suspend fun play(waveform: FloatArray)
    suspend fun pullSamples(): FloatArray
    suspend fun close()
}

interface PairingKeyPair {
    val publicKey: ByteArray
    val secretKey: ByteArray
}

private data class PairingX25519KeyPair(
    override val publicKey: ByteArray,
    override val secretKey: ByteArray,
) : PairingKeyPair

data class PairingAttempt(
    val ephemeralPrivateKey: ByteArray,
    val ephemeralPublicKey: ByteArray,
    val sessionNonce: ByteArray,
    val outgoingPayload: ChirpPayload,
)

fun createPairingAttempt(deviceType: Byte): PairingAttempt {
    val kp = generatePairingKeyPair()
    val nonce = ByteArray(16)
    secureRandom().nextBytes(nonce)
    return PairingAttempt(
        ephemeralPrivateKey = kp.secretKey,
        ephemeralPublicKey = kp.publicKey,
        sessionNonce = nonce,
        outgoingPayload = ChirpPayload(
            version = CHIRP_PROTOCOL_VERSION,
            deviceType = deviceType,
            ephemeralPublicKey = kp.publicKey,
            sessionNonce = nonce,
        ),
    )
}

data class ChirpExchangeResult(
    val remoteChirp: ChirpPayload,
)

data class ChirpExchangeOptions(
    val attempt: PairingAttempt,
    val timeoutMs: Long = DEFAULT_CHIRP_TIMEOUT_MS,
    val onDiscard: ((ChirpDecodeError.Kind) -> Unit)? = null,
    val onChirpPlayed: ((Int) -> Unit)? = null,
)

enum class ChirpExchangeErrorKind {
    TIMEOUT,
    VERSION_MISMATCH,
    MIC_DENIED,
    AUDIO_UNAVAILABLE,
}

class ChirpExchangeError(
    val kind: ChirpExchangeErrorKind,
    message: String? = null,
) : Error(message ?: kind.name.lowercase())

private const val CHIRP_LOOP_INTERVAL_MS = 4_000L
private const val DEFAULT_CHIRP_TIMEOUT_MS = 90_000L
private const val LISTEN_FRAME_SAMPLES = 4096

fun generatePairingKeyPair(): PairingKeyPair {
    val priv = X25519PrivateKeyParameters(secureRandom())
    return PairingX25519KeyPair(
        publicKey = priv.generatePublicKey().encoded,
        secretKey = priv.encoded,
    )
}

private fun secureRandom(): SecureRandom = SecureRandom()

private fun sameBytes(a: ByteArray, b: ByteArray): Boolean {
    if (a.size != b.size) return false
    for (i in a.indices) {
        if (a[i] != b[i]) return false
    }
    return true
}

/**
 * Run the two-chirp exchange against a provided audio interface.
 *
 * This is a direct port of the browser coordinator logic. The codec and
 * audio bridge are injectable so the protocol can be unit-tested before the
 * Android-native ggwave bridge lands.
 */
suspend fun runChirpExchange(
    io: AudioIO,
    opts: ChirpExchangeOptions,
    codecFactory: ChirpCodecFactory = ::createChirpCodec,
): ChirpExchangeResult {
    return coroutineScope {
        val codec = codecFactory(
            ChirpCodecOptions(
                sampleRate = io.sampleRate,
                protocol = GgwaveProtocol.ULTRASOUND_FASTEST,
                volumePercent = 10,
            ),
        )
        val ourBytes = encodeChirpPayload(opts.attempt.outgoingPayload)
        val ourWaveform = codec.encode(ourBytes)

        val deadline = System.currentTimeMillis() + opts.timeoutMs
        val stopped = AtomicBoolean(false)

        val playLoop = async {
            var playCount = 0
            while (isActive && !stopped.get() && System.currentTimeMillis() < deadline) {
                try {
                    io.play(ourWaveform)
                    playCount++
                    opts.onChirpPlayed?.invoke(playCount)
                } catch (_: Throwable) {
                    // Play failures are non-fatal; keep listening.
                }
                if (stopped.get()) break
                val remaining = deadline - System.currentTimeMillis()
                val wait = minOf(CHIRP_LOOP_INTERVAL_MS, maxOf(0L, remaining))
                if (wait <= 0L) break
                delay(wait)
            }
        }

        try {
            while (System.currentTimeMillis() < deadline && !stopped.get()) {
                val samples = io.pullSamples()
                if (samples.isEmpty()) {
                    delay(50)
                    continue
                }

                val decoded = codec.decode(samples)
                if (decoded == null) continue
                if (decoded.size != CHIRP_PAYLOAD_LENGTH) {
                    opts.onDiscard?.invoke(ChirpDecodeError.Kind.LENGTH)
                    continue
                }

                try {
                    val payload = decodeChirpPayload(decoded)
                    if (sameBytes(payload.ephemeralPublicKey, opts.attempt.ephemeralPublicKey)) {
                        continue
                    }
                    if (payload.version != CHIRP_PROTOCOL_VERSION) {
                        throw ChirpExchangeError(ChirpExchangeErrorKind.VERSION_MISMATCH)
                    }
                    stopped.set(true)
                    return@coroutineScope ChirpExchangeResult(remoteChirp = payload)
                } catch (err: ChirpExchangeError) {
                    throw err
                } catch (err: ChirpDecodeError) {
                    if (err.kind == ChirpDecodeError.Kind.VERSION) {
                        throw ChirpExchangeError(ChirpExchangeErrorKind.VERSION_MISMATCH)
                    }
                    opts.onDiscard?.invoke(err.kind)
                    continue
                }
            }
            throw ChirpExchangeError(ChirpExchangeErrorKind.TIMEOUT)
        } finally {
            stopped.set(true)
            codec.free()
            playLoop.cancelAndJoin()
            try {
                io.close()
            } catch (_: Throwable) {
                // ignore
            }
        }
    }
}

suspend fun openAndroidAudioIO(): AudioIO = openPcmFloatAudioIO()

suspend fun runChirpPairing(
    deps: ChirpPairingDeps,
    role: PairingRole = PairingRole.RESPONDER,
    codecFactory: ChirpCodecFactory = ::createChirpCodec,
): PairedDevice {
    deps.onProgress(PairingProgress(phase = PairingPhase.IDLE))

    val localIdentity = deps.fetchLocalIdentity()

    if (role == PairingRole.RESPONDER && localIdentity.deviceType != DeviceType.EXTENSION) {
        throw ChirpPairingFailure(PairingPhase.HANDSHAKE_FAILED, "responder requires extension identity")
    }
    if (role == PairingRole.INITIATOR && localIdentity.deviceType == DeviceType.EXTENSION) {
        throw ChirpPairingFailure(PairingPhase.HANDSHAKE_FAILED, "initiator requires non-extension identity")
    }

    deps.onProgress(PairingProgress(phase = PairingPhase.LISTENING))

    val deviceTypeByte = when (localIdentity.deviceType) {
        DeviceType.ANDROID -> deviceTypeToByte(DeviceType.ANDROID)
        DeviceType.IOS -> deviceTypeToByte(DeviceType.IOS)
        DeviceType.DESKTOP -> deviceTypeToByte(DeviceType.DESKTOP)
        DeviceType.EXTENSION -> deviceTypeToByte(DeviceType.EXTENSION)
    }
    val attempt = createPairingAttempt(deviceTypeByte)

    val remoteChirpPublicKey: ByteArray
    try {
        val audio = deps.openAudio()
        val result = runChirpExchange(
            io = audio,
            opts = ChirpExchangeOptions(
                attempt = attempt,
                timeoutMs = deps.chirpTimeoutMs,
            ),
            codecFactory = codecFactory,
        )
        remoteChirpPublicKey = result.remoteChirp.ephemeralPublicKey
    } catch (err: ChirpExchangeError) {
        val phase = when (err.kind) {
            ChirpExchangeErrorKind.TIMEOUT -> PairingPhase.TIMEOUT
            ChirpExchangeErrorKind.VERSION_MISMATCH -> PairingPhase.VERSION_MISMATCH
            ChirpExchangeErrorKind.MIC_DENIED -> PairingPhase.MIC_DENIED
            ChirpExchangeErrorKind.AUDIO_UNAVAILABLE -> PairingPhase.HANDSHAKE_FAILED
        }
        deps.onProgress(PairingProgress(phase = phase, errorMessage = err.message))
        throw ChirpPairingFailure(phase, err.message ?: phase.wireValue)
    }

    deps.onProgress(PairingProgress(phase = PairingPhase.HANDSHAKE))

    if (role == PairingRole.RESPONDER) {
        return runResponderHandshake(deps, attempt, localIdentity)
    }
    return runInitiatorHandshake(deps, remoteChirpPublicKey, localIdentity)
}

private suspend fun runResponderHandshake(
    deps: ChirpPairingDeps,
    attempt: PairingAttempt,
    localIdentity: LocalChirpIdentity,
): PairedDevice {
    val subnet = deps.getLocalSubnet()
    if (subnet.isNullOrBlank()) {
        deps.onProgress(PairingProgress(phase = PairingPhase.HANDSHAKE_FAILED, errorMessage = "no_subnet"))
        throw ChirpPairingFailure(PairingPhase.HANDSHAKE_FAILED, "could not determine local subnet")
    }

    val peerEndpoint = firstResponsivePeer(deps, subnet, CHIRP_PAIRING_PORT)
    if (peerEndpoint == null) {
        deps.onProgress(PairingProgress(phase = PairingPhase.HANDSHAKE_FAILED, errorMessage = "peer_not_found"))
        throw ChirpPairingFailure(PairingPhase.HANDSHAKE_FAILED, "could not reach paired peer on subnet")
    }

    val transport = deps.openPairingTransport(peerEndpoint)
    val result = try {
        runPairingResponder(
            ourEphemeralPrivateKey = attempt.ephemeralPrivateKey,
            ourEphemeralPublicKey = attempt.ephemeralPublicKey,
            ourIdentity = localIdentity,
            transport = transport,
        )
    } catch (err: Throwable) {
        try {
            transport.close()
        } catch (_: Throwable) {
            // ignore
        }
        deps.onProgress(PairingProgress(phase = PairingPhase.HANDSHAKE_FAILED, errorMessage = err.message))
        throw ChirpPairingFailure(PairingPhase.HANDSHAKE_FAILED, err.message ?: "handshake failed")
    }
    transport.close()

    val paired = gateOnSAS(deps, result)
    activatePeerLicenseIfAny(deps, result)
    deps.completePairing(paired)
    deps.cacheLastKnownIp(paired.longTermPublicKeyHex.orEmpty(), peerEndpoint.ip)
    deps.onProgress(PairingProgress(phase = PairingPhase.PAIRED, peerDeviceName = paired.deviceName))
    return paired
}

private suspend fun gateOnSAS(
    deps: ChirpPairingDeps,
    result: PairingHandshakeResult,
): PairedDevice {
    val confirmed = try {
        deps.confirmSAS(result.sas)
    } catch (err: Throwable) {
        deps.onProgress(PairingProgress(phase = PairingPhase.HANDSHAKE_FAILED, errorMessage = err.message))
        throw ChirpPairingFailure(PairingPhase.HANDSHAKE_FAILED, err.message ?: "handshake failed")
    }
    if (!confirmed) {
        deps.onProgress(PairingProgress(phase = PairingPhase.HANDSHAKE_FAILED, errorMessage = "sas_rejected"))
        throw ChirpPairingFailure(
            PairingPhase.HANDSHAKE_FAILED,
            "pairing rejected - the verification codes did not match",
        )
    }
    return result.device
}

/**
 * If the (now SAS-confirmed) peer advertised a Pro license, hand it to the
 * deps to verify + activate. Best-effort: a bad/duplicate license never fails
 * the pairing — the device link still succeeds.
 */
private suspend fun activatePeerLicenseIfAny(
    deps: ChirpPairingDeps,
    result: PairingHandshakeResult,
) {
    val license = result.peerLicense ?: return
    try {
        deps.activatePeerLicense(license)
    } catch (_: Throwable) {
        // Invalid signature, already-Pro, or storage error — pairing stands.
    }
}

private suspend fun runInitiatorHandshake(
    deps: ChirpPairingDeps,
    remoteChirpPublicKey: ByteArray,
    localIdentity: LocalChirpIdentity,
): PairedDevice {
    val server = try {
        deps.openPairingServer()
    } catch (err: Throwable) {
        deps.onProgress(PairingProgress(phase = PairingPhase.HANDSHAKE_FAILED, errorMessage = err.message))
        throw ChirpPairingFailure(PairingPhase.HANDSHAKE_FAILED, err.message ?: "handshake failed")
    }

    val result = try {
        runPairingInitiator(remoteChirpPublicKey, localIdentity, server.transport)
    } catch (err: Throwable) {
        try {
            server.transport.close()
        } catch (_: Throwable) {
            // ignore
        }
        deps.onProgress(PairingProgress(phase = PairingPhase.HANDSHAKE_FAILED, errorMessage = err.message))
        throw ChirpPairingFailure(PairingPhase.HANDSHAKE_FAILED, err.message ?: "handshake failed")
    }
    server.transport.close()

    val paired = gateOnSAS(deps, result)
    activatePeerLicenseIfAny(deps, result)
    deps.completePairing(paired)
    deps.cacheLastKnownIp(paired.longTermPublicKeyHex.orEmpty(), server.peerIp)
    deps.onProgress(PairingProgress(phase = PairingPhase.PAIRED, peerDeviceName = paired.deviceName))
    return paired
}

private suspend fun firstResponsivePeer(
    deps: ChirpPairingDeps,
    subnet: String,
    port: Int,
): PairingEndpoint? {
    val probeTimeoutMs = 250L
    val batch = 16
    for (start in 1..254 step batch) {
        val end = minOf(start + batch - 1, 254)
        val ips = (start..end).map { "$subnet.$it" }
        val results = coroutineScope {
            ips.map { ip ->
                async {
                    try {
                        val transport = withTimeoutOrNull(probeTimeoutMs) {
                            deps.openPairingTransport(PairingEndpoint(ip = ip, port = port))
                        }
                        if (transport != null) {
                            transport.close()
                            PairingEndpoint(ip = ip, port = port)
                        } else {
                            null
                        }
                    } catch (_: Throwable) {
                        null
                    }
                }
            }.mapNotNull { it.await() }
        }
        val hit = results.firstOrNull()
        if (hit != null) return hit
    }
    return null
}

private suspend fun runPairingInitiator(
    remoteEphemeralPublicKey: ByteArray,
    ourIdentity: LocalChirpIdentity,
    transport: PairingTransport,
): PairingHandshakeResult {
    val session = noiseNKInitiator(
        remoteEphemeralPublicKey,
        transport::send,
        transport::receive,
    )

    transport.send(
        session.encrypt(
            encodePairingFrame(
                PairingFrame(
                    type = "PAIR_IDENTITY",
                    deviceId = ourIdentity.deviceId,
                    deviceName = ourIdentity.deviceName,
                    deviceType = ourIdentity.deviceType,
                    longTermPublicKeyHex = ourIdentity.longTermPublicKeyHex,
                    signingPublicKeyHex = ourIdentity.signingPublicKeyHex,
                    wsPort = ourIdentity.wsPort,
                    license = ourIdentity.license,
                ),
            ),
        ),
    )

    val theirCipher = transport.receive()
    val theirPlain = session.decrypt(theirCipher)
    val theirFrame = decodePairingFrame(theirPlain)
    return PairingHandshakeResult(
        device = frameToPairedDevice(theirFrame),
        sas = SASDerivation.derive(session.handshakeHash),
        peerLicense = theirFrame.license,
    )
}

private suspend fun runPairingResponder(
    ourEphemeralPrivateKey: ByteArray,
    ourEphemeralPublicKey: ByteArray,
    ourIdentity: LocalChirpIdentity,
    transport: PairingTransport,
): PairingHandshakeResult {
    val session = noiseNKResponder(
        ourEphemeralPrivateKey,
        ourEphemeralPublicKey,
        transport::send,
        transport::receive,
    )

    val theirCipher = transport.receive()
    val theirPlain = session.decrypt(theirCipher)
    val theirFrame = decodePairingFrame(theirPlain)

    transport.send(
        session.encrypt(
            encodePairingFrame(
                PairingFrame(
                    type = "PAIR_IDENTITY",
                    deviceId = ourIdentity.deviceId,
                    deviceName = ourIdentity.deviceName,
                    deviceType = ourIdentity.deviceType,
                    longTermPublicKeyHex = ourIdentity.longTermPublicKeyHex,
                    signingPublicKeyHex = ourIdentity.signingPublicKeyHex,
                    wsPort = ourIdentity.wsPort,
                    license = ourIdentity.license,
                ),
            ),
        ),
    )

    return PairingHandshakeResult(
        device = frameToPairedDevice(theirFrame),
        sas = SASDerivation.derive(session.handshakeHash),
        peerLicense = theirFrame.license,
    )
}

internal data class PairingFrame(
    val type: String,
    val deviceId: String,
    val deviceName: String,
    val deviceType: DeviceType,
    val longTermPublicKeyHex: String,
    val signingPublicKeyHex: String? = null,
    val wsPort: Int? = null,
    // Optional signed Pro license blob — additive; peers that don't understand
    // it ignore it. Activated by the receiver only after SAS confirmation.
    val license: String? = null,
)

internal fun encodePairingFrame(frame: PairingFrame): ByteArray {
    val builder = StringBuilder()
    builder.append('{')
    builder.append("\"type\":\"").append(frame.type).append('\"')
    builder.append(",\"device_id\":\"").append(escapeJson(frame.deviceId)).append('\"')
    builder.append(",\"device_name\":\"").append(escapeJson(frame.deviceName)).append('\"')
    builder.append(",\"device_type\":\"").append(frame.deviceType.wireValue).append('\"')
    builder.append(",\"long_term_public_key_hex\":\"").append(frame.longTermPublicKeyHex).append('\"')
    if (!frame.signingPublicKeyHex.isNullOrBlank()) {
        builder.append(",\"signing_public_key_hex\":\"").append(frame.signingPublicKeyHex).append('\"')
    }
    if (frame.wsPort != null) {
        builder.append(",\"ws_port\":").append(frame.wsPort)
    }
    if (!frame.license.isNullOrBlank()) {
        builder.append(",\"license\":\"").append(escapeJson(frame.license)).append('\"')
    }
    builder.append('}')
    return builder.toString().encodeToByteArray()
}

internal fun decodePairingFrame(bytes: ByteArray): PairingFrame {
    val text = bytes.decodeToString()
    val parsed = simpleJsonObject(text)
    require(parsed["type"] == "PAIR_IDENTITY") { "expected PAIR_IDENTITY frame" }
    val deviceId = parsed["device_id"] ?: error("device_id missing")
    val deviceName = parsed["device_name"] ?: error("device_name missing")
    val deviceType = when (parsed["device_type"]) {
        "extension" -> DeviceType.EXTENSION
        "android" -> DeviceType.ANDROID
        "ios" -> DeviceType.IOS
        "desktop" -> DeviceType.DESKTOP
        else -> error("device_type missing")
    }
    val longTermPublicKeyHex = parsed["long_term_public_key_hex"] ?: error("long_term_public_key_hex missing")
    require(longTermPublicKeyHex.matches(Regex("^[0-9a-f]{64}$"))) {
        "long_term_public_key_hex is not 32-byte hex"
    }
    val signingPublicKeyHex = parsed["signing_public_key_hex"]?.takeIf { it.isNotBlank() }
    if (signingPublicKeyHex != null) {
        require(signingPublicKeyHex.matches(Regex("^[0-9a-f]{64}$"))) {
            "signing_public_key_hex is not 32-byte hex"
        }
    }
    val wsPort = parsed["ws_port"]?.toIntOrNull()
    val license = parsed["license"]?.takeIf { it.isNotBlank() }
    return PairingFrame(
        type = "PAIR_IDENTITY",
        deviceId = deviceId,
        deviceName = deviceName,
        deviceType = deviceType,
        longTermPublicKeyHex = longTermPublicKeyHex,
        signingPublicKeyHex = signingPublicKeyHex,
        wsPort = wsPort,
        license = license,
    )
}

private fun frameToPairedDevice(frame: PairingFrame): PairedDevice {
    return PairedDevice(
        deviceId = frame.deviceId,
        deviceName = frame.deviceName,
        deviceType = frame.deviceType.wireValue,
        longTermPublicKeyHex = frame.longTermPublicKeyHex,
        signingPublicKeyHex = frame.signingPublicKeyHex,
        trustSource = "direct",
        wsPort = frame.wsPort,
        pairedAt = System.currentTimeMillis(),
    )
}

private fun escapeJson(value: String): String {
    return buildString(value.length) {
        for (ch in value) {
            when (ch) {
                '\\' -> append("\\\\")
                '"' -> append("\\\"")
                '\b' -> append("\\b")
                '\u000C' -> append("\\f")
                '\n' -> append("\\n")
                '\r' -> append("\\r")
                '\t' -> append("\\t")
                else -> append(ch)
            }
        }
    }
}

private fun simpleJsonObject(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).trim()
    if (body.isEmpty()) return emptyMap()
    val out = linkedMapOf<String, String>()
    var i = 0
    while (i < body.length) {
        while (i < body.length && body[i].isWhitespace()) i++
        require(i < body.length && body[i] == '"') { "expected string key" }
        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 valueStart = i
            while (i < body.length && body[i] != '"') i++
            value = body.substring(valueStart, i)
            i++
        } else {
            val valueStart = i
            while (i < body.length && body[i] != ',') i++
            value = body.substring(valueStart, i).trim()
        }
        out[key] = value
        while (i < body.length && body[i].isWhitespace()) i++
        if (i < body.length && body[i] == ',') i++
    }
    return out
}