//
//  MLDSA65.swift
//  Montana Protocol — macOS
//
//  ML-DSA-65 (FIPS 204) Post-Quantum Signatures via liboqs
//  Layer 1: Keys, Addresses, Signatures
//
//  Key sizes:
//    Private key: 4032 bytes
//    Public key:  1952 bytes
//    Signature:   3309 bytes
//
//  Address format: mt + SHA256(pubkey)[:18].hex() + checksum(4 hex) = 42 chars
//  Checksum: SHA256(SHA256("mt" + payload))[:2].hex()
//

import Foundation

/// ML-DSA-65 Post-Quantum Signature Scheme
/// Self-Custody implementation using liboqs
public struct MLDSA65 {

    // MARK: - Constants (FIPS 204)

    /// Private key size: 4032 bytes
    public static let privateKeySize = 4032

    /// Public key size: 1952 bytes
    public static let publicKeySize = 1952

    /// Signature size: 3309 bytes
    public static let signatureSize = 3309

    /// Address length: 42 characters (mt + 36 hex payload + 4 hex checksum)
    public static let addressLength = 42

    /// Serial queue for thread-safe liboqs operations
    private static let cryptoQueue = DispatchQueue(label: "network.montana.crypto", qos: .userInitiated)

    /// Thread-safe OQS initialization (dispatch_once semantics via Swift static let)
    private static let oqsInit: Void = {
        OQS_init()
    }()

    /// Global deterministic RNG state (protected by cryptoQueue)
    private static var deterministicRNG: DeterministicCTRDRBG?

    // MARK: - Deterministic RNG (SHA256-CTR DRBG)

    /// Counter-mode DRBG using SHA-256 for deterministic key generation
    private class DeterministicCTRDRBG {
        private var seed: [UInt8]
        private var counter: UInt64 = 0

        init(seed: Data) {
            self.seed = [UInt8](seed)
        }

        func generate(buf: UnsafeMutablePointer<UInt8>, count: Int) {
            var offset = 0
            while offset < count {
                // block = SHA256(seed || counter_be)
                var input = seed
                withUnsafeBytes(of: counter.bigEndian) { input.append(contentsOf: $0) }
                var hash = [UInt8](repeating: 0, count: 32)
                input.withUnsafeBufferPointer { ptr in
                    CC_SHA256(ptr.baseAddress, CC_LONG(input.count), &hash)
                }
                let needed = min(32, count - offset)
                for i in 0..<needed {
                    buf[offset + i] = hash[i]
                }
                offset += needed
                counter += 1
            }
        }

        func reset() {
            // Zero the seed
            for i in 0..<seed.count { seed[i] = 0 }
            counter = 0
        }
    }

    /// C callback for liboqs custom RNG
    private static let deterministicRNGCallback: @convention(c) (UnsafeMutablePointer<UInt8>?, Int) -> Void = { buf, len in
        guard let buf = buf else { return }
        MLDSA65.deterministicRNG?.generate(buf: buf, count: len)
    }

    /// C callback that returns all zeros (for deterministic signing, rnd=0)
    private static let zeroRNGCallback: @convention(c) (UnsafeMutablePointer<UInt8>?, Int) -> Void = { buf, len in
        guard let buf = buf else { return }
        memset(buf, 0, len)
    }

    // MARK: - Key Generation (Random)

    /// Generate ML-DSA-65 keypair using system secure RNG
    /// - Returns: (privateKey, publicKey) tuple
    public static func generateKeypair() -> (privateKey: Data, publicKey: Data)? {
        return cryptoQueue.sync {
            _ = oqsInit

            guard let sig = OQS_SIG_new(OQS_SIG_alg_ml_dsa_65) else {
                return nil
            }
            defer { OQS_SIG_free(sig) }

            var publicKey = Data(count: Int(sig.pointee.length_public_key))
            var privateKey = Data(count: Int(sig.pointee.length_secret_key))

            let result = publicKey.withUnsafeMutableBytes { pubPtr in
                privateKey.withUnsafeMutableBytes { privPtr in
                    OQS_SIG_keypair(sig,
                                   pubPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                                   privPtr.baseAddress?.assumingMemoryBound(to: UInt8.self))
                }
            }

            guard result == OQS_SUCCESS else {
                return nil
            }

            return (privateKey, publicKey)
        }
    }

    // MARK: - Key Generation (Deterministic from Seed)

    /// Generate ML-DSA-65 keypair deterministically from a 64-byte seed
    /// Injects SHA256-CTR DRBG into liboqs, generates keypair, restores system RNG
    /// - Parameter seed: 64-byte seed (from PBKDF2)
    /// - Returns: (privateKey, publicKey) tuple — always identical for same seed
    public static func generateKeypairFromSeed(seed: Data) -> (privateKey: Data, publicKey: Data)? {
        guard seed.count == 64 else { return nil }

        return cryptoQueue.sync {
            _ = oqsInit

            // Inject deterministic RNG
            deterministicRNG = DeterministicCTRDRBG(seed: seed)
            OQS_randombytes_custom_algorithm(deterministicRNGCallback)

            defer {
                // Restore system RNG and zero state
                OQS_randombytes_switch_algorithm(OQS_RAND_alg_system)
                deterministicRNG?.reset()
                deterministicRNG = nil
            }

            guard let sig = OQS_SIG_new(OQS_SIG_alg_ml_dsa_65) else {
                return nil
            }
            defer { OQS_SIG_free(sig) }

            var publicKey = Data(count: Int(sig.pointee.length_public_key))
            var privateKey = Data(count: Int(sig.pointee.length_secret_key))

            let result = publicKey.withUnsafeMutableBytes { pubPtr in
                privateKey.withUnsafeMutableBytes { privPtr in
                    OQS_SIG_keypair(sig,
                                   pubPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                                   privPtr.baseAddress?.assumingMemoryBound(to: UInt8.self))
                }
            }

            guard result == OQS_SUCCESS else {
                return nil
            }

            return (privateKey, publicKey)
        }
    }

    // MARK: - Signing (Deterministic — FIPS 204 rnd=0)

    /// Sign message deterministically with ML-DSA-65 private key
    /// Same message + same key = identical signature (FIPS 204 deterministic mode)
    /// - Parameters:
    ///   - message: Message to sign (must not be empty)
    ///   - privateKey: ML-DSA-65 private key (4032 bytes)
    /// - Returns: Signature (3309 bytes)
    public static func sign(message: Data, privateKey: Data) -> Data? {
        guard !message.isEmpty else {
            return nil
        }
        guard privateKey.count == privateKeySize else {
            return nil
        }

        return cryptoQueue.sync {
            _ = oqsInit

            // Inject zero RNG → deterministic signing (rnd = 0x00...00)
            OQS_randombytes_custom_algorithm(zeroRNGCallback)

            defer {
                // Restore system RNG
                OQS_randombytes_switch_algorithm(OQS_RAND_alg_system)
            }

            guard let sig = OQS_SIG_new(OQS_SIG_alg_ml_dsa_65) else {
                return nil
            }
            defer { OQS_SIG_free(sig) }

            var signature = Data(count: Int(sig.pointee.length_signature))
            var signatureLen: Int = 0

            let result = signature.withUnsafeMutableBytes { sigPtr in
                message.withUnsafeBytes { msgPtr in
                    privateKey.withUnsafeBytes { keyPtr in
                        OQS_SIG_sign(sig,
                                    sigPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                                    &signatureLen,
                                    msgPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                                    message.count,
                                    keyPtr.baseAddress?.assumingMemoryBound(to: UInt8.self))
                    }
                }
            }

            guard result == OQS_SUCCESS else {
                return nil
            }

            signature = signature.prefix(signatureLen)
            return signature
        }
    }

    // MARK: - Verification

    /// Verify ML-DSA-65 signature
    /// - Parameters:
    ///   - message: Original message (must not be empty)
    ///   - signature: Signature to verify (3309 bytes)
    ///   - publicKey: ML-DSA-65 public key (1952 bytes)
    /// - Returns: true if signature is valid
    public static func verify(message: Data, signature: Data, publicKey: Data) -> Bool {
        guard !message.isEmpty else {
            return false
        }
        guard signature.count == signatureSize else {
            return false
        }
        guard publicKey.count == publicKeySize else {
            return false
        }

        return cryptoQueue.sync {
            _ = oqsInit

            guard let sig = OQS_SIG_new(OQS_SIG_alg_ml_dsa_65) else {
                return false
            }
            defer { OQS_SIG_free(sig) }

            let result = message.withUnsafeBytes { msgPtr in
                signature.withUnsafeBytes { sigPtr in
                    publicKey.withUnsafeBytes { keyPtr in
                        OQS_SIG_verify(sig,
                                      msgPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                                      message.count,
                                      sigPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                                      signature.count,
                                      keyPtr.baseAddress?.assumingMemoryBound(to: UInt8.self))
                    }
                }
            }

            return result == OQS_SUCCESS
        }
    }

    // MARK: - Address Generation (with Checksum)

    /// Generate Montana address from public key
    /// Format: mt + SHA256(pubkey)[:18].hex() + checksum = 42 chars
    /// Checksum: SHA256(SHA256("mt" + payload))[:2].hex()
    public static func generateAddress(from publicKey: Data) -> String {
        let hash = sha256(publicKey)
        let payload = hash.prefix(18).map { String(format: "%02x", $0) }.joined()
        let checksum = computeChecksum("mt" + payload)
        return "mt" + payload + checksum
    }

    /// Validate a Montana address (format + checksum)
    /// - Parameter address: Address string to validate
    /// - Returns: true if format and checksum are valid
    public static func validateAddress(_ address: String) -> Bool {
        guard address.count == addressLength else { return false }
        guard address.hasPrefix("mt") else { return false }

        let body = String(address.dropFirst(2))
        guard body.allSatisfy({ "0123456789abcdef".contains($0) }) else { return false }

        let payload = String(body.prefix(36))
        let checksum = String(body.suffix(4))
        let expected = computeChecksum("mt" + payload)

        return checksum == expected
    }

    /// Compute 4-char hex checksum: SHA256(SHA256(prefix + payload))[:2]
    private static func computeChecksum(_ input: String) -> String {
        guard let data = input.data(using: .utf8) else { return "0000" }
        let hash1 = sha256(data)
        let hash2 = sha256(hash1)
        return hash2.prefix(2).map { String(format: "%02x", $0) }.joined()
    }

    // MARK: - Memory Safety

    /// Zero out sensitive data in place
    public static func zeroMemory(_ data: inout Data) {
        data.withUnsafeMutableBytes { ptr in
            if let baseAddress = ptr.baseAddress {
                memset(baseAddress, 0, ptr.count)
            }
        }
        data = Data()
    }

    // MARK: - Helpers

    /// SHA-256 hash (internal, used for address + checksum)
    static func sha256(_ data: Data) -> Data {
        var hash = [UInt8](repeating: 0, count: 32)
        data.withUnsafeBytes {
            _ = CC_SHA256($0.baseAddress, CC_LONG(data.count), &hash)
        }
        return Data(hash)
    }
}