import Foundation

// ═══════════════════════════════════════════════════════════════
// Montana Protocol — Layer 1 Cryptography: 4 Checks + 16 Tests
// ML-DSA-65 (FIPS 204) Post-Quantum Signatures
//
// Check 1: Generation (seed → deterministic keys)
// Check 2: Addressing (pubkey → mt... with checksum)
// Check 3: Sign & Verify (deterministic signatures)
// Check 4: Isolation (private key never exposed)
//
// Servers: Amsterdam (72.56.102.240), Almaty (91.200.148.93)
// ═══════════════════════════════════════════════════════════════

@main
struct CryptoTest {
    static var passed = 0
    static var failed = 0
    static var total = 16

    static func test(_ num: Int, _ name: String, _ check: () -> Bool) {
        let result = check()
        if result {
            print("  [\(num)/\(total)] PASS: \(name)")
            passed += 1
        } else {
            print("  [\(num)/\(total)] FAIL: \(name)")
            failed += 1
        }
    }

    // HTTP helper for server tests
    static func httpRequest(url: String, method: String = "GET", body: [String: Any]? = nil, headers: [String: String] = [:]) -> (data: Data?, statusCode: Int) {
        guard let u = URL(string: url) else { return (nil, 0) }
        var req = URLRequest(url: u, timeoutInterval: 10)
        req.httpMethod = method
        for (k, v) in headers {
            req.setValue(v, forHTTPHeaderField: k)
        }
        if let body = body {
            req.setValue("application/json", forHTTPHeaderField: "Content-Type")
            req.httpBody = try? JSONSerialization.data(withJSONObject: body)
        }

        let sem = DispatchSemaphore(value: 0)
        var respData: Data?
        var statusCode = 0

        let task = URLSession.shared.dataTask(with: req) { data, response, _ in
            respData = data
            statusCode = (response as? HTTPURLResponse)?.statusCode ?? 0
            sem.signal()
        }
        task.resume()
        _ = sem.wait(timeout: .now() + 15)
        return (respData, statusCode)
    }

    static func main() {
        print("═══════════════════════════════════════════════════════")
        print("  Montana Layer 1 — 4 Checks, 16 Mainnet Tests")
        print("  Algorithm: ML-DSA-65 (FIPS 204)")
        print("  Key sizes: priv=4032, pub=1952, sig=3309")
        print("═══════════════════════════════════════════════════════")

        // ═════════════════════════════════════════════════════════
        // CHECK 1: GENERATION — seed → deterministic keys
        // ═════════════════════════════════════════════════════════
        print("\n══ CHECK 1: GENERATION ══")

        // TEST 1: Generate 24-word mnemonic
        print("\n[TEST 1] Generate Mnemonic")
        guard let words = MontanaSeed.generateMnemonic() else {
            print("  FATAL: generateMnemonic() returned nil")
            Foundation.exit(1)
        }
        test(1, "Mnemonic: 24 words generated") {
            words.count == 24
        }

        // TEST 2: Validate mnemonic
        print("\n[TEST 2] Validate Mnemonic")
        test(2, "Mnemonic checksum valid") {
            MontanaSeed.validateMnemonic(words)
        }

        // TEST 3: Derive keypair from mnemonic
        print("\n[TEST 3] Derive Keypair from Mnemonic")
        guard let kp1 = MontanaSeed.keypairFromMnemonic(words) else {
            print("  FATAL: keypairFromMnemonic() returned nil")
            Foundation.exit(1)
        }
        test(3, "Keygen from seed: priv=4032, pub=1952") {
            kp1.privateKey.count == 4032 && kp1.publicKey.count == 1952
        }

        // TEST 4: DETERMINISTIC — same seed = same keys (CRITICAL)
        print("\n[TEST 4] Deterministic Keys — Same Seed = Same Keys")
        guard let kp2 = MontanaSeed.keypairFromMnemonic(words) else {
            print("  FATAL: second keypairFromMnemonic() returned nil")
            Foundation.exit(1)
        }
        test(4, "Same mnemonic → identical privateKey AND publicKey") {
            kp1.privateKey == kp2.privateKey && kp1.publicKey == kp2.publicKey
        }

        // ═════════════════════════════════════════════════════════
        // CHECK 2: ADDRESSING — pubkey → mt... with checksum
        // ═════════════════════════════════════════════════════════
        print("\n══ CHECK 2: ADDRESSING ══")

        // TEST 5: Address format
        print("\n[TEST 5] Address Format")
        let addr = MLDSA65.generateAddress(from: kp1.publicKey)
        test(5, "Address: mt prefix, 42 chars, lowercase hex") {
            addr.hasPrefix("mt") &&
            addr.count == 42 &&
            String(addr.dropFirst(2)).allSatisfy { "0123456789abcdef".contains($0) }
        }

        // TEST 6: Deterministic address
        print("\n[TEST 6] Deterministic Address")
        let addr2 = MLDSA65.generateAddress(from: kp1.publicKey)
        test(6, "Same pubkey → same address") {
            addr == addr2
        }

        // TEST 7: Address checksum validation
        print("\n[TEST 7] Address Checksum Validation")
        test(7, "Valid address passes checksum validation") {
            MLDSA65.validateAddress(addr)
        }

        // TEST 8: Tampered address fails validation
        print("\n[TEST 8] Tampered Address Rejection")
        var tamperedAddr = addr
        // Change one character in the middle
        let idx = tamperedAddr.index(tamperedAddr.startIndex, offsetBy: 10)
        let oldChar = tamperedAddr[idx]
        let newChar: Character = oldChar == "a" ? "b" : "a"
        tamperedAddr.replaceSubrange(idx...idx, with: String(newChar))
        test(8, "Changed 1 char → checksum validation fails") {
            !MLDSA65.validateAddress(tamperedAddr)
        }

        // TEST 9: Cannot reverse address → pubkey
        print("\n[TEST 9] Address Irreversibility")
        guard let kp3 = MLDSA65.generateKeypair() else {
            print("  FATAL: random keygen failed")
            Foundation.exit(1)
        }
        let addr3 = MLDSA65.generateAddress(from: kp3.publicKey)
        test(9, "Different keys → different addresses, both valid") {
            addr != addr3 && MLDSA65.validateAddress(addr3)
        }

        // ═════════════════════════════════════════════════════════
        // CHECK 3: SIGN & VERIFY — deterministic signatures
        // ═════════════════════════════════════════════════════════
        print("\n══ CHECK 3: SIGN & VERIFY ══")

        // TEST 10: Sign + Verify roundtrip
        print("\n[TEST 10] Sign + Verify Roundtrip")
        let genesisMsg = "Montana Genesis".data(using: .utf8)!
        guard let sig1 = MLDSA65.sign(message: genesisMsg, privateKey: kp1.privateKey) else {
            print("  FATAL: sign() returned nil")
            Foundation.exit(1)
        }
        test(10, "Sign(Montana Genesis) → Verify = true, sig=3309") {
            MLDSA65.verify(message: genesisMsg, signature: sig1, publicKey: kp1.publicKey) &&
            sig1.count == 3309
        }

        // TEST 11: DETERMINISTIC SIGNATURE — same key + same message = identical signature
        print("\n[TEST 11] Deterministic Signature — Same Key + Same Message = Identical")
        guard let sig2 = MLDSA65.sign(message: genesisMsg, privateKey: kp1.privateKey) else {
            print("  FATAL: second sign() returned nil")
            Foundation.exit(1)
        }
        test(11, "Two signatures of same message are byte-identical") {
            sig1 == sig2
        }

        // TEST 12: THE CANONICAL TEST
        // Seed → keys → sign → "delete" → recover from same seed → sign → identical
        print("\n[TEST 12] CANONICAL: Seed → Sign → Delete → Recover → Sign = Identical")
        let sig1Hex = sig1.map { String(format: "%02x", $0) }.joined()
        // "Delete" keys (forget them)
        // Recover from same mnemonic
        guard let kpRecovered = MontanaSeed.keypairFromMnemonic(words) else {
            print("  FATAL: recovery keypairFromMnemonic() returned nil")
            Foundation.exit(1)
        }
        guard let sigRecovered = MLDSA65.sign(message: genesisMsg, privateKey: kpRecovered.privateKey) else {
            print("  FATAL: recovered sign() returned nil")
            Foundation.exit(1)
        }
        let sigRecoveredHex = sigRecovered.map { String(format: "%02x", $0) }.joined()
        test(12, "Recovered signature is byte-identical to original") {
            sig1Hex == sigRecoveredHex &&
            kpRecovered.publicKey == kp1.publicKey &&
            kpRecovered.privateKey == kp1.privateKey
        }

        // TEST 13: Wrong key rejection
        print("\n[TEST 13] Wrong Key Rejection")
        test(13, "Verify with wrong pubkey = false") {
            !MLDSA65.verify(message: genesisMsg, signature: sig1, publicKey: kp3.publicKey)
        }

        // TEST 14: Tampered message rejection
        print("\n[TEST 14] Tampered Message Rejection")
        let tampered = "Montana Genesi5".data(using: .utf8)!
        test(14, "Verify tampered message = false") {
            !MLDSA65.verify(message: tampered, signature: sig1, publicKey: kp1.publicKey)
        }

        // ═════════════════════════════════════════════════════════
        // CHECK 4: ISOLATION + SERVER VERIFICATION
        // ═════════════════════════════════════════════════════════
        print("\n══ CHECK 4: ISOLATION + SERVERS ══")

        // TEST 15: Server Health — Amsterdam & Almaty
        print("\n[TEST 15] Server Health")
        let (_, amsStatus) = httpRequest(url: "http://72.56.102.240/api/health")
        let (_, almStatus) = httpRequest(url: "http://91.200.148.93/api/health")
        test(15, "Amsterdam=200, Almaty=200") {
            amsStatus == 200 && almStatus == 200
        }

        // TEST 16: Server PQ-Signed Balance — Amsterdam
        print("\n[TEST 16] Server PQ-Signed Balance — Amsterdam")
        let ts = Int(Date().timeIntervalSince1970)
        let balanceMsg = "BALANCE:\(addr):\(ts)"
        guard let balanceData = balanceMsg.data(using: .utf8),
              let balanceSig = MLDSA65.sign(message: balanceData, privateKey: kp1.privateKey) else {
            print("  FATAL: balance sign failed")
            Foundation.exit(1)
        }
        let pubHex = kp1.publicKey.map { String(format: "%02x", $0) }.joined()
        let sigHex = balanceSig.map { String(format: "%02x", $0) }.joined()

        let (balRespData, balRespStatus) = httpRequest(
            url: "http://72.56.102.240/api/balance/\(addr)",
            method: "GET",
            headers: [
                "X-Address": addr,
                "X-Timestamp": "\(ts)",
                "X-Signature": sigHex,
                "X-Public-Key": pubHex
            ]
        )

        var pqVerified = false
        if let data = balRespData,
           let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] {
            pqVerified = json["pq_verified"] as? Bool ?? false
            let bal = json["balance"] as? Double ?? -1
            print("    Response: balance=\(bal), pq_verified=\(pqVerified)")
        }

        test(16, "Amsterdam PQ balance: status=200, pq_verified=true") {
            balRespStatus == 200 && pqVerified
        }

        // ═══════════════════════════════════════════════════════
        // RESULTS
        // ═══════════════════════════════════════════════════════
        print("\n" + String(repeating: "═", count: 55))
        print("  Results: \(passed)/\(total) passed, \(failed) failed")
        print(String(repeating: "═", count: 55))

        if failed == 0 {
            print("\n  ML-DSA-65 LAYER 1: ALL \(total) TESTS PASSED")
            print("  Mnemonic: \(words.prefix(3).joined(separator: " "))... (\(words.count) words)")
            print("  Address:  \(addr)")
            print("  Checksum: \(MLDSA65.validateAddress(addr) ? "VALID" : "INVALID")")
            print("  Deterministic: sig1 == sigRecovered = \(sig1Hex == sigRecoveredHex)")
        } else {
            print("\n  FAILED — \(failed) test(s)")
            Foundation.exit(1)
        }
    }
}