import Foundation
import CryptoKit

/// iOS-side ECIES P-256 keybox (Path A minimal envelope).
///
/// This is **not** byte-compatible with Android's Tink-generated `keybox:v1:`.
/// It is a self-contained iOS envelope that proves the CryptoKit composition
/// (P256.KeyAgreement + HKDF-SHA256 + AES-128-GCM + HMAC-SHA256) is correct in
/// isolation. Cross-platform interop requires either reverse-engineering Tink's
/// envelope (Path B) or a server-side helper (deferred to Batch 5).
///
/// Wire format: `keybox:v1:{urlsafe-base64-no-padding(JSON)}`
/// where JSON = `{"v":1,"pub":"<base64url-65-byte-uncompressed-P-256>",
///               "ct":"<base64url-AES-128-GCM-ciphertext>",
///               "mac":"<base64url-HMAC-SHA256>"}`
public struct Keybox: Sendable {
    /// 65-byte uncompressed P-256 public key (0x04 || X || Y).
    public let ephemeralPublicKey: Data
    /// AES-128-GCM ciphertext.
    public let ciphertext: Data
    /// HMAC-SHA256 tag over `pub || ct`.
    public let mac: Data

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

/// ECIES P-256 keybox wrapper/unwrapper using CryptoKit.
public enum KeyboxCrypto {
    public static let keyboxPrefix = "keybox:v1:"
    public static let keyboxVersion = 1

    /// Wraps a plaintext blob for a recipient's P-256 public key.
    ///
    /// Steps:
    /// 1. Generate an ephemeral P-256 keypair.
    /// 2. ECDH with the recipient's public key → shared secret.
    /// 3. HKDF-SHA256 over the shared secret with `info` to derive 64 bytes.
    ///    - First 16 bytes: AES-128-GCM key.
    ///    - Second 16 bytes: HMAC-SHA256 key.
    ///    - Remaining 32 bytes unused (reserved).
    /// 4. AES-128-GCM encrypt the plaintext with a random 12-byte nonce.
    /// 5. HMAC-SHA256 over `ephemeralPub || ciphertext` with the MAC key.
    public static func wrap(
        plaintext: Data,
        recipientPublicKey: P256.KeyAgreement.PublicKey,
        info: Data
    ) throws -> Keybox {
        let ephemeral = P256.KeyAgreement.PrivateKey()
        let sharedSecret = try ephemeral.sharedSecretFromKeyAgreement(with: recipientPublicKey)
        let sharedKey = sharedSecret.x963DerivedSymmetricKey(
            using: SHA256.self,
            sharedInfo: Data(),
            outputByteCount: 32
        )
        let derived = try HKDF<SHA256>.deriveKey(
            inputKeyMaterial: sharedKey,
            salt: Data(),
            info: info,
            outputByteCount: 64
        )
        var derivedBytes = [UInt8](repeating: 0, count: 64)
        derived.withUnsafeBytes { raw in
            derivedBytes.replaceSubrange(0..<raw.count, with: raw)
        }
        let aesKey = SymmetricKey(data: Data(derivedBytes[0..<16]))
        let macKey = Data(derivedBytes[16..<32])

        let nonce = AES.GCM.Nonce()
        let sealed = try AES.GCM.seal(plaintext, using: aesKey, nonce: nonce)
        guard let ct = sealed.combined else {
            throw KeyboxError.missingCombinedCiphertext
        }

        let pub = ephemeral.publicKey.x963Representation
        var macData = pub
        macData.append(ct)
        let mac = HMAC<SHA256>.authenticationCode(for: macData, using: SymmetricKey(data: macKey))

        return Keybox(
            ephemeralPublicKey: pub,
            ciphertext: ct,
            mac: Data(mac)
        )
    }

    /// Unwraps a keybox with the recipient's private key.
    public static func unwrap(
        _ keybox: Keybox,
        recipientPrivateKey: P256.KeyAgreement.PrivateKey,
        info: Data
    ) throws -> Data {
        let sharedSecret = try recipientPrivateKey.sharedSecretFromKeyAgreement(
            with: P256.KeyAgreement.PublicKey(x963Representation: keybox.ephemeralPublicKey)
        )
        let sharedKey = sharedSecret.x963DerivedSymmetricKey(
            using: SHA256.self,
            sharedInfo: Data(),
            outputByteCount: 32
        )
        let derived = try HKDF<SHA256>.deriveKey(
            inputKeyMaterial: sharedKey,
            salt: Data(),
            info: info,
            outputByteCount: 64
        )
        var derivedBytes = [UInt8](repeating: 0, count: 64)
        derived.withUnsafeBytes { raw in
            derivedBytes.replaceSubrange(0..<raw.count, with: raw)
        }
        let aesKey = SymmetricKey(data: Data(derivedBytes[0..<16]))
        let macKey = Data(derivedBytes[16..<32])

        var macData = keybox.ephemeralPublicKey
        macData.append(keybox.ciphertext)
        let expectedMAC = HMAC<SHA256>.authenticationCode(for: macData, using: SymmetricKey(data: macKey))
        guard keybox.mac == Data(expectedMAC) else {
            throw KeyboxError.macMismatch
        }

        let sealedBox = try AES.GCM.SealedBox(combined: keybox.ciphertext)
        return try AES.GCM.open(sealedBox, using: aesKey)
    }

    /// Encodes a keybox to the `keybox:v1:` wire string.
    public static func encode(_ keybox: Keybox) throws -> String {
        let envelope: [String: Any] = [
            "v": keyboxVersion,
            "pub": urlsafeBase64NoPadding(keybox.ephemeralPublicKey),
            "ct": urlsafeBase64NoPadding(keybox.ciphertext),
            "mac": urlsafeBase64NoPadding(keybox.mac)
        ]
        let json = try JSONSerialization.data(withJSONObject: envelope, options: [.sortedKeys])
        return keyboxPrefix + urlsafeBase64NoPadding(json)
    }

    /// Decodes a `keybox:v1:` wire string back to a `Keybox`.
    public static func decode(_ blob: String) throws -> Keybox {
        guard blob.hasPrefix(keyboxPrefix) else {
            throw KeyboxError.missingPrefix
        }
        let b64 = String(blob.dropFirst(keyboxPrefix.count))
        guard let data = urlsafeBase64Decode(b64),
              let envelope = try JSONSerialization.jsonObject(with: data) as? [String: Any],
              let version = envelope["v"] as? Int,
              version == keyboxVersion,
              let pubB64 = envelope["pub"] as? String,
              let ctB64 = envelope["ct"] as? String,
              let macB64 = envelope["mac"] as? String,
              let pub = urlsafeBase64Decode(pubB64),
              let ct = urlsafeBase64Decode(ctB64),
              let mac = urlsafeBase64Decode(macB64) else {
            throw KeyboxError.invalidEnvelope
        }
        return Keybox(ephemeralPublicKey: pub, ciphertext: ct, mac: mac)
    }

    public enum KeyboxError: Error {
        case missingCombinedCiphertext
        case macMismatch
        case missingPrefix
        case invalidEnvelope
    }
}

private func urlsafeBase64NoPadding(_ data: Data) -> String {
    data.base64EncodedString()
        .replacingOccurrences(of: "+", with: "-")
        .replacingOccurrences(of: "/", with: "_")
        .trimmingCharacters(in: CharacterSet(charactersIn: "="))
}

private func urlsafeBase64Decode(_ s: String) -> Data? {
    var b64 = s
        .replacingOccurrences(of: "-", with: "+")
        .replacingOccurrences(of: "_", with: "/")
    let padding = (4 - b64.count % 4) % 4
    b64.append(String(repeating: "=", count: padding))
    return Data(base64Encoded: b64)
}