package com.peachpasswords.sync

import java.security.MessageDigest
import javax.inject.Inject
import javax.inject.Singleton
import kotlinx.coroutines.CoroutineScope
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 var job: Job? = null
    private var socket: RelayBinarySocket? = null

    fun start(
        scope: CoroutineScope,
        ourStaticPrivateKey: ByteArray,
        ourStaticPublicKey: ByteArray,
        localIdentity: LocalDeviceIdentity,
        sourceDeps: AdoptionSourceDeps,
        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))

                val adopted = withTimeout(windowMs) {
                    runRelayAdoptionSourceSession(
                        socket = connected,
                        expectedToken = token,
                        relayUrl = relayUrl,
                        ourStaticPrivateKey = ourStaticPrivateKey,
                        ourStaticPublicKey = ourStaticPublicKey,
                        localIdentity = localIdentity,
                        sourceDeps = sourceDeps,
                        confirm = { sas, request ->
                            trySend(
                                RelayAdoptionEvent.AwaitingPeerConfirmation(
                                    sas = sas,
                                    deviceName = capDeviceName(request.deviceName.ifBlank { "New device" }),
                                ),
                            )
                            // Showing a fresh invitation from an unlocked vault
                            // is source-side authorization. The scanning device's
                            // explicit SAS confirmation is the single commit gate.
                            true
                        },
                    )
                }
                if (adopted == null) {
                    trySend(RelayAdoptionEvent.Rejected("adoption_rejected"))
                } else {
                    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() {
        job?.cancel()
        job = null
        socket?.close()
        socket = null
    }
}

/**
 * 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: AdoptionSourceDeps,
    confirm: suspend (sas: String, request: AdoptionFrame.Request) -> Boolean,
): 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(isAdoptRequest(request)) { "expected ADOPT_REQUEST" }
        request as AdoptionFrame.Request
        if (!confirm(SASDerivation.derive(noise.handshakeHash), request)) return null

        val pairState = AdoptionPairState(
            rendezvousSecretHex = Rendezvous.bytesToHex(
                Rendezvous.deriveSecret(noise.handshakeHash),
            ),
            relayUrl = Rendezvous.normalizeRelayUrl(relayUrl),
        )
        val boundDeps = object : AdoptionSourceDeps {
            override val sourceIdentity: LocalDeviceIdentity = localIdentity
            override val pairState: AdoptionPairState = pairState
            override suspend fun rootKeyHex(): String = sourceDeps.rootKeyHex()
            override suspend fun vaultId(): String = sourceDeps.vaultId()
            override suspend fun license(): String? = sourceDeps.license()
            override suspend fun applyLicense(license: String) {
                sourceDeps.applyLicense(license)
            }
            override suspend fun onConsent(request: AdoptionFrame.Request): Boolean = true
            override suspend fun upsertPairedDevice(device: PairedDevice) {
                sourceDeps.upsertPairedDevice(device)
            }
            override suspend fun vaultSnapshot() = sourceDeps.vaultSnapshot()
            override suspend fun outgoingDeviceManifest() = sourceDeps.outgoingDeviceManifest()
            override suspend fun markSynced(deviceId: String) = sourceDeps.markSynced(deviceId)
            override suspend fun duressSlotJson() = sourceDeps.duressSlotJson()
        }
        val adopted = grantAdoption(channel, request, boundDeps) ?: return null
        // If the receiver advertised snapshot support, send the vault inline.
        if (request.snapshotCapable == true) {
            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()
    }
}