import Foundation
import CryptoKit

/// Seals under one per-session key with IV = 4 random epoch bytes ‖ 8-byte
/// big-endian counter. One sealer per key per room session; the counter never
/// repeats within an epoch and no API accepts a caller-supplied IV.
public struct Sealer: Sendable {
    let key: SymmetricKey
    public let epoch: Data
    public private(set) var counter: UInt64

    init(key: SymmetricKey) {
        self.init(key: key, epoch: Bytes.random(4), counter: 0)
    }

    /// Deterministic sealer — for tests and vector reproduction only.
    init(key: SymmetricKey, epoch: Data, counter: UInt64 = 0) {
        precondition(epoch.count == 4, "epoch is 4 bytes")
        self.key = key
        self.epoch = epoch
        self.counter = counter
    }

    mutating func nextIV() -> Data {
        var iv = Data(capacity: 12)
        iv.append(epoch)
        var be = counter.bigEndian
        withUnsafeBytes(of: &be) { iv.append(contentsOf: $0) }
        counter &+= 1
        return iv
    }

    /// Seals a JSON value (serialized `JSON.stringify`-style) into a JSON envelope.
    public mutating func seal(_ value: JSONValue, aad: Data) -> Envelope {
        let iv = nextIV()
        return Envelope(iv: iv, ciphertext: AEAD.seal(value.serializedData, key: key, iv: iv, aad: aad))
    }

    /// Seals raw bytes into the binary layout `0x01 ‖ iv ‖ ct‖tag`.
    public mutating func sealBinary(_ data: Data, aad: Data) -> Data {
        let iv = nextIV()
        return AEAD.binary(prefix: BinaryEnvelopePrefix.session, iv: iv, ciphertext: AEAD.seal(data, key: key, iv: iv, aad: aad))
    }
}

/// Room-wide at-rest cipher (mailbox payloads): random 12-byte IVs, because
/// many sessions share the key and counters would collide. Layout
/// `0x02 ‖ iv ‖ ct‖tag`.
public struct AtRestCipher: Sendable {
    let key: SymmetricKey
    let aad: Data

    init(key: SymmetricKey, aad: Data) {
        self.key = key
        self.aad = aad
    }

    public func seal(_ data: Data) -> Data {
        seal(data, iv: Bytes.random(12))
    }

    /// Deterministic variant for tests and vector reproduction.
    func seal(_ data: Data, iv: Data) -> Data {
        AEAD.binary(prefix: BinaryEnvelopePrefix.atRest, iv: iv, ciphertext: AEAD.seal(data, key: key, iv: iv, aad: aad))
    }

    /// `nil` on any failure.
    public func open(_ data: Data) -> Data? {
        guard let (iv, ct) = AEAD.splitBinary(data, expectedPrefix: BinaryEnvelopePrefix.atRest) else { return nil }
        return AEAD.open(ct, key: key, iv: iv, aad: aad)
    }
}