//
//  MontanaSeed.swift
//  Montana Protocol — macOS
//
//  BIP-39 Mnemonic (24 words) → Deterministic ML-DSA-65 Keypair
//  Standard: BIP-39 + PBKDF2-SHA512 + SHA256-CTR DRBG
//
//  Flow: entropy(256 bit) → 24 words → PBKDF2 → 64-byte seed → ML-DSA-65 keypair
//

import Foundation

public struct MontanaSeed {

    // MARK: - Constants

    /// Entropy size for 24-word mnemonic: 256 bits = 32 bytes
    static let entropySize = 32

    /// PBKDF2 iterations (BIP-39 standard)
    static let pbkdf2Iterations: UInt32 = 2048

    /// Derived seed size: 512 bits = 64 bytes
    static let seedSize = 64

    // MARK: - Generate Mnemonic

    /// Generate a new 24-word BIP-39 mnemonic from secure random entropy
    /// - Returns: Array of 24 words
    public static func generateMnemonic() -> [String]? {
        // 1. Generate 256 bits of secure random entropy
        var entropy = [UInt8](repeating: 0, count: entropySize)
        let status = SecRandomCopyBytes(kSecRandomDefault, entropySize, &entropy)
        guard status == errSecSuccess else { return nil }

        return mnemonicFromEntropy(entropy)
    }

    /// Convert entropy bytes to mnemonic words
    /// - Parameter entropy: 32 bytes of entropy
    /// - Returns: 24 mnemonic words
    static func mnemonicFromEntropy(_ entropy: [UInt8]) -> [String]? {
        guard entropy.count == entropySize else { return nil }

        // 2. Compute checksum: SHA-256(entropy), take first 8 bits
        var hash = [UInt8](repeating: 0, count: 32)
        entropy.withUnsafeBufferPointer { ptr in
            CC_SHA256(ptr.baseAddress, CC_LONG(entropy.count), &hash)
        }
        let checksumByte = hash[0] // first 8 bits for 256-bit entropy

        // 3. Combine entropy + checksum = 264 bits
        var bits = [Bool]()
        bits.reserveCapacity(264)
        for byte in entropy {
            for j in (0..<8).reversed() {
                bits.append((byte >> j) & 1 == 1)
            }
        }
        for j in (0..<8).reversed() {
            bits.append((checksumByte >> j) & 1 == 1)
        }

        // 4. Split into 24 groups of 11 bits → word indices
        let wordlist = BIP39Wordlist.english
        guard wordlist.count == 2048 else { return nil }

        var words = [String]()
        words.reserveCapacity(24)
        for i in 0..<24 {
            var index = 0
            for j in 0..<11 {
                if bits[i * 11 + j] {
                    index |= (1 << (10 - j))
                }
            }
            guard index < 2048 else { return nil }
            words.append(wordlist[index])
        }

        return words
    }

    // MARK: - Validate Mnemonic

    /// Validate a BIP-39 mnemonic (24 words, checksum verification)
    /// - Parameter words: Array of mnemonic words
    /// - Returns: true if valid
    public static func validateMnemonic(_ words: [String]) -> Bool {
        guard words.count == 24 else { return false }

        let wordlist = BIP39Wordlist.english
        guard wordlist.count == 2048 else { return false }

        // Convert words to bit array
        var bits = [Bool]()
        bits.reserveCapacity(264)
        for word in words {
            guard let index = wordlist.firstIndex(of: word.lowercased()) else {
                return false
            }
            for j in (0..<11).reversed() {
                bits.append((index >> j) & 1 == 1)
            }
        }

        guard bits.count == 264 else { return false }

        // Extract entropy (first 256 bits) and checksum (last 8 bits)
        var entropy = [UInt8](repeating: 0, count: 32)
        for i in 0..<32 {
            var byte: UInt8 = 0
            for j in 0..<8 {
                if bits[i * 8 + j] {
                    byte |= UInt8(1 << (7 - j))
                }
            }
            entropy[i] = byte
        }

        var checksumByte: UInt8 = 0
        for j in 0..<8 {
            if bits[256 + j] {
                checksumByte |= UInt8(1 << (7 - j))
            }
        }

        // Verify checksum: SHA-256(entropy)[0] == checksumByte
        var hash = [UInt8](repeating: 0, count: 32)
        entropy.withUnsafeBufferPointer { ptr in
            CC_SHA256(ptr.baseAddress, CC_LONG(entropy.count), &hash)
        }

        return hash[0] == checksumByte
    }

    // MARK: - Mnemonic → Seed (PBKDF2-SHA512)

    /// Derive 64-byte seed from mnemonic using PBKDF2-SHA512 (BIP-39 standard)
    /// - Parameter words: 24 mnemonic words
    /// - Returns: 64-byte seed
    public static func mnemonicToSeed(_ words: [String]) -> Data? {
        guard validateMnemonic(words) else { return nil }

        let mnemonic = words.joined(separator: " ")
        let salt = "mnemonic" // BIP-39 standard salt (no passphrase)

        guard let mnemonicData = mnemonic.data(using: .utf8),
              let saltData = salt.data(using: .utf8) else { return nil }

        var derivedKey = [UInt8](repeating: 0, count: seedSize)

        let status = mnemonicData.withUnsafeBytes { mnemonicPtr in
            saltData.withUnsafeBytes { saltPtr in
                CCKeyDerivationPBKDF(
                    CCPBKDFAlgorithm(kCCPBKDF2),
                    mnemonicPtr.baseAddress?.assumingMemoryBound(to: Int8.self),
                    mnemonicData.count,
                    saltPtr.baseAddress?.assumingMemoryBound(to: UInt8.self),
                    saltData.count,
                    CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA512),
                    pbkdf2Iterations,
                    &derivedKey,
                    seedSize
                )
            }
        }

        guard status == kCCSuccess else { return nil }
        return Data(derivedKey)
    }

    // MARK: - Seed → ML-DSA-65 Keypair (Deterministic)

    /// Derive ML-DSA-65 keypair deterministically from 64-byte seed
    /// Uses SHA256-CTR as deterministic PRNG injected into liboqs
    /// - Parameter seed: 64-byte seed from PBKDF2
    /// - Returns: (privateKey, publicKey) tuple
    public static func deriveKeypair(from seed: Data) -> (privateKey: Data, publicKey: Data)? {
        guard seed.count == seedSize else { return nil }
        return MLDSA65.generateKeypairFromSeed(seed: seed)
    }

    // MARK: - Full Flow: Mnemonic → Keypair

    /// Generate keypair from mnemonic words
    /// - Parameter words: 24 mnemonic words
    /// - Returns: (privateKey, publicKey) tuple
    public static func keypairFromMnemonic(_ words: [String]) -> (privateKey: Data, publicKey: Data)? {
        guard let seed = mnemonicToSeed(words) else { return nil }
        defer {
            var mutableSeed = seed
            MLDSA65.zeroMemory(&mutableSeed)
        }
        return deriveKeypair(from: seed)
    }
}