// HybridCrypto.swift
// XWingMLKEM768X25519 (ML-KEM-768 + X25519) — iOS 26+
// KEM encapsulation → HKDF-SHA256 → AES-256-GCM

import Foundation
import CryptoKit

// MARK: - Public Key

public struct HybridPublicKey: Codable, Equatable, Hashable {
    public let rawBytes: Data   // 1216 bytes — XWing public key
    public let peerID: String

    public var encoded: Data { rawBytes }

    public static func from(encoded: Data, peerID: String) -> HybridPublicKey? {
        guard encoded.count == 1216 else { return nil }
        return HybridPublicKey(rawBytes: Data(encoded), peerID: peerID)
    }
}

// MARK: - Private Key

public struct HybridPrivateKey {
    public let xwing: XWingMLKEM768X25519.PrivateKey
}

// MARK: - Encrypted Message (pairwise — KEM-based)

public struct EncryptedMessage: Codable {
    public let kemCiphertext: Data   // XWing encapsulated ciphertext (~1120 bytes)
    public let aesCiphertext: Data   // nonce(12) + ciphertext + tag(16)
}

// MARK: - Crypto Engine

public enum HybridCrypto {

    public static func generateKeyPair(peerID: String) throws -> (publicKey: HybridPublicKey, privateKey: HybridPrivateKey) {
        let priv = try XWingMLKEM768X25519.PrivateKey()
        let pub  = HybridPublicKey(rawBytes: priv.publicKey.rawRepresentation, peerID: peerID)
        return (pub, HybridPrivateKey(xwing: priv))
    }

    // MARK: Pairwise encrypt (for DMs and group key distribution)

    public static func encrypt(message: Data, recipientPublicKey: HybridPublicKey) throws -> EncryptedMessage {
        let recipientPub  = try XWingMLKEM768X25519.PublicKey(rawRepresentation: recipientPublicKey.rawBytes)
        let encapResult   = try recipientPub.encapsulate()

        let symmetricKey  = deriveAESKey(from: encapResult.sharedSecret)

        let nonce  = AES.GCM.Nonce()
        let sealed = try AES.GCM.seal(message, using: symmetricKey, nonce: nonce)

        var aesCT = Data()
        aesCT.append(sealed.nonce.withUnsafeBytes { Data($0) })
        aesCT.append(sealed.ciphertext)
        aesCT.append(sealed.tag)

        return EncryptedMessage(kemCiphertext: encapResult.encapsulated, aesCiphertext: aesCT)
    }

    public static func decrypt(message: EncryptedMessage, recipientPrivateKey: HybridPrivateKey) throws -> Data {
        let sharedSecret = try recipientPrivateKey.xwing.decapsulate(message.kemCiphertext)
        let symmetricKey = deriveAESKey(from: sharedSecret)

        return try aesGCMDecrypt(aesCiphertext: message.aesCiphertext, key: symmetricKey)
    }

    // MARK: Group symmetric encrypt/decrypt

    public static func encryptGroup(plaintext: Data, groupKey: SymmetricKey) throws -> Data {
        let nonce  = AES.GCM.Nonce()
        let sealed = try AES.GCM.seal(plaintext, using: groupKey, nonce: nonce)

        var aesCT = Data()
        aesCT.append(sealed.nonce.withUnsafeBytes { Data($0) })
        aesCT.append(sealed.ciphertext)
        aesCT.append(sealed.tag)
        return aesCT
    }

    public static func decryptGroup(aesCiphertext: Data, groupKey: SymmetricKey) throws -> Data {
        try aesGCMDecrypt(aesCiphertext: aesCiphertext, key: groupKey)
    }

    // MARK: - Internals

    private static func deriveAESKey(from sharedSecret: SymmetricKey) -> SymmetricKey {
        // Extract raw bytes then run HKDF-SHA256
        let raw = sharedSecret.withUnsafeBytes { Data($0) }
        return HKDF<SHA256>.deriveKey(
            inputKeyMaterial: SymmetricKey(data: raw),
            info: Data("btmessage-v1".utf8),
            outputByteCount: 32
        )
    }

    private static func aesGCMDecrypt(aesCiphertext: Data, key: SymmetricKey) throws -> Data {
        guard aesCiphertext.count >= 28 else { throw CryptoError.decryptionFailed }
        let nonceData  = Data(aesCiphertext.prefix(12))
        let tag        = Data(aesCiphertext.suffix(16))
        let ciphertext = Data(aesCiphertext[12..<(aesCiphertext.count - 16)])

        let nonce     = try AES.GCM.Nonce(data: nonceData)
        let sealedBox = try AES.GCM.SealedBox(nonce: nonce, ciphertext: ciphertext, tag: tag)
        return try AES.GCM.open(sealedBox, using: key)
    }
}

// MARK: - Errors

public enum CryptoError: Error, LocalizedError {
    case decryptionFailed
    case invalidKey

    public var errorDescription: String? {
        switch self {
        case .decryptionFailed: return "Decryption failed"
        case .invalidKey:       return "Invalid key"
        }
    }
}