import Foundation
import CryptoKit

/// X3DH key agreement protocol (Signal Protocol)
enum X3DH {

    // MARK: - Pre-Key Generation

    /// Generate a signed pre-key (SPK).
    /// Returns (private, public, signature, id).
    /// Matches Python: generate_signed_prekey(identity_private)
    static func generateSignedPrekey(
        identityPrivate: Curve25519.Signing.PrivateKey
    ) throws -> (privateKey: Curve25519.KeyAgreement.PrivateKey,
                  publicKey: Curve25519.KeyAgreement.PublicKey,
                  signature: Data,
                  id: String) {
        let (spkPriv, spkPub) = X25519Crypto.generateKeypair()
        let spkPubBytes = X25519Crypto.serializePublic(spkPub)
        let signature = try Ed25519Crypto.sign(identityPrivate, data: spkPubBytes)
        return (spkPriv, spkPub, signature, UUID().uuidString)
    }

    /// Generate a batch of one-time pre-keys.
    /// Matches Python: generate_one_time_prekeys(count=50)
    static func generateOneTimePrekeys(count: Int = 50) -> [(privateKey: Curve25519.KeyAgreement.PrivateKey,
                                                               publicKey: Curve25519.KeyAgreement.PublicKey,
                                                               id: String)] {
        (0..<count).map { _ in
            let (priv, pub) = X25519Crypto.generateKeypair()
            return (priv, pub, UUID().uuidString)
        }
    }

    // MARK: - X3DH Initiate (Alice)

    /// Initiator side of X3DH.
    /// Returns (sharedSecret, ephemeralPrivate, ephemeralPublic).
    /// Matches Python: x3dh_initiate(ik_private_ed, ik_public_remote_ed, spk_remote, spk_signature, opk_remote?)
    static func initiate(
        ikPrivateEd: Curve25519.Signing.PrivateKey,
        ikPublicRemoteEd: Curve25519.Signing.PublicKey,
        spkRemote: Curve25519.KeyAgreement.PublicKey,
        spkSignature: Data,
        opkRemote: Curve25519.KeyAgreement.PublicKey? = nil
    ) throws -> (sharedSecret: Data,
                  ephemeralPrivate: Curve25519.KeyAgreement.PrivateKey,
                  ephemeralPublic: Curve25519.KeyAgreement.PublicKey) {
        // Verify SPK signature
        let spkRemoteBytes = X25519Crypto.serializePublic(spkRemote)
        guard Ed25519Crypto.verify(ikPublicRemoteEd, signature: spkSignature, data: spkRemoteBytes) else {
            throw CryptoError.x3dhFailed("Invalid SPK signature")
        }

        // Convert identity keys to X25519
        let ikX25519Private = try X25519Crypto.fromEd25519Private(ikPrivateEd)
        let ikX25519Remote = try X25519Crypto.fromEd25519Public(ikPublicRemoteEd)

        // Generate ephemeral keypair
        let (ekPriv, ekPub) = X25519Crypto.generateKeypair()

        // Debug: print key inputs (matching Python x3dh_respond)
        #if DEBUG
        print("DEBUG x3dh_initiate: ik_remote_ed = \(Ed25519Crypto.serializePublic(ikPublicRemoteEd).hexString)")
        print("DEBUG x3dh_initiate: ik_x25519_remote = \(X25519Crypto.serializePublic(ikX25519Remote).hexString)")
        print("DEBUG x3dh_initiate: ek_pub = \(X25519Crypto.serializePublic(ekPub).hexString)")
        print("DEBUG x3dh_initiate: spk_remote = \(spkRemoteBytes.hexString)")
        #endif

        // DH computations
        let dh1 = try X25519Crypto.dh(ikX25519Private, spkRemote)   // IK_A, SPK_B
        let dh2 = try X25519Crypto.dh(ekPriv, ikX25519Remote)        // EK_A, IK_B
        let dh3 = try X25519Crypto.dh(ekPriv, spkRemote)             // EK_A, SPK_B

        // Debug: print DH outputs
        #if DEBUG
        print("DEBUG x3dh_initiate: dh1 = \(dh1.hexString)")
        print("DEBUG x3dh_initiate: dh2 = \(dh2.hexString)")
        print("DEBUG x3dh_initiate: dh3 = \(dh3.hexString)")
        #endif

        var dhConcat = dh1 + dh2 + dh3
        if let opk = opkRemote {
            let dh4 = try X25519Crypto.dh(ekPriv, opk)               // EK_A, OPK_B
            #if DEBUG
            print("DEBUG x3dh_initiate: dh4 = \(dh4.hexString)")
            #endif
            dhConcat += dh4
        }

        // Derive shared secret
        let sharedSecret = CryptoUtils.hkdfDerive(
            inputKey: dhConcat,
            salt: Data(repeating: 0x00, count: 32),
            info: Data(Constants.x3dhInfo.utf8),
            length: 32
        )
        #if DEBUG
        print("DEBUG x3dh_initiate: shared_secret = \(sharedSecret.hexString)")
        #endif

        return (sharedSecret, ekPriv, ekPub)
    }

    // MARK: - X3DH Respond (Bob)

    /// Responder side of X3DH.
    /// Returns sharedSecret.
    /// Matches Python: x3dh_respond(ik_private_ed, spk_private, ik_remote_ed, ek_remote, opk_private?)
    static func respond(
        ikPrivateEd: Curve25519.Signing.PrivateKey,
        spkPrivate: Curve25519.KeyAgreement.PrivateKey,
        ikRemoteEd: Curve25519.Signing.PublicKey,
        ekRemote: Curve25519.KeyAgreement.PublicKey,
        opkPrivate: Curve25519.KeyAgreement.PrivateKey? = nil
    ) throws -> Data {
        let ikX25519Private = try X25519Crypto.fromEd25519Private(ikPrivateEd)
        let ikX25519Remote = try X25519Crypto.fromEd25519Public(ikRemoteEd)

        let dh1 = try X25519Crypto.dh(spkPrivate, ikX25519Remote)   // SPK_B, IK_A
        let dh2 = try X25519Crypto.dh(ikX25519Private, ekRemote)     // IK_B, EK_A
        let dh3 = try X25519Crypto.dh(spkPrivate, ekRemote)          // SPK_B, EK_A

        var dhConcat = dh1 + dh2 + dh3
        if let opk = opkPrivate {
            let dh4 = try X25519Crypto.dh(opk, ekRemote)             // OPK_B, EK_A
            dhConcat += dh4
        }

        let sharedSecret = CryptoUtils.hkdfDerive(
            inputKey: dhConcat,
            salt: Data(repeating: 0x00, count: 32),
            info: Data(Constants.x3dhInfo.utf8),
            length: 32
        )

        return sharedSecret
    }
}