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

/**
 * 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.
private const val SNAPSHOT_FRAME_WIRE_CEILING = MAX_SNAPSHOT_CHUNK_BYTES + 256 * 1024

/** 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) {}

    /** Seal a fresh local vault under the receiver's password, reusing [rootKeyHex]. Returns the vault id. */
    suspend fun createVaultFromRootKey(rootKeyHex: String): String
    suspend fun deriveLocalIdentity(): LocalDeviceIdentity
    /** Only called after a successful, SAS-confirmed direct pairing. */
    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>) {}
    /** 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 upsertPairedDevice(device: PairedDevice)
    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) {}
}

private 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"

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,
    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,
        license = receiverLicense,
        snapshotCapable = true,
    ))

    val confirmed = deps.confirmSas()
    sendFrame(channel, AdoptionFrame.Confirm(confirmed))
    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

    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
    }

    deps.validateSource(grant.source)
    val vaultId = deps.createVaultFromRootKey(grant.rootKeyHex)
    if (vaultId != grant.vaultId) {
        deps.onError("ADOPT_VAULT_ID_MISMATCH")
        return null
    }

    if (receiverLicense == null) {
        grant.license?.let { deps.applyLicense(it) }
    }

    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 = grant.source.deviceId,
        deviceName = grant.source.deviceName,
        deviceType = grant.source.deviceType,
        longTermPublicKeyHex = grant.source.longTermPublicKeyHex,
        signingPublicKeyHex = grant.source.signingPublicKeyHex,
        trustSource = "direct",
        wsPort = grant.source.wsPort,
        rendezvousSecretHex = deps.pairState?.rendezvousSecretHex,
        relayUrl = deps.pairState?.relayUrl ?: grant.source.relayUrl,
        pairedAt = nowMs(),
    )
    // A physically confirmed direct re-pair is the sole path allowed to lift
    // terminal revocation. Clear before upsert so storage can never expose the
    // source as both paired and revoked, even transiently.
    deps.clearPairedDeviceRevocation(sourceDevice.deviceId)
    deps.upsertPairedDevice(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)
    if (!isAdoptConfirm(confirmation) || !(confirmation as AdoptionFrame.Confirm).confirmed) {
        return null
    }

    var sourceLicense = deps.license()
    if (sourceLicense == null && request.license != null) {
        sourceLicense = runCatching {
            deps.applyLicense(request.license)
            deps.license()
        }.getOrNull()
    }
    if (sourceLicense == null) {
        sendFrame(channel, AdoptionFrame.Reject("ADOPT_PRO_REQUIRED"))
        throw IllegalStateException("ADOPT_PRO_REQUIRED")
    }

    val rootKeyHex = deps.rootKeyHex()
    if (!validateRootKeyHex(rootKeyHex)) throw IllegalStateException("ADOPT_INVALID_ROOT_KEY")

    sendFrame(channel, AdoptionFrame.Grant(
        rootKeyHex = rootKeyHex,
        vaultId = deps.vaultId(),
        license = sourceLicense,
        source = AdoptionFrame.Source(
            deviceId = deps.sourceIdentity.deviceId,
            deviceName = deps.sourceIdentity.deviceName,
            deviceType = deps.sourceIdentity.deviceType.wireValue,
            longTermPublicKeyHex = deps.sourceIdentity.longTermPublicKeyHex,
            signingPublicKeyHex = deps.sourceIdentity.signingPublicKeyHex,
            wsPort = deps.sourceIdentity.wsPort,
            relayUrl = deps.pairState?.relayUrl,
        ),
    ))

    val identity = receiveFrame(channel)
    if (!isAdoptIdentity(identity)) return null
    identity as AdoptionFrame.Identity

    val 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(),
    )
    deps.upsertPairedDevice(device)
    sendFrame(channel, AdoptionFrame.Complete)
    if (!isAdoptApplied(receiveFrame(channel))) 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.
 * Entries are split so each chunk's serialized entries map stays under
 * [MAX_SNAPSHOT_CHUNK_BYTES]; tombstones ride the final chunk. If the vault
 * needs more than [MAX_SNAPSHOT_CHUNKS] chunks, only the first that many are
 * sent and the fallback sync pipeline fills the remainder.
 */
suspend fun sendVaultSnapshot(
    channel: PairingTransport,
    entries: Map<String, String>,
    tombstones: Map<String, Long>,
) {
    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 chunkCount = minOf(chunks.size, MAX_SNAPSHOT_CHUNKS)
    for (index in 0 until chunkCount) {
        val isFinal = index == chunkCount - 1
        sendFrame(channel, AdoptionFrame.Snapshot(
            chunkIndex = index,
            chunkCount = chunkCount,
            entries = chunks[index],
            tombstones = if (isFinal) tombstones else emptyMap(),
        ))
    }
}

/**
 * Receiver side: receive and reassemble every snapshot chunk. Validates the
 * encoded frame size before parsing, then enforces sequential 0-based
 * chunk_index ordering, a stable chunk_count within [MAX_SNAPSHOT_CHUNKS], and
 * structural validity of each frame. Throws on any validation failure so the
 * caller can fall back to the sync pipeline.
 */
suspend fun receiveVaultSnapshot(
    channel: PairingTransport,
): Pair<Map<String, String>, Map<String, Long>> {
    val entries = LinkedHashMap<String, String>()
    val tombstones = LinkedHashMap<String, Long>()
    var chunkCount = -1
    var received = 0
    while (chunkCount == -1 || received < chunkCount) {
        val raw = channel.receive()
        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")
        }
        entries.putAll(frame.entries)
        tombstones.putAll(frame.tombstones)
        received++
    }
    return entries to tombstones
}

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