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

        job = scope.launch(Dispatchers.IO) {
            try {
                val connected = OkHttpRelaySocket.connect(
                    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 = 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(
                            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 onConsent(request: AdoptionFrame.Request): Boolean = true
            override suspend fun upsertPairedDevice(device: PairedDevice) {
                sourceDeps.upsertPairedDevice(device)
            }
        }
        return grantAdoption(channel, request, boundDeps)
    } finally {
        noise.close()
    }
}