package com.peachpasswords.sync

import java.security.MessageDigest
import javax.inject.Inject
import javax.inject.Singleton
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.callbackFlow
import kotlinx.coroutines.launch
import kotlinx.coroutines.withTimeout
import kotlinx.coroutines.withTimeoutOrNull
import org.json.JSONObject

sealed interface RelayAdoptionEvent {
    /** The source is already occupying the non-buffering room; it is safe to show this. */
    data class InvitationReady(val invitation: String) : RelayAdoptionEvent
    data class AwaitingPeerConfirmation(
        val sas: String,
        val deviceName: String,
    ) : RelayAdoptionEvent
    data class Adopted(val deviceName: String) : RelayAdoptionEvent
    data object Expired : RelayAdoptionEvent
    data class Rejected(val reason: String) : RelayAdoptionEvent
}

@Singleton
class RelayAdoptionHost @Inject constructor(
    private val relaySettings: RelaySettingsStore,
    private val generation: SyncGenerationPort,
    private val deviceStore: PairedDeviceRegistry,
) {
    private var job: Job? = null
    private var socket: RelayBinarySocket? = null
    private val consent = SourceConsentGate()

    /**
     * The vault owner's answer to the source-side SAS prompt. §10.2 requires a
     * human decision on the device that hands over the root — without this the
     * grant was sent as soon as the code had been drawn (PB-26-116).
     */
    fun confirmSas(matches: Boolean) = consent.resolve(matches)

    fun start(
        scope: CoroutineScope,
        ourStaticPrivateKey: ByteArray,
        ourStaticPublicKey: ByteArray,
        localIdentity: LocalDeviceIdentity,
        sourceDeps: AdoptionSourceV2Deps,
        windowMs: Long = PAIRING_WINDOW_MS,
    ): Flow<RelayAdoptionEvent> = callbackFlow {
        // Pairing is also how this mesh's relay reaches a new peer: the address
        // in the invitation is the one the receiver will dial for ever after.
        // Handing out Peach's relay while the user has asked for a private one
        // would push the same lie one layer down, into the mesh.
        val relayUrl = when (val decision = relaySettings.relayInForce()) {
            is RelayDecision.Refuse -> {
                trySend(RelayAdoptionEvent.Rejected(decision.reason.wire()))
                close()
                return@callbackFlow
            }

            is RelayDecision.Dial -> decision.relayUrl
        }
        val token = Rendezvous.newPairingToken()
        val invitation = createPairingInvitation(
            relayUrl,
            token,
            noiseStaticPublicKeyHex = Rendezvous.bytesToHex(ourStaticPublicKey),
        )

        job = scope.launch(Dispatchers.IO) {
            try {
                val connected = OkHttpRelaySocket.connectPairing(
                    Rendezvous.buildUrl(relayUrl, token),
                )
                socket = connected
                // The relay buffers nothing, so the invitation is not surfaced
                // until the source WebSocket is open and waiting in the room.
                trySend(RelayAdoptionEvent.InvitationReady(invitation))

                var peerOutdated = false
                val adopted = withTimeout(windowMs) {
                    runRelayAdoptionSourceSession(
                        socket = connected,
                        expectedToken = token,
                        relayUrl = relayUrl,
                        ourStaticPrivateKey = ourStaticPrivateKey,
                        ourStaticPublicKey = ourStaticPublicKey,
                        localIdentity = localIdentity,
                        sourceDeps = sourceDeps,
                        // §11.6: once this device has reset its generation, a v1
                        // peer is repair-required exactly as an outdated sync
                        // peer is. The invitation still goes out; the refusal
                        // happens at the first frame.
                        postReset = { generation.isPostReset(System.currentTimeMillis()) },
                        onPeerOutdated = { peerOutdated = true },
                        confirm = { sas, request ->
                            trySend(
                                RelayAdoptionEvent.AwaitingPeerConfirmation(
                                    sas = sas,
                                    deviceName = capDeviceName(request.deviceName.ifBlank { "New device" }),
                                ),
                            )
                            // §10.2: the device that owns the vault decides.
                            // Cancellation (screen exit, stop(), window expiry)
                            // ends the await without a grant, which is the
                            // fail-closed outcome.
                            consent.await()
                        },
                    )
                }
                if (peerOutdated) {
                    trySend(RelayAdoptionEvent.Rejected(PEER_PROTOCOL_OUTDATED))
                } else if (adopted == null) {
                    trySend(RelayAdoptionEvent.Rejected("adoption_rejected"))
                } else {
                    onDirectV2PairingEstablished(
                        generation,
                        hasDirectDevice = deviceStore.getAll().any { it.trustSource == "direct" },
                        nowMs = System.currentTimeMillis(),
                    )
                    trySend(
                        RelayAdoptionEvent.Adopted(
                            capDeviceName(adopted.deviceName?.takeIf(String::isNotBlank) ?: "New device"),
                        ),
                    )
                }
            } catch (error: kotlinx.coroutines.TimeoutCancellationException) {
                trySend(RelayAdoptionEvent.Expired)
            } catch (error: kotlinx.coroutines.CancellationException) {
                // stop() cancels this job when the user leaves the screen or
                // restarts pairing. That is not a pairing failure — reporting it
                // as one raced the new session and replaced the freshly shown
                // code with "Couldn't pair: StandaloneCoroutine was cancelled".
                // Cancellation must propagate, never be swallowed.
                throw error
            } catch (error: Throwable) {
                trySend(RelayAdoptionEvent.Rejected(error.message ?: "relay_adoption_failed"))
            } finally {
                socket?.close()
                socket = null
                close()
            }
        }

        awaitClose { stop() }
    }

    fun stop() {
        // Leaving the screen is a denial, not a silent timeout: an unanswered
        // source prompt must never leave the session waiting on a grant.
        consent.resolve(false)
        job?.cancel()
        job = null
        socket?.close()
        socket = null
    }
}

/**
 * The source-side consent gate.
 *
 * §10.2 puts a human decision on the device that owns the vault: the root must
 * not be read until someone there approves the SAS the other device is showing.
 * Extracted as a seam so the await/resolve contract is provable without a relay
 * socket — a callback that returned true immediately is exactly the defect this
 * gate exists to prevent (PB-26-116).
 */
internal class SourceConsentGate {
    private val lock = Any()
    private var pending: CompletableDeferred<Boolean>? = null

    /** Suspend until [resolve]. A gate with no pending await resolves nothing. */
    suspend fun await(): Boolean {
        val deferred = synchronized(lock) {
            CompletableDeferred<Boolean>().also { pending = it }
        }
        return try {
            deferred.await()
        } finally {
            synchronized(lock) { if (pending === deferred) pending = null }
        }
    }

    /** Answer the live prompt. A second call, or one with no prompt, is a no-op. */
    fun resolve(matches: Boolean) {
        val deferred = synchronized(lock) { pending.also { pending = null } } ?: return
        deferred.complete(matches)
    }
}

/**
 * Source half of first-contact adoption through an opaque relay splice.
 *
 * The relay token selects only a room. Authentication and confidentiality
 * begin with Noise NK; the root grant remains behind both SAS decisions in
 * [grantAdoption].
 */
internal suspend fun runRelayAdoptionSourceSession(
    socket: RelayBinarySocket,
    expectedToken: String,
    relayUrl: String,
    ourStaticPrivateKey: ByteArray,
    ourStaticPublicKey: ByteArray,
    localIdentity: LocalDeviceIdentity,
    sourceDeps: AdoptionSourceV2Deps,
    confirm: suspend (sas: String, request: AdoptionFrame.Request) -> Boolean,
    postReset: () -> Boolean = { false },
    onPeerOutdated: () -> Unit = {},
): PairedDevice? {
    val hello = JSONObject(String(socket.receive(), Charsets.UTF_8))
    require(hello.optString("type") == "PEACH_HELLO") { "expected PEACH_HELLO" }
    require(
        MessageDigest.isEqual(
            expectedToken.toByteArray(Charsets.US_ASCII),
            hello.optString("token").toByteArray(Charsets.US_ASCII),
        ),
    ) { "pairing token mismatch" }

    socket.send(
        JSONObject()
            .put("type", "PEACH_HELLO_ACK")
            .put(
                "ephemeralPublicKeyHex",
                ourStaticPublicKey.joinToString("") { "%02x".format(it) },
            )
            .toString()
            .toByteArray(Charsets.UTF_8),
    )

    val noise = noiseNKResponder(
        ourStaticPrivateKey,
        ourStaticPublicKey,
        socket::send,
        socket::receive,
    )
    try {
        val channel = object : PairingTransport {
            override suspend fun send(msg: ByteArray) = socket.send(noise.encrypt(msg))
            override suspend fun receive(): ByteArray = noise.decrypt(socket.receive())
            override suspend fun close() = socket.close()
        }

        val request = receiveFrame(channel)
        require(request is AdoptionFrame.Request || request is AdoptionFrame.RequestV2) { "expected ADOPT_REQUEST" }
        // §11.6: refused before any SAS prompt and before any grant. §10.3
        // freezes the wire reason set to the two v1 peers already understand, so
        // this is a local refusal rather than a new ADOPT_REJECT reason.
        if (request is AdoptionFrame.Request && !allowLegacyV1(postReset())) {
            onPeerOutdated()
            return null
        }
        if (request is AdoptionFrame.Request && !confirm(SASDerivation.derive(noise.handshakeHash), request)) return null

        val pairState = AdoptionPairState(
            rendezvousSecretHex = Rendezvous.bytesToHex(
                Rendezvous.deriveSecret(noise.handshakeHash),
            ),
            relayUrl = Rendezvous.normalizeRelayUrl(relayUrl),
        )
        // §10.5: the whole forwarding surface lives in one binder so a dropped
        // callback is caught by AdoptionAdapterCoverageTest rather than by a peer.
        val boundDeps = bindSourceSessionDeps(
            sourceDeps = sourceDeps,
            sourceIdentity = localIdentity,
            pairState = pairState,
            sas = SASDerivation.derive(noise.handshakeHash),
            snapshotCapable = sourceDeps.snapshotCapable(),
            confirm = confirm,
        )
        val adopted = when (request) {
            is AdoptionFrame.RequestV2 -> grantAdoptionV2(channel, request, boundDeps)
            is AdoptionFrame.Request -> grantAdoption(channel, request, boundDeps)
            else -> null
        } ?: return null
        // If the receiver advertised snapshot support, send the vault inline.
        val receiverWantsSnapshot = when (request) {
            is AdoptionFrame.Request -> request.snapshotCapable == true
            is AdoptionFrame.RequestV2 -> request.snapshotCapable == true
            else -> false
        }
        if (receiverWantsSnapshot) {
            try {
                val snapshot = boundDeps.vaultSnapshot()
                if (snapshot != null) {
                    val duressSlot = runCatching { boundDeps.duressSlotJson() }.getOrNull()
                    sendVaultSnapshot(channel, snapshot.first, snapshot.second, duressSlot)
                    val applied = receiveFrame(channel)
                    if (applied is AdoptionFrame.SnapshotApplied) {
                        val manifest = boundDeps.outgoingDeviceManifest()
                        if (manifest != null) {
                            sendFrame(channel, AdoptionFrame.DeviceManifestFrame(manifest))
                            withTimeoutOrNull(5_000) { receiveFrame(channel) }
                        }
                        boundDeps.markSynced(adopted.deviceId)
                    }
                }
            } catch (error: Exception) {
                if (error is SnapshotDuressFrameTooLargeException) throw error
                if (error is kotlinx.coroutines.CancellationException) throw error
                PeachSyncLog.w("adoption: snapshot send failed — reason=${relayLogReason(error)}")
            }
        }
        return adopted
    } finally {
        noise.close()
    }
}