import Foundation
import CryptoKit

/// The JSON envelope carried inside relay `signal` frames and data-channel text
/// messages: `{"v":1,"iv":"<base64>","ct":"<base64 ciphertext‖tag>"}`, keys in
/// exactly that order.
public struct Envelope: Sendable, Hashable {
    public static let version = 1
    public let iv: Data
    /// AES-GCM ciphertext followed by the 16-byte tag.
    public let ciphertext: Data

    public init(iv: Data, ciphertext: Data) {
        self.iv = iv
        self.ciphertext = ciphertext
    }

    /// The wire form, key order `v`, `iv`, `ct`.
    public var json: JSONValue {
        .object(JSONObject([("v", .int(Envelope.version)), ("iv", .string(Bytes.base64(iv))), ("ct", .string(Bytes.base64(ciphertext)))]))
    }

    /// Strict decoding: `v` must be exactly 1 and `iv`/`ct` strings; `nil` otherwise.
    public init?(json: JSONValue) {
        guard case .object(let o) = json,
              o["v"]?.intValue == Envelope.version,
              let ivS = o["iv"]?.stringValue, let ctS = o["ct"]?.stringValue,
              let iv = Bytes.fromBase64(ivS), let ct = Bytes.fromBase64(ctS)
        else { return nil }
        self.iv = iv
        self.ciphertext = ct
    }
}

/// Byte prefixes of the two binary envelope layouts: `prefix ‖ iv(12) ‖ ct‖tag`.
public enum BinaryEnvelopePrefix {
    /// Session-key binary envelope (data channel).
    public static let session: UInt8 = 0x01
    /// Room-wide at-rest envelope (mailbox).
    public static let atRest: UInt8 = 0x02
    /// Smallest possible envelope: prefix + IV + tag.
    public static let minimumLength = 1 + 12 + 16
}

enum AEAD {
    static func seal(_ plaintext: Data, key: SymmetricKey, iv: Data, aad: Data) -> Data {
        // A 12-byte nonce never fails to construct; a programming error here is unrecoverable.
        let box = try! AES.GCM.seal(plaintext, using: key, nonce: AES.GCM.Nonce(data: iv), authenticating: aad)
        return box.ciphertext + box.tag
    }

    /// `nil` on any failure — the caller drops silently (PROTOCOL.md).
    static func open(_ ciphertextAndTag: Data, key: SymmetricKey, iv: Data, aad: Data) -> Data? {
        guard iv.count == 12, ciphertextAndTag.count >= 16 else { return nil }
        let ct = ciphertextAndTag.prefix(ciphertextAndTag.count - 16)
        let tag = ciphertextAndTag.suffix(16)
        guard let nonce = try? AES.GCM.Nonce(data: iv),
              let box = try? AES.GCM.SealedBox(nonce: nonce, ciphertext: ct, tag: tag)
        else { return nil }
        return try? AES.GCM.open(box, using: key, authenticating: aad)
    }

    static func binary(prefix: UInt8, iv: Data, ciphertext: Data) -> Data {
        var out = Data(capacity: 1 + iv.count + ciphertext.count)
        out.append(prefix)
        out.append(iv)
        out.append(ciphertext)
        return out
    }

    static func splitBinary(_ data: Data, expectedPrefix: UInt8) -> (iv: Data, ciphertext: Data)? {
        guard data.count >= BinaryEnvelopePrefix.minimumLength, data.first == expectedPrefix else { return nil }
        let start = data.startIndex
        return (Data(data[start + 1 ..< start + 13]), Data(data[(start + 13)...]))
    }
}