package com.peachpasswords.sync

import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable

/**
 * Adoption transcript v2.
 *
 * v1 hands over the root-bearing `ADOPT_GRANT` before the receiver's identity
 * has been exchanged or validated, so the source's pending record is keyed on a
 * session label that is explicitly not an identity — the source can list and
 * revoke nothing durable. v2 places a validated identity before every
 * root-bearing frame.
 *
 * Frozen contract: `peach-browser/docs/sync-protocol.md` section 10.
 */

/** Adoption transcript version. A v1 request carries no version field at all. */
const val ADOPTION_VERSION = 2

private const val REASON_PRO_REQUIRED = "ADOPT_PRO_REQUIRED"
private const val REASON_DEVICE_ID_COLLISION = "ADOPT_DEVICE_ID_COLLISION"

/**
 * §11.6: legacy v1 adoption is allowed only before a secure reset. A v1 peer's
 * identity is derived from the legacy root and is therefore forgeable by any
 * former root holder — the exact property the reset exists to retire. Both
 * relay source sessions call this one predicate, so the rule cannot hold on one
 * pairing entry point and not the other (PB-26-117).
 */
internal fun allowLegacyV1(postReset: Boolean): Boolean = !postReset

/**
 * The v2 grant as it appears on the wire.
 *
 * It is deliberately NOT part of [AdoptionFrame]: the frozen v1 `Grant` carries
 * a required `source`, a sealed hierarchy can bind only one class to the
 * `ADOPT_GRANT` discriminator, and widening the v1 type to satisfy a v2 shape
 * would let v1-only code accept a grant it can never receive. The v2 transcript
 * therefore reads and writes this shape directly.
 */
@Serializable
data class AdoptGrantV2(
    val type: String,
    @SerialName("root_key_hex") val rootKeyHex: String,
    @SerialName("vault_id") val vaultId: String,
    val license: String? = null,
)

private const val GRANT_TYPE = "ADOPT_GRANT"

/**
 * What the RECEIVER (joining device) needs to adopt a vault over the v2
 * transcript: the v1 surface plus the v2 additions. Extending the v1 interface
 * is what makes the union compile-time visible — delegating through a cast
 * would fail only once a peer triggered the missing callback.
 */
interface AdoptionReceiverV2Deps : AdoptionReceiverDeps {
    /** The receiver's durable device identity. Available with no vault present. */
    fun localIdentity(): LocalDeviceIdentity

    suspend fun rotateLocalIdentity(): LocalDeviceIdentity

    suspend fun validateSource(prepare: AdoptionFrame.Prepare)

    suspend fun hasSourceIdentityCollision(prepare: AdoptionFrame.Prepare): Boolean
}

/** What the SOURCE (existing device) needs to grant its vault over the v2 transcript. */
interface AdoptionSourceV2Deps : AdoptionSourceDeps {
    suspend fun onConsent(request: AdoptionFrame.RequestV2): Boolean

    // Re-declared abstract: the v1 defaults would silently accept an unrotated
    // identity or throw only after a peer asked for the callback.
    override suspend fun rotateSourceIdentity(): LocalDeviceIdentity

    override fun hasDeviceIdentityCollision(device: PairedDevice): Boolean

    /**
     * Durably record the validated receiver identity before any root-bearing
     * frame. Named apart from the v1 member, which is keyed on a request label.
     */
    suspend fun persistPendingIdentity(device: PairedDevice, nowMs: Long): Boolean

    /** Promote the pending identity to a direct paired device. */
    suspend fun completePendingIdentity(device: PairedDevice): Boolean

    fun snapshotCapable(): Boolean
}

/**
 * Bind a source session's deps to the live session facts.
 *
 * Both relay source adapters — the phone showing a QR, and the phone scanning an
 * empty invitation — build the same object; the only differences are which
 * identity is presented and whether snapshot support is advertised, so both are
 * parameters rather than two near-identical anonymous objects.
 *
 * Extracted from the session functions so §10.5 can be proven: the test dirties
 * every member here and asserts each one reaches its sink. A callback that is
 * dropped, or one that returns without calling through, fails that test.
 */
internal fun bindSourceSessionDeps(
    sourceDeps: AdoptionSourceV2Deps,
    sourceIdentity: LocalDeviceIdentity,
    pairState: AdoptionPairState,
    sas: String,
    snapshotCapable: Boolean,
    confirm: suspend (sas: String, request: AdoptionFrame.Request) -> Boolean,
): AdoptionSourceV2Deps = object : AdoptionSourceV2Deps {
    override val sourceIdentity = sourceIdentity
    override val pairState = pairState
    override suspend fun rootKeyHex() = sourceDeps.rootKeyHex()
    override suspend fun vaultId() = sourceDeps.vaultId()
    override suspend fun license() = sourceDeps.license()
    override suspend fun applyLicense(license: String) = sourceDeps.applyLicense(license)
    override suspend fun onConsent(request: AdoptionFrame.Request) = true

    // The transcript sends ADOPT_CONFIRM_REQUIRED before calling this, so the
    // receiver's own prompt is already live when we ask our human to compare the
    // same SAS. The request is only a prompt label here — but the call MUST
    // reach `confirm`: returning true without it was a real defect, because the
    // source-side prompt then never appeared.
    override suspend fun onConsent(request: AdoptionFrame.RequestV2): Boolean =
        confirm(
            sas,
            AdoptionFrame.Request(
                sessionRequestId = request.deviceId,
                deviceName = request.deviceName,
                deviceType = request.deviceType,
                license = request.license,
                snapshotCapable = request.snapshotCapable,
            ),
        )

    override suspend fun persistPendingHolder(request: AdoptionFrame.Request, nowMs: Long) =
        sourceDeps.persistPendingHolder(request, nowMs)

    override suspend fun completePendingHolder(request: AdoptionFrame.Request, device: PairedDevice) =
        sourceDeps.completePendingHolder(request, device)

    override suspend fun persistPendingIdentity(device: PairedDevice, nowMs: Long) =
        sourceDeps.persistPendingIdentity(device, nowMs)

    override suspend fun completePendingIdentity(device: PairedDevice) =
        sourceDeps.completePendingIdentity(device)

    override suspend fun upsertPairedDevice(device: PairedDevice) =
        sourceDeps.upsertPairedDevice(device)

    override fun hasDeviceIdentityCollision(device: PairedDevice) =
        sourceDeps.hasDeviceIdentityCollision(device)

    override suspend fun rotateSourceIdentity() = sourceDeps.rotateSourceIdentity()
    override fun snapshotCapable() = snapshotCapable
    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()
}

private fun AdoptionFrame.RequestV2.valid(): Boolean =
    adoptionVersion == ADOPTION_VERSION &&
        deviceId.isNotEmpty() &&
        Hex64.matches(longTermPublicKeyHex) &&
        Hex64.matches(signingPublicKeyHex) &&
        isDeviceType(deviceType)

private fun AdoptionFrame.Prepare.valid(): Boolean =
    deviceId.isNotEmpty() &&
        Hex64.matches(longTermPublicKeyHex) &&
        Hex64.matches(signingPublicKeyHex) &&
        isDeviceType(deviceType) &&
        Hex64.matches(vaultId)

/**
 * RECEIVER side: present our durable identity, then adopt the vault the source
 * prepares. Returns null on any protocol failure, with the reason via
 * [AdoptionReceiverV2Deps.onError].
 */
suspend fun requestAdoptionV2(
    channel: PairingTransport,
    deps: AdoptionReceiverV2Deps,
    // Called with the accepted ADOPT_PREPARE once it has been validated, so the
    // caller can read the source's capabilities before the root is granted.
    onPrepare: ((AdoptionFrame.Prepare) -> Unit)? = null,
    nowMs: Long = System.currentTimeMillis(),
): PairedDevice? {
    val receiverLicense = deps.license()

    fun requestFrame(identity: LocalDeviceIdentity): AdoptionFrame.RequestV2 =
        AdoptionFrame.RequestV2(
            adoptionVersion = ADOPTION_VERSION,
            deviceId = identity.deviceId,
            deviceName = identity.deviceName,
            deviceType = identity.deviceType.wireValue,
            longTermPublicKeyHex = identity.longTermPublicKeyHex,
            signingPublicKeyHex = identity.signingPublicKeyHex,
            snapshotCapable = true,
        )

    // A malformed root key must not escape as an exception; every failure here
    // is the same protocol outcome.
    fun rootKeyMatchesVault(rootKeyHex: String, vaultId: String): Boolean {
        if (!validateRootKeyHex(rootKeyHex) || !Hex64.matches(vaultId)) return false
        val rootKeyBytes = hexToPubkey(rootKeyHex)
        return try {
            computeVaultId(rootKeyBytes) == vaultId
        } finally {
            rootKeyBytes.fill(0)
        }
    }

    sendFrame(channel, requestFrame(deps.localIdentity()))

    if (receiveFrame(channel) !is AdoptionFrame.ConfirmRequired) {
        deps.onError("ADOPT_UNEXPECTED_FRAME")
        return null
    }

    val confirmed = deps.confirmSas()
    // The Pro key travels only after the SAS decision (PB-26-121): a peer that
    // controls the scanned invitation reads the request frame before the user
    // has compared anything.
    sendFrame(channel, AdoptionFrame.Confirm(confirmed, license = receiverLicense))
    if (!confirmed) {
        deps.onError("ADOPT_SAS_REJECTED")
        return null
    }

    // The source rejects our identity once if it is already in its terminal
    // ledger; we rotate exactly once and present the fresh identity.
    var prepared = receiveFrame(channel)
    if (prepared is AdoptionFrame.Reject && prepared.reason == REASON_DEVICE_ID_COLLISION) {
        sendFrame(channel, requestFrame(deps.rotateLocalIdentity()))
        prepared = receiveFrame(channel)
    }
    if (prepared is AdoptionFrame.Reject) {
        deps.onError(prepared.reason)
        return null
    }
    if (prepared !is AdoptionFrame.Prepare || !prepared.valid()) {
        deps.onError("ADOPT_UNEXPECTED_FRAME")
        return null
    }

    deps.validateSource(prepared)

    // The source's identity may be in our terminal ledger. Reject once, accept
    // only a rotated preparation, and treat a second collision as terminal.
    if (deps.hasSourceIdentityCollision(prepared)) {
        sendFrame(channel, AdoptionFrame.Reject(REASON_DEVICE_ID_COLLISION))
        val rotated = receiveFrame(channel)
        if (rotated !is AdoptionFrame.Prepare || !rotated.valid()) {
            deps.onError("ADOPT_SOURCE_COLLISION_NOT_RESOLVED")
            return null
        }
        deps.validateSource(rotated)
        if (deps.hasSourceIdentityCollision(rotated)) {
            sendFrame(channel, AdoptionFrame.Reject(REASON_DEVICE_ID_COLLISION))
            deps.onError("ADOPT_SOURCE_COLLISION_NOT_RESOLVED")
            return null
        }
        prepared = rotated
    }

    if (receiverLicense == null && prepared.license == null) {
        deps.onError(REASON_PRO_REQUIRED)
        return null
    }

    onPrepare?.invoke(prepared)

    sendFrame(channel, AdoptionFrame.Ready)

    val raw = String(channel.receive(), Charsets.UTF_8)
    // Either the v2 grant or a terminal rejection.
    val rejection = runCatching {
        AdoptJson.decodeFromString(AdoptionFrame.serializer(), raw)
    }.getOrNull() as? AdoptionFrame.Reject
    if (rejection != null) {
        deps.onError(rejection.reason)
        return null
    }
    val grant = runCatching {
        AdoptJson.decodeFromString(AdoptGrantV2.serializer(), raw)
    }.getOrNull()
    if (grant == null || grant.type != GRANT_TYPE) {
        deps.onError("ADOPT_UNEXPECTED_FRAME")
        return null
    }
    // The source may not change the vault between preparing and granting, and
    // the granted root key must actually derive that vault id.
    if (grant.vaultId != prepared.vaultId || !rootKeyMatchesVault(grant.rootKeyHex, grant.vaultId)) {
        deps.onError("ADOPT_VAULT_ID_MISMATCH")
        return null
    }

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

    // The source is recorded from the validated preparation, never from the
    // grant: the preparation is what we validated against our terminal ledger.
    val sourceDevice = PairedDevice(
        deviceId = prepared.deviceId,
        deviceName = prepared.deviceName,
        deviceType = prepared.deviceType,
        longTermPublicKeyHex = prepared.longTermPublicKeyHex,
        signingPublicKeyHex = prepared.signingPublicKeyHex,
        trustSource = "direct",
        wsPort = prepared.wsPort,
        rendezvousSecretHex = deps.pairState?.rendezvousSecretHex,
        relayUrl = deps.pairState?.relayUrl ?: prepared.relayUrl,
        pairedAt = nowMs,
    )
    deps.upsertPairedDeviceAndClearRevocation(sourceDevice)
    sendFrame(channel, AdoptionFrame.Applied)

    if (receiveFrame(channel) !is AdoptionFrame.Complete) {
        deps.onError("ADOPT_COMPLETION_MISSING")
        return null
    }
    return sourceDevice
}

/**
 * SOURCE side: answer a v2 request. The receiver's validated identity is
 * durably recorded before any root-bearing frame, and is promoted to a direct
 * paired device only once the receiver acknowledges the vault.
 */
suspend fun grantAdoptionV2(
    channel: PairingTransport,
    request: AdoptionFrame.RequestV2,
    deps: AdoptionSourceV2Deps,
    nowMs: Long = System.currentTimeMillis(),
): PairedDevice? {
    if (!request.valid()) return null

    // Acknowledge before asking the human, so the receiver can raise its own
    // SAS prompt instead of waiting on this device's consent dialog.
    sendFrame(channel, AdoptionFrame.ConfirmRequired)
    if (!deps.onConsent(request)) return null

    val confirmation = receiveFrame(channel)
    if (confirmation !is AdoptionFrame.Confirm || !confirmation.confirmed) return null

    fun receiverFrom(identity: AdoptionFrame.RequestV2): PairedDevice =
        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,
        )

    // The receiver's identity may collide with a paired device or a terminal
    // revocation. Reject once, accept only a rotated request, then stop.
    var receiver = receiverFrom(request)
    if (deps.hasDeviceIdentityCollision(receiver)) {
        sendFrame(channel, AdoptionFrame.Reject(REASON_DEVICE_ID_COLLISION))
        val rotated = receiveFrame(channel)
        if (rotated !is AdoptionFrame.RequestV2 || !rotated.valid()) return null
        receiver = receiverFrom(rotated)
        if (deps.hasDeviceIdentityCollision(receiver)) {
            sendFrame(channel, AdoptionFrame.Reject(REASON_DEVICE_ID_COLLISION))
            return null
        }
    }

    var license = 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 = confirmation.license ?: request.license
    if (license == null && receiverLicense != null) {
        deps.applyLicense(receiverLicense)
        license = deps.license()
    }
    if (license == null) {
        sendFrame(channel, AdoptionFrame.Reject(REASON_PRO_REQUIRED))
        return null
    }

    // No root-bearing frame before this commit: afterwards the receiver is
    // listable and revocable even if the session dies mid-grant.
    if (!deps.persistPendingIdentity(receiver, nowMs)) return null

    val rootKeyHex = deps.rootKeyHex()
    if (!validateRootKeyHex(rootKeyHex)) return null

    var sourceIdentity = deps.sourceIdentity
    var sourceRotations = 0
    while (true) {
        sendFrame(
            channel,
            AdoptionFrame.Prepare(
                deviceId = sourceIdentity.deviceId,
                deviceName = sourceIdentity.deviceName,
                deviceType = sourceIdentity.deviceType.wireValue,
                longTermPublicKeyHex = sourceIdentity.longTermPublicKeyHex,
                signingPublicKeyHex = sourceIdentity.signingPublicKeyHex,
                vaultId = deps.vaultId(),
                wsPort = sourceIdentity.wsPort,
                relayUrl = deps.pairState?.relayUrl,
                license = license,
                snapshotCapable = deps.snapshotCapable(),
            ),
        )

        val next = receiveFrame(channel)
        if (next is AdoptionFrame.Reject &&
            next.reason == REASON_DEVICE_ID_COLLISION &&
            sourceRotations == 0
        ) {
            sourceIdentity = deps.rotateSourceIdentity()
            sourceRotations++
            continue
        }
        if (next !is AdoptionFrame.Ready) return null
        break
    }

    // Sent with its own serializer: the v2 grant carries no source, so it is
    // not a member of the AdoptionFrame hierarchy (see AdoptGrantV2).
    channel.send(
        AdoptJson.encodeToString(
            AdoptGrantV2.serializer(),
            AdoptGrantV2(
                type = GRANT_TYPE,
                rootKeyHex = rootKeyHex,
                vaultId = deps.vaultId(),
                license = license,
            ),
        ).toByteArray(Charsets.UTF_8),
    )

    if (receiveFrame(channel) !is AdoptionFrame.Applied) return null

    // Promotion is legal only now: the receiver holds the vault and has said so.
    if (!deps.completePendingIdentity(receiver)) return null
    deps.upsertPairedDevice(receiver)
    sendFrame(channel, AdoptionFrame.Complete)
    return receiver
}