package com.peachpasswords.sync

import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json

/**
 * Vault adoption frame protocol — exact port of the browser `shared/vault-adopt.ts`.
 *
 * Adoption lets a brand-new device join an existing vault: the SOURCE (an already
 * set-up device) hands its root key to the RECEIVER over the pairing channel so
 * both devices end up with the same vault id and can sync. The receiver seals an
 * empty local vault under its own password but reuses the adopted root key
 * ([VaultRepository.createVaultFromRootKey]).
 *
 * Frames travel as JSON with a `type` discriminator and SNAKE_CASE keys — unlike
 * the camelCase [SyncMessage] sync wire — so the keys must match the browser
 * byte-for-byte for cross-platform adoption. Optional fields are omitted when
 * absent (encodeDefaults = false) to mirror `JSON.stringify` dropping `undefined`.
 */
@Serializable
sealed class AdoptionFrame {
    /**
     * [sessionRequestId] carries the wire's `device_id`, but it is NOT an
     * identity: the requester generates it locally before the handshake that
     * establishes any long-term key, so it cannot be the device's real one
     * (see `PEACH_RELAY_AND_DEVICEID_FIXES.md` Part 2). It is a label for the
     * consent prompt and nothing else — never persist it, never compare it
     * against a previous attempt's, never key trust on it. The only identity
     * that may be persisted is `long_term_public_key_hex`, which arrives later
     * in ADOPT_IDENTITY / ADOPT_GRANT.source.
     */
    @Serializable
    @SerialName("ADOPT_REQUEST")
    data class Request(
        @SerialName("device_id") val sessionRequestId: String,
        @SerialName("device_name") val deviceName: String,
        @SerialName("device_type") val deviceType: String,
        val license: String? = null,
        @SerialName("snapshot_capable") val snapshotCapable: Boolean? = null,
    ) : AdoptionFrame()

    @Serializable
    @SerialName("ADOPT_CONFIRM")
    data class Confirm(
        val confirmed: Boolean,
        /**
         * The receiver's Pro key. It rides the confirm frame — sent only after
         * the SAS comparison — rather than the first request: a peer that
         * controls the scanned invitation could otherwise read the key and walk
         * away without ever being authenticated (PB-26-121).
         */
        val license: String? = null,
    ) : AdoptionFrame()

    @Serializable
    @SerialName("ADOPT_GRANT")
    data class Grant(
        @SerialName("root_key_hex") val rootKeyHex: String,
        @SerialName("vault_id") val vaultId: String,
        val license: String? = null,
        val source: Source,
        // v1-only receiver capability metadata.
        @SerialName("snapshot_capable") val snapshotCapable: Boolean = false,
    ) : AdoptionFrame()

    @Serializable
    @SerialName("ADOPT_IDENTITY")
    data class Identity(
        @SerialName("device_id") val deviceId: String,
        @SerialName("device_name") val deviceName: String,
        @SerialName("device_type") val deviceType: String,
        @SerialName("long_term_public_key_hex") val longTermPublicKeyHex: String,
        @SerialName("signing_public_key_hex") val signingPublicKeyHex: String? = null,
    ) : AdoptionFrame()

    /**
     * Encrypted source-to-receiver acknowledgement emitted only after the source
     * has durably recorded the receiver. The receiver cannot truthfully report a
     * successful pair before this arrives.
     */
    @Serializable
    @SerialName("ADOPT_COMPLETE")
    data object Complete : AdoptionFrame()

    /** Encrypted receiver-to-source acknowledgement after the source row is durable. */
    @Serializable
    @SerialName("ADOPT_APPLIED")
    data object Applied : AdoptionFrame()

    /** Neither peer proved Pro, so the pairing must not commit. */
    @Serializable
    @SerialName("ADOPT_REJECT")
    data class Reject(
        val reason: String,
    ) : AdoptionFrame()

    /**
     * One chunk of the inline vault snapshot sent by the source after
     * ADOPT_APPLIED over the same authenticated Noise NK channel. Entries are
     * opaque base64 ciphertext envelopes; tombstones ride the final chunk only.
     */
    @Serializable
    @SerialName("ADOPT_SNAPSHOT")
    data class Snapshot(
        @SerialName("chunk_index") val chunkIndex: Int,
        @SerialName("chunk_count") val chunkCount: Int,
        val entries: Map<String, String>,
        val tombstones: Map<String, Long>,
        @SerialName("duress_slot") val duressSlot: String? = null,
    ) : AdoptionFrame()

    /** Receiver acknowledgement sent after all snapshot chunks are applied. */
    @Serializable
    @SerialName("ADOPT_SNAPSHOT_APPLIED")
    data class SnapshotApplied(
        @SerialName("entries_received") val entriesReceived: Int,
        @SerialName("tombstones_received") val tombstonesReceived: Int,
    ) : AdoptionFrame()

    /** Optional signed device graph appended after the snapshot (manifest v3+). */
    @Serializable
    @SerialName("ADOPT_DEVICE_MANIFEST")
    data class DeviceManifestFrame(
        val manifest: SignedDeviceManifest,
    ) : AdoptionFrame()

    /** Receiver acknowledgement for an applied device manifest. */
    @Serializable
    @SerialName("ADOPT_DEVICE_MANIFEST_APPLIED")
    data object DeviceManifestApplied : AdoptionFrame()
    @Serializable
    @SerialName("ADOPT_REQUEST_V2")
    data class RequestV2(
        @SerialName("adoptionVersion") val adoptionVersion: Int,
        @SerialName("device_id") val deviceId: String,
        @SerialName("device_name") val deviceName: String,
        @SerialName("device_type") val deviceType: String,
        @SerialName("long_term_public_key_hex") val longTermPublicKeyHex: String,
        @SerialName("signing_public_key_hex") val signingPublicKeyHex: String,
        val license: String? = null,
        @SerialName("snapshot_capable") val snapshotCapable: Boolean? = null,
    ) : AdoptionFrame()

    @Serializable @SerialName("ADOPT_CONFIRM_REQUIRED")
    data object ConfirmRequired : AdoptionFrame()

    @Serializable @SerialName("ADOPT_PREPARE")
    data class Prepare(
        @SerialName("device_id") val deviceId: String,
        @SerialName("device_name") val deviceName: String,
        @SerialName("device_type") val deviceType: String,
        @SerialName("long_term_public_key_hex") val longTermPublicKeyHex: String,
        @SerialName("signing_public_key_hex") val signingPublicKeyHex: String,
        @SerialName("vault_id") val vaultId: String,
        @SerialName("ws_port") val wsPort: Int? = null,
        @SerialName("relay_url") val relayUrl: String? = null,
        val license: String? = null,
        @SerialName("snapshot_capable") val snapshotCapable: Boolean? = null,
    ) : AdoptionFrame()
    @Serializable @SerialName("ADOPT_READY")
    data object Ready : AdoptionFrame()
    @Serializable
    data class Source(
        @SerialName("device_id") val deviceId: String,
        @SerialName("device_name") val deviceName: String,
        @SerialName("device_type") val deviceType: String,
        @SerialName("long_term_public_key_hex") val longTermPublicKeyHex: String,
        @SerialName("signing_public_key_hex") val signingPublicKeyHex: String? = null,
        @SerialName("ws_port") val wsPort: Int? = null,
        @SerialName("relay_url") val relayUrl: String? = null,
    )
}

data class AdoptionPairState(
    val rendezvousSecretHex: String,
    val relayUrl: String? = null,
)

/**
 * Receiver-side reassembly result from [receiveVaultSnapshot]. Carries the
 * aggregated entries and tombstones plus the optional browser duress slot
 * JSON (present on the final chunk only, when the source has duress configured).
 */
data class ReceivedSnapshot(
    val entries: Map<String, String>,
    val tombstones: Map<String, Long>,
    val duressSlot: String? = null,
)

/**
 * Wire guardrails for the inline snapshot. A chunk's serialized entries map
 * stays under [MAX_SNAPSHOT_CHUNK_BYTES]; the whole transfer is capped at
 * [MAX_SNAPSHOT_CHUNKS] frames so a compromised source cannot exhaust receiver
 * memory (mirrors SYNC_LIMITS.MAX_FRAME_BYTES in the sync coordinator).
 */
const val MAX_SNAPSHOT_CHUNK_BYTES = 512 * 1024 // 512 KiB
const val MAX_SNAPSHOT_CHUNKS = 200

// Headroom added to the per-chunk entries budget when validating an encoded
// frame on the wire: JSON envelope plus the final chunk's tombstone map.
internal const val SNAPSHOT_FRAME_WIRE_CEILING = MAX_SNAPSHOT_CHUNK_BYTES + 256 * 1024
internal val MAX_PADDED_SNAPSHOT_FRAME_BYTES = SyncPadding.bucket(SNAPSHOT_FRAME_WIRE_CEILING + 4)

/**
 * PB-26-206 aggregate budget shared by sender and receiver — exact mirror of
 * the browser's `vault-adopt.ts` ceilings. The per-frame guard above bounds
 * one frame; this bounds the whole transfer independently of the 200-frame
 * count, so many individually valid frames cannot OOM the receiver.
 */
const val MAX_ADOPT_SNAPSHOT_BYTES = 10 * 1024 * 1024 // 10 MiB
const val MAX_ADOPT_SNAPSHOT_ENTRIES = 25_000
const val MAX_ADOPT_SNAPSHOT_TOMBSTONES = 25_000
/** Same value as [SyncLimits.MAX_ID_LENGTH]; named here so vectors read identically on both clients. */
const val MAX_ADOPT_SNAPSHOT_ID_CHARS = SyncLimits.MAX_ID_LENGTH
/**
 * Maximum encrypted entry envelope: plaintext ceiling plus the existing
 * base64 overhead allowance. Int arithmetic gives 699,114 — the browser's
 * float comparison accepts and rejects at exactly the same lengths.
 */
const val MAX_ADOPT_SNAPSHOT_ENVELOPE_BYTES = SyncLimits.MAX_ENTRY_ENVELOPE_BYTES

/**
 * Running aggregate for one snapshot transfer. Counts ride alongside the maps
 * (never re-derived per frame) so the bound stays O(n).
 */
class AdoptSnapshotBudget {
    val entries = LinkedHashMap<String, String>()
    val tombstones = LinkedHashMap<String, Long>()
    var duressSlot: String? = null
    /** Sum of encoded wire-frame bytes consumed so far. */
    var bytes: Int = 0
    var entryCount: Int = 0
    var tombstoneCount: Int = 0
}

fun createAdoptSnapshotBudget(): AdoptSnapshotBudget = AdoptSnapshotBudget()

/**
 * PB-26-206 shared preflight/merge helper used by BOTH the sender (preflight
 * before the first frame goes out) and the receiver (per frame, after the
 * byte-budget check, before decoded maps are merged).
 *
 * One bounded pass per attacker-controlled map: each key/value is validated
 * and projected counts are enforced inside the same loop that merges, so the
 * maps are never iterated count→size→merge repeatedly. Throws before merging
 * the offending frame, leaving no partial adopted state.
 */
fun applyAdoptSnapshotFrame(
    acc: AdoptSnapshotBudget,
    entries: Map<String, String>,
    tombstones: Map<String, Long>,
    duressSlot: String?,
    frameBytes: Int,
) {
    if (frameBytes > MAX_ADOPT_SNAPSHOT_BYTES - acc.bytes) {
        throw IllegalStateException("ADOPT_SNAPSHOT_OVER_BUDGET")
    }
    for ((id, envelope) in entries) {
        if (id.length > MAX_ADOPT_SNAPSHOT_ID_CHARS) throw IllegalStateException("ADOPT_SNAPSHOT_ID_TOO_LONG")
        if (envelope.length > MAX_ADOPT_SNAPSHOT_ENVELOPE_BYTES) {
            throw IllegalStateException("ADOPT_SNAPSHOT_ENVELOPE_TOO_LARGE")
        }
        if (acc.entryCount >= MAX_ADOPT_SNAPSHOT_ENTRIES) {
            throw IllegalStateException("ADOPT_SNAPSHOT_TOO_MANY_ENTRIES")
        }
        acc.entryCount++
        acc.entries[id] = envelope
    }
    for ((id, ts) in tombstones) {
        if (id.length > MAX_ADOPT_SNAPSHOT_ID_CHARS) throw IllegalStateException("ADOPT_SNAPSHOT_ID_TOO_LONG")
        // Long is always finite, so non-finite hostile JSON dies at the
        // strict codec decode before reaching this merge (see V7).
        if (acc.tombstoneCount >= MAX_ADOPT_SNAPSHOT_TOMBSTONES) {
            throw IllegalStateException("ADOPT_SNAPSHOT_TOO_MANY_TOMBSTONES")
        }
        acc.tombstoneCount++
        acc.tombstones[id] = ts
    }
    if (!duressSlot.isNullOrEmpty()) {
        // The duress slot rides inside wire frames already bounded by the
        // frame ceiling, and its raw bytes count against the aggregate above.
        if (duressSlot.length > SNAPSHOT_FRAME_WIRE_CEILING) {
            throw IllegalStateException("ADOPT_SNAPSHOT_DURESS_FRAME_TOO_LARGE")
        }
        acc.duressSlot = duressSlot
    }
    acc.bytes += frameBytes
}

/** What the RECEIVER (joining device) needs to adopt a vault from a source. */
interface AdoptionReceiverDeps {
    val deviceName: String
    val deviceType: DeviceType
    fun generateDeviceId(): String
    /** Root-key transfer cannot begin until this device confirms the SAS. */
    suspend fun confirmSas(): Boolean
    /**
     * Local safety policy evaluated after the authenticated grant is validated
     * but before any vault state changes.
     */
    suspend fun validateSource(source: AdoptionFrame.Source) {}

    /**
     * True when the source's identity (from ADOPT_GRANT) is in the local
     * terminal revocation ledger. When true, the receiver sends
     * ADOPT_DEVICE_ID_COLLISION so the source rotates and resends the grant.
     */
    suspend fun hasSourceIdentityCollision(source: AdoptionFrame.Source): Boolean = false

    /** Seal a fresh local vault under the receiver's password, reusing [rootKeyHex]. Returns the vault id. */
    suspend fun createVaultFromRootKey(rootKeyHex: String): String
    /** Build the receiver's local identity after the adopted vault exists. */
    suspend fun deriveLocalIdentity(): LocalDeviceIdentity
    /** Persist the direct pairing and lift its terminal revocation as one store operation. */
    suspend fun upsertPairedDeviceAndClearRevocation(device: PairedDevice) {
        upsertPairedDevice(device)
        clearPairedDeviceRevocation(device.deviceId)
    }
    suspend fun clearPairedDeviceRevocation(deviceId: String) {}
    suspend fun upsertPairedDevice(device: PairedDevice)
    suspend fun applyLicense(license: String) {}
    /** Return this receiver's locally validated Pro key, if present. */
    suspend fun license(): String? = null
    val pairState: AdoptionPairState? get() = null
    fun onError(error: String)
    /** Apply encrypted entries + tombstones from the inline snapshot. */
    suspend fun applySnapshot(entries: Map<String, String>, tombstones: Map<String, Long>, duressSlot: String? = null) {}
    /**
     * Apply the source's signed device graph so the receiver graduates into the
     * mesh. The [authenticatedRemoteStaticPublicKey] is the source's Noise
     * transport static key (NOT the signing key); [applyIncomingDeviceManifest]
     * resolves it to the source's [PairedDevice] locally and uses the stored
     * signing key to verify the manifest signature. Default no-op preserves
     * the contract for callers that haven't wired the mesh path yet.
     */
    suspend fun applyIncomingDeviceManifest(
        manifest: SignedDeviceManifest,
        authenticatedRemoteStaticPublicKey: ByteArray,
    ) {}
    /** Mark the paired device as synced after a successful snapshot. */
    suspend fun markSynced(deviceId: String) {}
}

/** What the SOURCE (existing device) needs to grant its vault to a receiver. */
interface AdoptionSourceDeps {
    val sourceIdentity: LocalDeviceIdentity
    suspend fun rootKeyHex(): String
    suspend fun vaultId(): String
    suspend fun license(): String? = null
    /** Validate and durably store a Pro key offered by the receiver. */
    suspend fun applyLicense(license: String) {}
    /**
     * Ask the user to approve handing this vault to the requesting device. The
     * request's name/type/label are for the human to read; the decision must not
     * be made by recognising [AdoptionFrame.Request.sessionRequestId], which
     * carries no continuity between attempts.
     */
    suspend fun onConsent(request: AdoptionFrame.Request): Boolean
    suspend fun persistPendingHolder(request: AdoptionFrame.Request, nowMs: Long): Boolean = true
    suspend fun completePendingHolder(request: AdoptionFrame.Request, device: PairedDevice): Boolean = true
    suspend fun upsertPairedDevice(device: PairedDevice)
    fun hasDeviceIdentityCollision(device: PairedDevice): Boolean = false
    /**
     * Rotate the source's identity after the receiver reports a collision on
     * the source's own identity. Returns the fresh identity for the resent
     * ADOPT_GRANT.
     */
    suspend fun rotateSourceIdentity(): LocalDeviceIdentity = error("rotateSourceIdentity not wired")
    val pairState: AdoptionPairState? get() = null
    /** Return all encrypted entries + tombstones for the inline snapshot. */
    suspend fun vaultSnapshot(): Pair<Map<String, String>, Map<String, Long>>? = null
    /** Return the signed device manifest for the graph, if available. */
    suspend fun outgoingDeviceManifest(): SignedDeviceManifest? = null
    /** Mark the paired device as synced after a successful snapshot. */
    suspend fun markSynced(deviceId: String) {}
    /** Browser-shaped sealed duress slot JSON for the snapshot, if duress is configured. */
    suspend fun duressSlotJson(): String? = null
}

internal val Hex64 = Regex("^[0-9a-f]{64}$")

fun validateRootKeyHex(rootKeyHex: String?): Boolean = rootKeyHex != null && Hex64.matches(rootKeyHex)

fun isAdoptRequest(frame: AdoptionFrame?): Boolean =
    frame is AdoptionFrame.Request &&
        frame.sessionRequestId.isNotEmpty() &&
        isDeviceType(frame.deviceType)

fun isAdoptGrant(frame: AdoptionFrame?): Boolean =
    frame is AdoptionFrame.Grant &&
        validateRootKeyHex(frame.rootKeyHex) &&
        Hex64.matches(frame.vaultId) &&
        isDeviceType(frame.source.deviceType) &&
        Hex64.matches(frame.source.longTermPublicKeyHex) &&
        (frame.source.relayUrl == null ||
            runCatching { Rendezvous.normalizeRelayUrl(frame.source.relayUrl) }.isSuccess)

fun isAdoptConfirm(frame: AdoptionFrame?): Boolean =
    frame is AdoptionFrame.Confirm

fun isAdoptIdentity(frame: AdoptionFrame?): Boolean =
    frame is AdoptionFrame.Identity &&
        frame.deviceId.isNotEmpty() &&
        isDeviceType(frame.deviceType) &&
        Hex64.matches(frame.longTermPublicKeyHex)

fun isAdoptComplete(frame: AdoptionFrame?): Boolean =
    frame is AdoptionFrame.Complete

fun isAdoptApplied(frame: AdoptionFrame?): Boolean =
    frame is AdoptionFrame.Applied

fun isAdoptReject(frame: AdoptionFrame?): Boolean =
    frame is AdoptionFrame.Reject &&
        (frame.reason == "ADOPT_PRO_REQUIRED" || frame.reason == "ADOPT_DEVICE_ID_COLLISION")

fun isAdoptSnapshot(frame: AdoptionFrame?): Boolean =
    frame is AdoptionFrame.Snapshot &&
        frame.chunkIndex >= 0 &&
        frame.chunkCount >= 1

/**
 * RECEIVER side: ask a source to adopt us. Sends ADOPT_REQUEST, verifies the
 * granted root key actually derives the advertised vault id (catching a lying or
 * corrupted source before it can seal a mismatched vault), seals the local vault,
 * applies any carried license, returns our identity, and records the source as a
 * paired device. Returns null on any protocol failure (reason via [deps.onError]).
 */
suspend fun requestAdoption(
    channel: PairingTransport,
    deps: AdoptionReceiverDeps,
    // Called once after a valid ADOPT_GRANT is decoded and validated, before
    // any side effects (vault creation, paired-device upsert). Used by the
    // receiver to inspect the grant's snapshot_capable flag and skip the 15s
    // inline-snapshot wait when the source has nothing to send. Defaults to no-op
    // for callers (tests, other helpers) that don't need to inspect the grant.
    onGrant: ((AdoptionFrame.Grant) -> Unit)? = null,
    nowMs: () -> Long = { System.currentTimeMillis() },
): PairedDevice? {
    val receiverLicense = deps.license()
    // Fresh every attempt, and deliberately not our real device id: no identity
    // exists yet at this point in the protocol.
    sendFrame(channel, AdoptionFrame.Request(
        sessionRequestId = deps.generateDeviceId(),
        deviceName = deps.deviceName,
        deviceType = deps.deviceType.wireValue,
        snapshotCapable = true,
    ))

    val confirmed = deps.confirmSas()
    // Sent only after the SAS decision (PB-26-121): a request-frame licence
    // reaches a peer the user has not authenticated yet.
    sendFrame(channel, AdoptionFrame.Confirm(confirmed, license = receiverLicense))
    if (!confirmed) {
        deps.onError("ADOPT_SAS_REJECTED")
        return null
    }

    val frame = receiveFrame(channel)
    if (isAdoptReject(frame)) {
        deps.onError((frame as AdoptionFrame.Reject).reason)
        return null
    }
    if (!isAdoptGrant(frame)) {
        deps.onError("ADOPT_UNEXPECTED_FRAME")
        return null
    }
    val grant = frame as AdoptionFrame.Grant
    onGrant?.invoke(grant)

    val rootKeyBytes = hexToPubkey(grant.rootKeyHex)
    try {
        if (computeVaultId(rootKeyBytes) != grant.vaultId) {
            deps.onError("ADOPT_VAULT_ID_MISMATCH")
            return null
        }
    } finally {
        rootKeyBytes.fill(0)
    }
    if (receiverLicense == null && grant.license == null) {
        deps.onError("ADOPT_PRO_REQUIRED")
        return null
    }

    var acceptedGrant = grant
    var acceptedSource = grant.source ?: run {
        deps.onError("ADOPT_UNEXPECTED_FRAME")
        return null
    }
    deps.validateSource(acceptedSource)

    // Check the source's identity against the local revocation ledger. A
    // previously-revoked device must present a fresh identity so that stale
    // signed revocations held by other mesh peers cannot remove it after
    // re-pair. If the source's identity is revoked, send ADOPT_DEVICE_ID_COLLISION
    // and wait for a fresh grant (one retry).
    if (deps.hasSourceIdentityCollision(acceptedSource)) {
        sendFrame(channel, AdoptionFrame.Reject("ADOPT_DEVICE_ID_COLLISION"))
        val retry = receiveFrame(channel)
        if (!isAdoptGrant(retry)) {
            deps.onError("ADOPT_SOURCE_COLLISION_NOT_RESOLVED")
            return null
        }
        acceptedGrant = retry as AdoptionFrame.Grant
        acceptedSource = acceptedGrant.source ?: run {
            deps.onError("ADOPT_UNEXPECTED_FRAME")
            return null
        }
        if (deps.hasSourceIdentityCollision(acceptedSource)) {
            sendFrame(channel, AdoptionFrame.Reject("ADOPT_DEVICE_ID_COLLISION"))
            deps.onError("ADOPT_SOURCE_COLLISION_NOT_RESOLVED")
            return null
        }
        // Re-validate the fresh grant's vault id.
        val freshRootBytes = hexToPubkey(acceptedGrant.rootKeyHex)
        try {
            if (computeVaultId(freshRootBytes) != acceptedGrant.vaultId) {
                deps.onError("ADOPT_VAULT_ID_MISMATCH")
                return null
            }
        } finally {
            freshRootBytes.fill(0)
        }
    }
    if (receiverLicense == null && acceptedGrant.license == null) {
        deps.onError("ADOPT_PRO_REQUIRED")
        return null
    }


    if (receiverLicense == null) {
        acceptedGrant.license?.let { deps.applyLicense(it) }
    }
    val vaultId = deps.createVaultFromRootKey(acceptedGrant.rootKeyHex)
    if (vaultId != acceptedGrant.vaultId) {
        deps.onError("ADOPT_VAULT_ID_MISMATCH")
        return null
    }


    val identity = deps.deriveLocalIdentity()
    sendFrame(channel, AdoptionFrame.Identity(
        deviceId = identity.deviceId,
        deviceName = identity.deviceName,
        deviceType = identity.deviceType.wireValue,
        longTermPublicKeyHex = identity.longTermPublicKeyHex,
        signingPublicKeyHex = identity.signingPublicKeyHex,
    ))

    if (!isAdoptComplete(receiveFrame(channel))) {
        deps.onError("ADOPT_COMPLETION_MISSING")
        return null
    }

    val sourceDevice = PairedDevice(
        deviceId = acceptedSource.deviceId,
        deviceName = acceptedSource.deviceName,
        deviceType = acceptedSource.deviceType,
        longTermPublicKeyHex = acceptedSource.longTermPublicKeyHex,
        signingPublicKeyHex = acceptedSource.signingPublicKeyHex,
        trustSource = "direct",
        wsPort = acceptedSource.wsPort,
        rendezvousSecretHex = deps.pairState?.rendezvousSecretHex,
        relayUrl = deps.pairState?.relayUrl ?: acceptedSource.relayUrl,
        pairedAt = nowMs(),
    )
    // A physically confirmed direct re-pair is the sole path allowed to lift
    // terminal revocation. The registry validates identity before either value
    // is changed, so a mismatch leaves the revocation intact.
    deps.upsertPairedDeviceAndClearRevocation(sourceDevice)
    sendFrame(channel, AdoptionFrame.Applied)
    return sourceDevice
}

/**
 * SOURCE side: respond to an ADOPT_REQUEST. Gates on user consent, hands over the
 * root key, then records the receiver's returned identity as a paired device.
 * Returns null if consent is denied or the receiver's identity frame is invalid.
 */
suspend fun grantAdoption(
    channel: PairingTransport,
    request: AdoptionFrame.Request,
    deps: AdoptionSourceDeps,
    nowMs: () -> Long = { System.currentTimeMillis() },
): PairedDevice? {
    if (!deps.onConsent(request)) return null

    val confirmation = receiveFrame(channel)
    val confirmedFrame = confirmation as AdoptionFrame.Confirm
    if (!isAdoptConfirm(confirmation) || !confirmedFrame.confirmed) {
        return null
    }

    var sourceLicense = deps.license()
    // The receiver's key arrives on the confirm frame, after both SAS decisions
    // (PB-26-121). request.license is read only for a peer built before that
    // change.
    val receiverLicense = confirmedFrame.license ?: request.license
    if (sourceLicense == null && receiverLicense != null) {
        sourceLicense = runCatching {
            deps.applyLicense(receiverLicense)
            deps.license()
        }.getOrNull()
    }
    if (sourceLicense == null) {
        sendFrame(channel, AdoptionFrame.Reject("ADOPT_PRO_REQUIRED"))
        throw IllegalStateException("ADOPT_PRO_REQUIRED")
    }

    if (!deps.persistPendingHolder(request, nowMs())) return null
    val rootKeyHex = deps.rootKeyHex()
    if (!validateRootKeyHex(rootKeyHex)) throw IllegalStateException("ADOPT_INVALID_ROOT_KEY")

    var sourceIdentity = deps.sourceIdentity
    var grantAttempts = 0
    while (true) {
        sendFrame(channel, AdoptionFrame.Grant(
            rootKeyHex = rootKeyHex,
            vaultId = deps.vaultId(),
            license = sourceLicense,
            source = AdoptionFrame.Source(
                deviceId = sourceIdentity.deviceId,
                deviceName = sourceIdentity.deviceName,
                deviceType = sourceIdentity.deviceType.wireValue,
                longTermPublicKeyHex = sourceIdentity.longTermPublicKeyHex,
                signingPublicKeyHex = sourceIdentity.signingPublicKeyHex,
                wsPort = sourceIdentity.wsPort,
                relayUrl = deps.pairState?.relayUrl,
            ),
            snapshotCapable = request.snapshotCapable == true,
        ))

        var identity = receiveFrame(channel)

        // Receiver may reject the source's identity if it is in the receiver's
        // revocation ledger. Rotate and resend (one retry).
        if (identity is AdoptionFrame.Reject &&
            identity.reason == "ADOPT_DEVICE_ID_COLLISION" &&
            grantAttempts == 0
        ) {
            sourceIdentity = deps.rotateSourceIdentity()
            grantAttempts++
            continue
        }

        if (!isAdoptIdentity(identity)) return null
        identity as AdoptionFrame.Identity

        var device = PairedDevice(
            deviceId = identity.deviceId,
            deviceName = identity.deviceName,
            deviceType = identity.deviceType,
            longTermPublicKeyHex = identity.longTermPublicKeyHex,
            signingPublicKeyHex = identity.signingPublicKeyHex,
            trustSource = "direct",
            rendezvousSecretHex = deps.pairState?.rendezvousSecretHex,
            relayUrl = deps.pairState?.relayUrl,
            pairedAt = nowMs(),
        )
        if (deps.hasDeviceIdentityCollision(device)) {
            sendFrame(channel, AdoptionFrame.Reject("ADOPT_DEVICE_ID_COLLISION"))
            identity = receiveFrame(channel)
            if (!isAdoptIdentity(identity)) return null
            identity as AdoptionFrame.Identity
            device = device.copy(
                deviceId = identity.deviceId,
                deviceName = identity.deviceName,
                deviceType = identity.deviceType,
                longTermPublicKeyHex = identity.longTermPublicKeyHex,
                signingPublicKeyHex = identity.signingPublicKeyHex,
            )
            if (deps.hasDeviceIdentityCollision(device)) {
                sendFrame(channel, AdoptionFrame.Reject("ADOPT_DEVICE_ID_COLLISION"))
                return null
            }
        }
        deps.upsertPairedDevice(device)
        sendFrame(channel, AdoptionFrame.Complete)
        if (!isAdoptApplied(receiveFrame(channel))) return null
        if (!deps.completePendingHolder(request, device)) return null
        return device
    }
}

/**
 * Wire codec. `classDiscriminator = "type"` reproduces the browser's flat
 * `{ "type": "ADOPT_…", … }` shape; `encodeDefaults = false` omits absent
 * optionals exactly like `JSON.stringify` drops `undefined`.
 */
val AdoptJson: Json = Json {
    classDiscriminator = "type"
    encodeDefaults = false
    ignoreUnknownKeys = true
}

suspend fun sendFrame(channel: PairingTransport, frame: AdoptionFrame) {
    channel.send(AdoptJson.encodeToString(AdoptionFrame.serializer(), frame).toByteArray(Charsets.UTF_8))
}

/** Decode one frame, tolerating a malformed/unknown payload by returning null. */
suspend fun receiveFrame(channel: PairingTransport): AdoptionFrame? =
    runCatching {
        AdoptJson.decodeFromString(AdoptionFrame.serializer(), String(channel.receive(), Charsets.UTF_8))
    }.getOrNull()

/**
 * Source side: chunk the vault snapshot and send it over the pairing channel.
 *
 * Wire shape (v2.9.9+, supplemental frame; mirrors browser v4.6.4):
 *   - N entries frames, each <= [MAX_SNAPSHOT_CHUNK_BYTES] serialized entries.
 *   - One supplemental frame always follows the entries. It carries the
 *     tombstones and, when configured, the duress slot. Sending the same frame
 *     count in both cases prevents a relay from learning duress presence.
 *     Keeping duress off the entries frame also prevents a sizeable sealed
 *     duress vault from pushing the entries frame over the wire ceiling.
 *
 * Frame budget ([MAX_SNAPSHOT_CHUNKS] = 200):
 *   - ONE of the 200 frames is always reserved for the supplemental frame;
 *     entries use at most 199 chunks. Excess entries are deferred to the sync
 *     fallback pipeline (sync carries entries and tombstones on Android too,
 *     so this is a graceful truncation). The supplemental frame is required
 *     even without duress so its presence cannot disclose configuration. When
 *     configured, it carries the duress slot; omitting it would leave the
 *     adopted device permanently without its duress vault.
 *
 * Oversized duress: duress that itself exceeds [SNAPSHOT_FRAME_WIRE_CEILING]
 * is rejected with a thrown error BEFORE any snapshot frame is sent. Hard stop;
 * the caller must split duress out of band. We do not silently truncate
 * duress; sync cannot restore it.
 *
 * Backward compat: pre-v2.9.9 senders put `duress_slot` on the entries
 * final chunk. `receiveVaultSnapshot` reads `duress_slot` off any chunk (a
 * duress-only frame has empty `entries` and the complete tombstone map),
 * so both shapes decode and both new->old / old->new directions interoperate.
 */
suspend fun sendVaultSnapshot(
    channel: PairingTransport,
    entries: Map<String, String>,
    tombstones: Map<String, Long>,
    duressSlot: String? = null,
) {
    val frames = buildAdoptSnapshotFrames(entries, tombstones, duressSlot)
    // PB-26-206 preflight with the SAME helper the receiver enforces: nothing
    // is sent until every frame fits the aggregate budget, so an over-budget
    // snapshot is never partially sent or silently truncated. The caller falls
    // through to the regular sync pipeline (deny-only: no wire-field change,
    // invitation prefix unchanged).
    val acc = createAdoptSnapshotBudget()
    val encoded = frames.map {
        AdoptJson.encodeToString(AdoptionFrame.serializer(), it).encodeToByteArray()
    }
    for (i in frames.indices) {
        val frame = frames[i]
        applyAdoptSnapshotFrame(acc, frame.entries, frame.tombstones, frame.duressSlot, encoded[i].size)
    }
    for (raw in encoded) channel.send(SyncPadding.pad(raw))
}

/**
 * Pure frame construction for [sendVaultSnapshot]: chunking, the duress-frame
 * reservation, and the oversized-duress hard stop. No I/O happens here, so
 * the preflight above can reject before the first byte goes out.
 */
private fun buildAdoptSnapshotFrames(
    entries: Map<String, String>,
    tombstones: Map<String, Long>,
    duressSlot: String?,
): List<AdoptionFrame.Snapshot> {
    val includeDuress = !duressSlot.isNullOrBlank()
    // Always reserve the supplemental frame so its presence cannot disclose
    // whether the source has a duress vault.
    val entryCap = MAX_SNAPSHOT_CHUNKS - 1

    val chunks = mutableListOf<Map<String, String>>()
    var current = LinkedHashMap<String, String>()
    var currentBytes = 2 // "{}"
    for ((id, envelope) in entries) {
        // +1 for the separating comma; an overestimate only shrinks chunks.
        val entryBytes = AdoptJson.encodeToString(mapOf(id to envelope)).length + 1
        if (current.isNotEmpty() && currentBytes + entryBytes >= MAX_SNAPSHOT_CHUNK_BYTES) {
            chunks.add(current)
            current = LinkedHashMap()
            currentBytes = 2
        }
        current[id] = envelope
        currentBytes += entryBytes
    }
    chunks.add(current)

    val entriesChunkCount = minOf(chunks.size, entryCap)
    val totalChunks = entriesChunkCount + 1
    val supplementalFrame = AdoptionFrame.Snapshot(
        chunkIndex = entriesChunkCount,
        chunkCount = totalChunks,
        entries = emptyMap(),
        tombstones = tombstones,
        duressSlot = duressSlot.takeIf { includeDuress },
    )

    if (
        AdoptJson.encodeToString(AdoptionFrame.serializer(), supplementalFrame)
            .encodeToByteArray().size > SNAPSHOT_FRAME_WIRE_CEILING
    ) {
        if (includeDuress) throw SnapshotDuressFrameTooLargeException()
        throw IllegalStateException("ADOPT_SNAPSHOT_SUPPLEMENTAL_FRAME_TOO_LARGE")
    }

    val frames = mutableListOf<AdoptionFrame.Snapshot>()
    for (index in 0 until entriesChunkCount) {
        frames.add(AdoptionFrame.Snapshot(
            chunkIndex = index,
            chunkCount = totalChunks,
            entries = chunks[index],
            tombstones = emptyMap(),
            duressSlot = null,
        ))
    }
    frames.add(supplementalFrame)
    return frames
}

internal class SnapshotDuressFrameTooLargeException :
    IllegalStateException("ADOPT_SNAPSHOT_DURESS_FRAME_TOO_LARGE")

/**
 * Receiver side: receive and reassemble every snapshot chunk. Adds raw frame
 * bytes to the aggregate counter before decoding (rejecting before JSON
 * deserialization when the next frame exceeds the remaining budget), then
 * enforces the per-frame wire ceiling, sequential 0-based chunk_index
 * ordering, a stable chunk_count within [MAX_SNAPSHOT_CHUNKS], and structural
 * validity of each frame. Merging goes through the shared aggregate helper:
 * every key/value and projected count is validated in one bounded pass before
 * anything lands in the accumulator, so a rejection leaves no partial adopted
 * state and the caller can fall back to the sync pipeline.
 */
suspend fun receiveVaultSnapshot(
    channel: PairingTransport,
): ReceivedSnapshot {
    val acc = createAdoptSnapshotBudget()
    var chunkCount = -1
    var received = 0
    while (chunkCount == -1 || received < chunkCount) {
        val wire = channel.receive()
        if (wire.size > MAX_PADDED_SNAPSHOT_FRAME_BYTES) {
            throw IllegalStateException("ADOPT_SNAPSHOT_CHUNK_TOO_LARGE")
        }
        val raw = runCatching { SyncPadding.unpad(wire) }
            .getOrElse { throw IllegalStateException("ADOPT_SNAPSHOT_INVALID_FRAME", it) }
        if (raw.size > MAX_ADOPT_SNAPSHOT_BYTES - acc.bytes) {
            throw IllegalStateException("ADOPT_SNAPSHOT_OVER_BUDGET")
        }
        if (raw.size > SNAPSHOT_FRAME_WIRE_CEILING) {
            throw IllegalStateException("ADOPT_SNAPSHOT_CHUNK_TOO_LARGE")
        }
        val frame = runCatching {
            AdoptJson.decodeFromString(AdoptionFrame.serializer(), String(raw, Charsets.UTF_8))
        }.getOrNull()
        if (!isAdoptSnapshot(frame)) throw IllegalStateException("ADOPT_SNAPSHOT_INVALID_FRAME")
        frame as AdoptionFrame.Snapshot
        if (chunkCount == -1) {
            chunkCount = frame.chunkCount
            if (chunkCount > MAX_SNAPSHOT_CHUNKS) {
                throw IllegalStateException("ADOPT_SNAPSHOT_TOO_MANY_CHUNKS")
            }
        } else if (frame.chunkCount != chunkCount) {
            throw IllegalStateException("ADOPT_SNAPSHOT_CHUNK_COUNT_MISMATCH")
        }
        if (frame.chunkIndex != received) {
            throw IllegalStateException("ADOPT_SNAPSHOT_CHUNK_OUT_OF_ORDER")
        }
        applyAdoptSnapshotFrame(acc, frame.entries, frame.tombstones, frame.duressSlot, raw.size)
        received++
    }
    return ReceivedSnapshot(acc.entries, acc.tombstones, acc.duressSlot)
}

fun isDeviceType(value: String): Boolean =
    DeviceType.entries.any { it.wireValue == value }