package com.peachpasswords.sync

import java.security.MessageDigest
import javax.inject.Inject
import javax.inject.Singleton
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.withTimeout
import kotlinx.coroutines.withTimeoutOrNull
import org.json.JSONObject
import com.peachpasswords.sync.hexToPubkey

/**
 * 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(
    private val generation: SyncGenerationPort,
    private val deviceStore: PairedDeviceRegistry,
) {
    /**
     * 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: AdoptionReceiverV2Deps,
        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.connectPairing(invitation.roomUrl())
        val adopted = adoptFromRelaySocket(
            socket = socket,
            invitation = invitation,
            deps = deps,
            confirm = confirm,
            windowMs = windowMs,
        )
        markRepairProgress()
        return adopted
    }

    /**
     * §11.3 REPAIRING → V2_ACTIVE. Any adoption that reaches here while this
     * device is still anchored is a v2 one: the post-reset source refuses v1 at
     * the first frame, and a v1 source cannot answer a v2 request.
     */
    private fun markRepairProgress() {
        onDirectV2PairingEstablished(
            generation,
            hasDirectDevice = deviceStore.getAll().any { it.trustSource == "direct" },
            nowMs = System.currentTimeMillis(),
        )
    }

    /**
     * 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: AdoptionSourceV2Deps,
        confirm: suspend (sas: String, request: AdoptionFrame.Request) -> Boolean,
        onPeerOutdated: () -> Unit = {},
        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.connectPairing(invitation.roomUrl())
        val granted = grantFromRelaySocket(
            socket = socket,
            invitation = invitation,
            sourceDeps = sourceDeps,
            confirm = confirm,
            // §11.6 on this entry point too: the probe lives here rather than
            // being passed in, so neither source session can be built without it.
            postReset = { generation.isPostReset(System.currentTimeMillis()) },
            onPeerOutdated = onPeerOutdated,
            windowMs = windowMs,
        )
        markRepairProgress()
        return granted
    }
}

/**
 * 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: AdoptionReceiverV2Deps,
    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: AdoptionSourceV2Deps,
    confirm: suspend (sas: String, request: AdoptionFrame.Request) -> Boolean,
    postReset: () -> Boolean = { false },
    onPeerOutdated: () -> Unit = {},
    windowMs: Long = PAIRING_WINDOW_MS,
): PairedDevice {
    var completed = false
    return try {
        val granted = withTimeout(windowMs) {
            runRelayAdoptionInitiatorSourceSession(
                socket = socket,
                invitation = invitation,
                sourceDeps = sourceDeps,
                confirm = confirm,
                postReset = postReset,
                onPeerOutdated = onPeerOutdated,
            ) ?: 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: AdoptionReceiverV2Deps,
    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 invitedStatic = invitation.noiseStaticPublicKeyHex
        ?.let(::hexToPubkey)
        ?: throw IllegalArgumentException("pairing_invitation_missing_static_key")
    val staticMatches = MessageDigest.isEqual(sourceStatic, invitedStatic)
    invitedStatic.fill(0)
    if (!staticMatches) {
        sourceStatic.fill(0)
        throw IllegalArgumentException("pairing_static_key_mismatch")
    }
    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 : AdoptionReceiverV2Deps {
            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))

            // The v2 transcript presents this device's durable identity, which
            // exists before the vault does; the v1 path still derives its own.
            override fun localIdentity(): LocalDeviceIdentity = receiverDeps.localIdentity()

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

            override suspend fun validateSource(prepare: AdoptionFrame.Prepare) {
                receiverDeps.validateSource(prepare)
            }

            override suspend fun hasSourceIdentityCollision(prepare: AdoptionFrame.Prepare): Boolean =
                receiverDeps.hasSourceIdentityCollision(prepare)

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

            override suspend fun hasSourceIdentityCollision(source: AdoptionFrame.Source): Boolean =
                receiverDeps.hasSourceIdentityCollision(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 upsertPairedDeviceAndClearRevocation(device: PairedDevice) {
                receiverDeps.upsertPairedDeviceAndClearRevocation(device)
            }

            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 applyIncomingDeviceManifest(
                manifest: SignedDeviceManifest,
                authenticatedRemoteStaticPublicKey: ByteArray,
            ) {
                receiverDeps.applyIncomingDeviceManifest(manifest, authenticatedRemoteStaticPublicKey)
            }
            override suspend fun markSynced(deviceId: String) {
                receiverDeps.markSynced(deviceId)
            }
        }
        // Capture the source's snapshot capability so we can skip the 15s
        // inline-snapshot wait when the source has nothing to send.
        var sourceCanSendSnapshot = false
        val adopted = requestAdoptionV2(
            channel,
            boundDeps,
            onPrepare = { prepare -> sourceCanSendSnapshot = prepare.snapshotCapable == true },
        ) ?: return null
        if (sourceCanSendSnapshot) {
            // Try to receive the inline vault snapshot.
            try {
                val snapshot = withTimeoutOrNull(15_000) {
                    receiveVaultSnapshot(channel)
                }
                if (snapshot != null) {
                    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. Apply before acking so
                    // the receiver graduates into the mesh pool on the same
                    // adoption session; a missing or invalid long-term key is
                    // accepted (we still ack) because the graph can arrive via
                    // the regular sync path on the next SYNC_DEVICE_MANIFEST.
                    try {
                        val manifestFrame = withTimeoutOrNull(5_000) { receiveFrame(channel) }
                        if (manifestFrame is AdoptionFrame.DeviceManifestFrame) {
                            val sourceLongTerm = adopted.longTermPublicKeyHex
                            if (sourceLongTerm != null && sourceLongTerm.length == 64) {
                                boundDeps.applyIncomingDeviceManifest(
                                    manifestFrame.manifest,
                                    hexToPubkey(sourceLongTerm),
                                )
                            }
                            sendFrame(channel, AdoptionFrame.DeviceManifestApplied)
                        }
                    } catch (_: Exception) {}
                    boundDeps.markSynced(adopted.deviceId)
                } else {
                    PeachSyncLog.i("adoption: snapshot receive timed out")
                }
            } catch (e: Exception) {
                PeachSyncLog.w("adoption: snapshot receive failed — reason=${relayLogReason(e)}")
            }
        }
        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: AdoptionSourceV2Deps,
    confirm: suspend (sas: String, request: AdoptionFrame.Request) -> Boolean,
    postReset: () -> Boolean = { false },
    onPeerOutdated: () -> Unit = {},
): 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"
    }
    // PB-26-103: authenticate the cleartext ACK against the static pinned in the
    // invitation before running Noise. This is the root-key-bearing direction —
    // the phone grants its vault here — so trusting the ACK lets a relay
    // terminate both sessions, present its own key to each side, and collect
    // ADOPT_GRANT.root_key_hex. The mirror check on the receiver path is at
    // runRelayAdoptionReceiverSession.
    val responderStatic = hexToPubkey(ack.optString("ephemeralPublicKeyHex"))
    val invitedStatic = invitation.noiseStaticPublicKeyHex
        ?.let(::hexToPubkey)
        ?: throw IllegalArgumentException("pairing_invitation_missing_static_key")
    val staticMatches = MessageDigest.isEqual(responderStatic, invitedStatic)
    invitedStatic.fill(0)
    if (!staticMatches) {
        responderStatic.fill(0)
        throw IllegalArgumentException("pairing_static_key_mismatch")
    }
    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 (request !is AdoptionFrame.RequestV2 && !isAdoptRequest(request)) return null
        // §11.6 through the same predicate the QR-display session uses. After a
        // secure reset a v1 peer is refused before the SAS prompt: its identity
        // derives from the legacy root, which is exactly what the reset retires.
        // Without this the rule held on one pairing entry point and not the
        // other (PB-26-117).
        if (request is AdoptionFrame.Request && !allowLegacyV1(postReset())) {
            onPeerOutdated()
            return null
        }

        val pairState = AdoptionPairState(
            rendezvousSecretHex = Rendezvous.bytesToHex(
                Rendezvous.deriveSecret(noise.handshakeHash),
            ),
            relayUrl = invitation.relayUrl,
        )
        val boundDeps = bindSourceSessionDeps(
            sourceDeps = sourceDeps,
            sourceIdentity = sourceDeps.sourceIdentity,
            pairState = pairState,
            sas = SASDerivation.derive(noise.handshakeHash),
            snapshotCapable = true,
            confirm = confirm,
        )
        // A v2 receiver introduces itself first; a v1 receiver still opens with
        // the label-only ADOPT_REQUEST and is served by the legacy path.
        val adopted = when (request) {
            is AdoptionFrame.RequestV2 -> grantAdoptionV2(channel, request, boundDeps)
            is AdoptionFrame.Request -> {
                if (!confirm(SASDerivation.derive(noise.handshakeHash), request)) return null
                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 CancellationException) throw error
                PeachSyncLog.w("adoption: snapshot send failed — reason=${relayLogReason(error)}")
            }
        }
        return adopted
    } finally {
        noise.close()
    }
}