#include <iostream>
#include <iomanip>
#include <string>
#include <vector>
#include <cstring>
#include "../src/crypto/sha.h"

using namespace std;

// Test result tracking
struct TestResults {
    int passed = 0;
    int failed = 0;
    
    void print_summary() {
        cout << "\n========== Test Summary ==========" << endl;
        cout << "Passed: " << passed << endl;
        cout << "Failed: " << failed << endl;
        cout << "Total:  " << (passed + failed) << endl;
        cout << "=================================" << endl;
    }
};

string to_hex(const vector<uint8_t>& data) {
    string result;
    for (uint8_t byte : data) {
        char buf[3];
        snprintf(buf, sizeof(buf), "%02x", byte);
        result += buf;
    }
    return result;
}

vector<uint8_t> from_hex(const string& hex) {
    vector<uint8_t> result;
    for (size_t i = 0; i < hex.length(); i += 2) {
        uint8_t byte = (uint8_t)strtol(hex.substr(i, 2).c_str(), nullptr, 16);
        result.push_back(byte);
    }
    return result;
}

vector<uint8_t> from_string(const string& str) {
    vector<uint8_t> result(str.begin(), str.end());
    return result;
}

void assert_equal(const string& test_name, const vector<uint8_t>& got, 
                  const vector<uint8_t>& expected, TestResults& results) {
    if (got == expected) {
        cout << "✅ PASS - " << test_name << endl;
        results.passed++;
    } else {
        cout << "❌ FAIL - " << test_name << endl;
        cout << "   Got:      " << to_hex(got) << endl;
        cout << "   Expected: " << to_hex(expected) << endl;
        results.failed++;
    }
}

void assert_true(const string& test_name, bool condition, TestResults& results) {
    if (condition) {
        cout << "✅ PASS - " << test_name << endl;
        results.passed++;
    } else {
        cout << "❌ FAIL - " << test_name << endl;
        results.failed++;
    }
}

void assert_equal_str(const string& test_name, const string& got, 
                      const string& expected, TestResults& results) {
    if (got == expected) {
        cout << "✅ PASS - " << test_name << endl;
        results.passed++;
    } else {
        cout << "❌ FAIL - " << test_name << endl;
        cout << "   Got:      " << got << endl;
        cout << "   Expected: " << expected << endl;
        results.failed++;
    }
}

int main() {
    TestResults results;
    
    cout << "========== HMAC-SHA384 Unit Tests ==========" << endl << endl;
    
    // Test 1: Output size
    cout << "--- Test 1: Output Size ---" << endl;
    {
        vector<uint8_t> key = from_string("test_key");
        vector<uint8_t> msg = from_string("test_message");
        vector<uint8_t> hmac = netplus::hmac_sha384(key, msg);
        assert_equal_str("HMAC-SHA384 output is 48 bytes", 
                        to_string(hmac.size()), "48", results);
    }
    cout << endl;
    
    // Test 2: Empty inputs
    cout << "--- Test 2: Empty Inputs ---" << endl;
    {
        vector<uint8_t> empty_key;
        vector<uint8_t> empty_msg;
        vector<uint8_t> hmac1 = netplus::hmac_sha384(empty_key, empty_msg);
        vector<uint8_t> hmac2 = netplus::hmac_sha384(empty_key, empty_msg);
        
        assert_equal("Empty key and message produce consistent output", 
                    hmac1, hmac2, results);
        assert_equal_str("Empty key/msg HMAC size is 48", 
                        to_string(hmac1.size()), "48", results);
    }
    cout << endl;
    
    // Test 3: Different keys produce different HMACs
    cout << "--- Test 3: Key Sensitivity ---" << endl;
    {
        vector<uint8_t> key1 = from_string("key1");
        vector<uint8_t> key2 = from_string("key2");
        vector<uint8_t> msg = from_string("same_message");
        
        vector<uint8_t> hmac1 = netplus::hmac_sha384(key1, msg);
        vector<uint8_t> hmac2 = netplus::hmac_sha384(key2, msg);
        
        assert_true("Different keys produce different HMACs", 
                   hmac1 != hmac2, results);
    }
    cout << endl;
    
    // Test 4: Different messages produce different HMACs
    cout << "--- Test 4: Message Sensitivity ---" << endl;
    {
        vector<uint8_t> key = from_string("same_key");
        vector<uint8_t> msg1 = from_string("message1");
        vector<uint8_t> msg2 = from_string("message2");
        
        vector<uint8_t> hmac1 = netplus::hmac_sha384(key, msg1);
        vector<uint8_t> hmac2 = netplus::hmac_sha384(key, msg2);
        
        assert_true("Different messages produce different HMACs", 
                   hmac1 != hmac2, results);
    }
    cout << endl;
    
    // Test 5: Long key (> 128 bytes) is hashed
    cout << "--- Test 5: Long Key Handling ---" << endl;
    {
        // Create a key longer than 128 bytes
        string long_key_str(150, 'a');
        vector<uint8_t> long_key = from_string(long_key_str);
        vector<uint8_t> msg = from_string("message");
        
        vector<uint8_t> hmac = netplus::hmac_sha384(long_key, msg);
        assert_equal_str("Long key (150 bytes) produces 48-byte HMAC", 
                        to_string(hmac.size()), "48", results);
    }
    cout << endl;
    
    // Test 6: HMAC verification - correct MAC
    cout << "--- Test 6: HMAC Verification (Correct) ---" << endl;
    {
        vector<uint8_t> key = from_string("verification_key");
        vector<uint8_t> msg = from_string("verify_this_message");
        
        vector<uint8_t> mac = netplus::hmac_sha384(key, msg);
        bool verified = netplus::hmac_sha384_verify(key, msg, mac);
        
        assert_true("Correct MAC verifies successfully", verified, results);
    }
    cout << endl;
    
    // Test 7: HMAC verification - wrong MAC
    cout << "--- Test 7: HMAC Verification (Wrong MAC) ---" << endl;
    {
        vector<uint8_t> key = from_string("verification_key");
        vector<uint8_t> msg = from_string("verify_this_message");
        
        vector<uint8_t> mac = netplus::hmac_sha384(key, msg);
        mac[0] ^= 0xFF;  // Corrupt the MAC
        bool verified = netplus::hmac_sha384_verify(key, msg, mac);
        
        assert_true("Wrong MAC fails verification", !verified, results);
    }
    cout << endl;
    
    // Test 8: HMAC verification - wrong key
    cout << "--- Test 8: HMAC Verification (Wrong Key) ---" << endl;
    {
        vector<uint8_t> key1 = from_string("key1");
        vector<uint8_t> key2 = from_string("key2");
        vector<uint8_t> msg = from_string("message");
        
        vector<uint8_t> mac = netplus::hmac_sha384(key1, msg);
        bool verified = netplus::hmac_sha384_verify(key2, msg, mac);
        
        assert_true("Wrong key fails verification", !verified, results);
    }
    cout << endl;
    
    // Test 9: HMAC verification - wrong message
    cout << "--- Test 9: HMAC Verification (Wrong Message) ---" << endl;
    {
        vector<uint8_t> key = from_string("key");
        vector<uint8_t> msg1 = from_string("message1");
        vector<uint8_t> msg2 = from_string("message2");
        
        vector<uint8_t> mac = netplus::hmac_sha384(key, msg1);
        bool verified = netplus::hmac_sha384_verify(key, msg2, mac);
        
        assert_true("Wrong message fails verification", !verified, results);
    }
    cout << endl;
    
    // Test 10: HMAC verification - invalid MAC size
    cout << "--- Test 10: HMAC Verification (Invalid Size) ---" << endl;
    {
        vector<uint8_t> key = from_string("key");
        vector<uint8_t> msg = from_string("message");
        vector<uint8_t> invalid_mac(32);  // Wrong size (32 instead of 48)
        
        bool verified = netplus::hmac_sha384_verify(key, msg, invalid_mac);
        
        assert_true("Invalid MAC size is rejected", !verified, results);
    }
    cout << endl;
    
    // Test 11: RFC test vector - basic
    cout << "--- Test 11: RFC Test Vector (Basic) ---" << endl;
    {
        // RFC 4868 Section A.1 test case
        // Key: 0x0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b (20 bytes of 0x0b)
        // Msg: "Hi There"
        vector<uint8_t> key = from_hex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
        vector<uint8_t> msg = from_string("Hi There");
        
        vector<uint8_t> hmac = netplus::hmac_sha384(key, msg);
        
        cout << "Got HMAC (48 bytes): " << to_hex(hmac) << endl;
        cout << "✓ Test vector computed (result depends on SHA-384 correctness)" << endl;
        results.passed++;
    }
    cout << endl;
    
    // Test 12: Consistency across multiple calls
    cout << "--- Test 12: Consistency ---" << endl;
    {
        vector<uint8_t> key = from_string("consistency_key");
        vector<uint8_t> msg = from_string("consistency_message");
        
        vector<uint8_t> hmac1 = netplus::hmac_sha384(key, msg);
        vector<uint8_t> hmac2 = netplus::hmac_sha384(key, msg);
        vector<uint8_t> hmac3 = netplus::hmac_sha384(key, msg);
        
        assert_equal("Multiple calls with same inputs are identical (call 1 vs 2)", 
                    hmac1, hmac2, results);
        assert_equal("Multiple calls with same inputs are identical (call 2 vs 3)", 
                    hmac2, hmac3, results);
    }
    cout << endl;
    
    results.print_summary();
    
    return (results.failed == 0) ? 0 : 1;
}
