package com.peachpasswords.sync

import javax.inject.Inject
import javax.inject.Singleton
import kotlinx.coroutines.withTimeout
import kotlinx.coroutines.withTimeoutOrNull
import org.json.JSONObject

/**
 * Receiver half of relay-first vault adoption.
 *
 * The invitation is reachability only. The relay never sees identity, SAS, or
 * vault material; those enter only after the Noise NK session is established.
 */
@Singleton
class RelayAdoptionReceiver @Inject constructor() {
    /**
     * Phone as receiver: scan a [VaultStatus.POPULATED] invitation, dial the
     * relay room, complete Noise, and receive the source's root key. This is
     * the historical pairing direction and stays byte-for-byte unchanged.
     */
    suspend fun adopt(
        encodedInvitation: String,
        deps: AdoptionReceiverDeps,
        confirm: suspend (sas: String) -> Boolean,
        windowMs: Long = PAIRING_WINDOW_MS,
    ): PairedDevice {
        val invitation = parsePairingInvitation(encodedInvitation)
            ?: throw IllegalArgumentException("invalid_pairing_invitation")
        require(invitation.vaultStatus == VaultStatus.POPULATED) {
            "receiver adopt() expects a populated invitation; got ${invitation.vaultStatus}"
        }
        val socket = OkHttpRelaySocket.connect(invitation.roomUrl())
        return adoptFromRelaySocket(
            socket = socket,
            invitation = invitation,
            deps = deps,
            confirm = confirm,
            windowMs = windowMs,
        )
    }

    /**
     * Phone as source: scan a [VaultStatus.EMPTY] invitation from an empty
     * device, dial the relay room, complete Noise, and hand our root key over.
     * Symmetric to [adopt]: the initiator runs Noise NK, but the payload
     * direction flips — the phone sends the grant, the empty device receives
     * it.
     */
    suspend fun grantFromInvitation(
        encodedInvitation: String,
        sourceDeps: AdoptionSourceDeps,
        confirm: suspend (sas: String, request: AdoptionFrame.Request) -> Boolean,
        windowMs: Long = PAIRING_WINDOW_MS,
    ): PairedDevice {
        val invitation = parsePairingInvitation(encodedInvitation)
            ?: throw IllegalArgumentException("invalid_pairing_invitation")
        require(invitation.vaultStatus == VaultStatus.EMPTY) {
            "grantFromInvitation() expects an empty invitation; got ${invitation.vaultStatus}"
        }
        val socket = OkHttpRelaySocket.connect(invitation.roomUrl())
        return grantFromRelaySocket(
            socket = socket,
            invitation = invitation,
            sourceDeps = sourceDeps,
            confirm = confirm,
            windowMs = windowMs,
        )
    }
}

/**
 * Complete one receiver session and choose the correct socket shutdown.
 *
 * The final receiver action is an encrypted ADOPT_APPLIED frame. OkHttp's
 * send() only queues it, so a successful session must close gracefully; an
 * immediate cancel can discard that frame after both devices have already
 * persisted one another, leaving the source to report a false failure.
 */
internal suspend fun adoptFromRelaySocket(
    socket: RelayBinarySocket,
    invitation: PairingInvitation,
    deps: AdoptionReceiverDeps,
    confirm: suspend (sas: String) -> Boolean,
    windowMs: Long = PAIRING_WINDOW_MS,
): PairedDevice {
    var completed = false
    return try {
        val adopted =
            withTimeout(windowMs) {
                runRelayAdoptionReceiverSession(
                    socket = socket,
                    invitation = invitation,
                    receiverDeps = deps,
                    confirm = confirm,
                ) ?: throw IllegalStateException("adoption_rejected")
            }
        completed = true
        adopted
    } finally {
        if (completed) {
            socket.closeAfterQueuedFrames()
        } else {
            socket.close()
        }
    }
}

/**
 * Source-initiator counterpart to [adoptFromRelaySocket]: phone dials the
 * relay room, runs Noise as initiator, and hands the root key over to the
 * empty responder. The final source action is an encrypted ADOPT_COMPLETE,
 * so a graceful close is required to land it.
 */
internal suspend fun grantFromRelaySocket(
    socket: RelayBinarySocket,
    invitation: PairingInvitation,
    sourceDeps: AdoptionSourceDeps,
    confirm: suspend (sas: String, request: AdoptionFrame.Request) -> Boolean,
    windowMs: Long = PAIRING_WINDOW_MS,
): PairedDevice {
    var completed = false
    return try {
        val granted = withTimeout(windowMs) {
            runRelayAdoptionInitiatorSourceSession(
                socket = socket,
                invitation = invitation,
                sourceDeps = sourceDeps,
                confirm = confirm,
            ) ?: throw IllegalStateException("adoption_rejected")
        }
        completed = true
        granted
    } finally {
        if (completed) {
            socket.closeAfterQueuedFrames()
        } else {
            socket.close()
        }
    }
}

/**
 * Pure live-session seam used by both production and the real Noise interop
 * test. The source must already occupy the relay room before this starts.
 */
internal suspend fun runRelayAdoptionReceiverSession(
    socket: RelayBinarySocket,
    invitation: PairingInvitation,
    receiverDeps: AdoptionReceiverDeps,
    confirm: suspend (sas: String) -> Boolean,
): PairedDevice? {
    socket.send(
        JSONObject()
            .put("type", "PEACH_HELLO")
            .put("token", invitation.token)
            .toString()
            .toByteArray(Charsets.UTF_8),
    )

    val ack = JSONObject(String(socket.receive(), Charsets.UTF_8))
    require(ack.optString("type") == "PEACH_HELLO_ACK") {
        "expected PEACH_HELLO_ACK"
    }
    val sourceStatic = hexToPubkey(ack.optString("ephemeralPublicKeyHex"))
    val noise = try {
        noiseNKInitiator(
            sourceStatic,
            socket::send,
            socket::receive,
        )
    } finally {
        sourceStatic.fill(0)
    }
    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 pairState = AdoptionPairState(
            rendezvousSecretHex = Rendezvous.bytesToHex(
                Rendezvous.deriveSecret(noise.handshakeHash),
            ),
            relayUrl = invitation.relayUrl,
        )
        val boundDeps = object : AdoptionReceiverDeps {
            override val deviceName = receiverDeps.deviceName
            override val deviceType = receiverDeps.deviceType
            override val pairState = pairState
            override fun generateDeviceId() = receiverDeps.generateDeviceId()
            override suspend fun confirmSas(): Boolean =
                confirm(SASDerivation.derive(noise.handshakeHash))

            override suspend fun validateSource(source: AdoptionFrame.Source) {
                receiverDeps.validateSource(source)
            }

            override suspend fun createVaultFromRootKey(rootKeyHex: String): String =
                receiverDeps.createVaultFromRootKey(rootKeyHex)

            override suspend fun deriveLocalIdentity(): LocalDeviceIdentity =
                receiverDeps.deriveLocalIdentity()

            override suspend fun clearPairedDeviceRevocation(deviceId: String) {
                receiverDeps.clearPairedDeviceRevocation(deviceId)
            }

            override suspend fun upsertPairedDevice(device: PairedDevice) {
                receiverDeps.upsertPairedDevice(device)
            }

            override suspend fun applyLicense(license: String) {
                receiverDeps.applyLicense(license)
            }

            override suspend fun license(): String? = receiverDeps.license()

            override fun onError(error: String) {
                receiverDeps.onError(error)
            }
            override suspend fun applySnapshot(entries: Map<String, String>, tombstones: Map<String, Long>, duressSlot: String?) {
                receiverDeps.applySnapshot(entries, tombstones, duressSlot)
            }
            override suspend fun markSynced(deviceId: String) {
                receiverDeps.markSynced(deviceId)
            }
        }
        val adopted = requestAdoption(channel, boundDeps) ?: return null
        // Try to receive the inline vault snapshot.
        try {
            val snapshot = withTimeout(15_000) {
                receiveVaultSnapshot(channel)
            }
            boundDeps.applySnapshot(snapshot.entries, snapshot.tombstones, snapshot.duressSlot)
            sendFrame(channel, AdoptionFrame.SnapshotApplied(
                entriesReceived = snapshot.entries.size,
                tombstonesReceived = snapshot.tombstones.size,
            ))
            // Optional device manifest from source
            try {
                val manifestFrame = withTimeout(5_000) { receiveFrame(channel) }
                if (manifestFrame is AdoptionFrame.DeviceManifestFrame) {
                    sendFrame(channel, AdoptionFrame.DeviceManifestApplied)
                }
            } catch (_: Exception) {}
            boundDeps.markSynced(adopted.deviceId)
        } catch (_: Exception) {
            // Fallback: sync pipeline fills the vault.
        }
        return adopted
    } finally {
        noise.close()
    }
}

/**
 * Phone-initiator source session for the "empty browser" invitation path.
 * The phone dials the relay room and waits for the empty responder to send
 * ADOPT_REQUEST, then hands over its root key through [grantAdoption]. The
 * SAS gate is the single commit signal — same model as the existing source
 * session — and onConsent is pre-cleared because the user already opted in
 * by scanning the QR.
 */
internal suspend fun runRelayAdoptionInitiatorSourceSession(
    socket: RelayBinarySocket,
    invitation: PairingInvitation,
    sourceDeps: AdoptionSourceDeps,
    confirm: suspend (sas: String, request: AdoptionFrame.Request) -> Boolean,
): PairedDevice? {
    socket.send(
        JSONObject()
            .put("type", "PEACH_HELLO")
            .put("token", invitation.token)
            .toString()
            .toByteArray(Charsets.UTF_8),
    )

    val ack = JSONObject(String(socket.receive(), Charsets.UTF_8))
    require(ack.optString("type") == "PEACH_HELLO_ACK") {
        "expected PEACH_HELLO_ACK"
    }
    val responderStatic = hexToPubkey(ack.optString("ephemeralPublicKeyHex"))
    val noise = try {
        noiseNKInitiator(
            responderStatic,
            socket::send,
            socket::receive,
        )
    } finally {
        responderStatic.fill(0)
    }
    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)
        if (!isAdoptRequest(request)) return null
        request as AdoptionFrame.Request
        if (!confirm(SASDerivation.derive(noise.handshakeHash), request)) return null

        val pairState = AdoptionPairState(
            rendezvousSecretHex = Rendezvous.bytesToHex(
                Rendezvous.deriveSecret(noise.handshakeHash),
            ),
            relayUrl = invitation.relayUrl,
        )
        val boundDeps = object : AdoptionSourceDeps {
            override val sourceIdentity = sourceDeps.sourceIdentity
            override val pairState = 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 (_: Exception) {
                // Fallback: receiver uses the sync pipeline.
            }
        }
        return adopted
    } finally {
        noise.close()
    }
}