import Foundation
import CryptoKit

/// X25519 Diffie-Hellman key agreement
enum X25519Crypto {

    // MARK: - Key Generation

    /// Generate X25519 keypair
    static func generateKeypair() -> (privateKey: Curve25519.KeyAgreement.PrivateKey, publicKey: Curve25519.KeyAgreement.PublicKey) {
        let privateKey = Curve25519.KeyAgreement.PrivateKey()
        return (privateKey, privateKey.publicKey)
    }

    // MARK: - Serialization

    /// Serialize X25519 private key to 32 raw bytes
    static func serializePrivate(_ key: Curve25519.KeyAgreement.PrivateKey) -> Data {
        key.rawData  // 32 bytes
    }

    /// Serialize X25519 public key to 32 raw bytes
    static func serializePublic(_ key: Curve25519.KeyAgreement.PublicKey) -> Data {
        key.rawData  // 32 bytes
    }

    /// Load X25519 private key from 32 raw bytes
    static func loadPrivate(_ data: Data) throws -> Curve25519.KeyAgreement.PrivateKey {
        guard data.count == 32 else {
            throw CryptoError.invalidKeyData("X25519 private key must be 32 bytes")
        }
        return try Curve25519.KeyAgreement.PrivateKey(rawRepresentation: data)
    }

    /// Load X25519 public key from 32 raw bytes
    static func loadPublic(_ data: Data) throws -> Curve25519.KeyAgreement.PublicKey {
        guard data.count == 32 else {
            throw CryptoError.invalidKeyData("X25519 public key must be 32 bytes")
        }
        return try Curve25519.KeyAgreement.PublicKey(rawRepresentation: data)
    }

    // MARK: - Diffie-Hellman

    /// Perform X25519 DH key agreement. Returns 32-byte shared secret.
    /// Matches Python: x25519_dh(private_key, public_key)
    static func dh(_ privateKey: Curve25519.KeyAgreement.PrivateKey, _ publicKey: Curve25519.KeyAgreement.PublicKey) throws -> Data {
        let sharedSecret = try privateKey.sharedSecretFromKeyAgreement(with: publicKey)
        // Extract raw bytes from SharedSecret
        return sharedSecret.withUnsafeBytes { Data($0) }
    }

    // MARK: - Ed25519 → X25519 Key Conversion

    /// Convert Ed25519 private key to X25519 private key.
    /// SHA-512(seed) → take first 32 bytes → clamp per RFC 7748
    /// Matches Python: ed25519_private_to_x25519(ed_private)
    static func fromEd25519Private(_ edPrivate: Curve25519.Signing.PrivateKey) throws -> Curve25519.KeyAgreement.PrivateKey {
        let raw = edPrivate.rawData  // 32 bytes seed
        // SHA-512 of the seed
        let hash = SHA512.hash(data: raw)
        var clamped = Data(hash.prefix(32))
        // Clamp per RFC 7748
        clamped[0] &= 248
        clamped[31] &= 127
        clamped[31] |= 64
        return try Curve25519.KeyAgreement.PrivateKey(rawRepresentation: clamped)
    }

    /// Convert Ed25519 public key to X25519 public key.
    /// Uses Montgomery birational map: u = (1+y)/(1-y) mod p
    /// Matches Python: ed25519_public_to_x25519(ed_public)
    static func fromEd25519Public(_ edPublic: Curve25519.Signing.PublicKey) throws -> Curve25519.KeyAgreement.PublicKey {
        let x25519Bytes = FieldArithmetic.ed25519PublicToX25519(edPublic.rawData)
        return try Curve25519.KeyAgreement.PublicKey(rawRepresentation: x25519Bytes)
    }
}