#include <iostream>
#include <iomanip>
#include <cstring>
#include <vector>
#include <sstream>

#include "crypto/sha.h"

using namespace netplus;

// Helper to convert bytes to hex string
static std::string to_hex(const std::vector<uint8_t>& v) {
    std::ostringstream oss;
    for (auto b : v) {
        oss << std::hex << std::setfill('0') << std::setw(2) << int(b);
    }
    return oss.str();
}

// Helper to print test result
static void print_result(const std::string& name, const std::string& got, const std::string& expected) {
    bool pass = (got == expected);
    std::cout << (pass ? "✅ PASS" : "❌ FAIL") << " - " << name << std::endl;
    if (!pass) {
        std::cout << "  Got:      " << got << std::endl;
        std::cout << "  Expected: " << expected << std::endl;
    }
}

int main() {
    int passed = 0, failed = 0;

    std::cout << "========== SHA-1 NIST FIPS 180-4 Test Vectors ==========" << std::endl;
    
    // SHA-1 Test 1: Empty string
    {
        std::vector<uint8_t> input;
        auto hash = sha1_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "da39a3ee5e6b4b0d3255bfef95601890afd80709";
        if (got == expected) passed++; else failed++;
        print_result("SHA-1: Empty string", got, expected);
    }

    // SHA-1 Test 2: "abc"
    {
        std::vector<uint8_t> input = {'a', 'b', 'c'};
        auto hash = sha1_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "a9993e364706816aba3e25717850c26c9cd0d89d";
        if (got == expected) passed++; else failed++;
        print_result("SHA-1: 'abc'", got, expected);
    }

    // SHA-1 Test 3: Long string
    {
        std::string input_str = "abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq";
        std::vector<uint8_t> input(input_str.begin(), input_str.end());
        auto hash = sha1_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "84983e441c3bd26ebaae4aa1f95129e5e54670f1";
        if (got == expected) passed++; else failed++;
        print_result("SHA-1: Long string", got, expected);
    }

    // SHA-1 Test 4: One million 'a's
    {
        std::vector<uint8_t> input(1000000, 'a');
        auto hash = sha1_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "34aa973cd4c4daa4f61eeb2bdbad27316534016f";
        if (got == expected) passed++; else failed++;
        print_result("SHA-1: One million 'a's", got, expected);
    }

    std::cout << std::endl << "========== SHA-256 NIST FIPS 180-4 Test Vectors ==========" << std::endl;
    
    // SHA-256 Test 1: Empty string
    {
        std::vector<uint8_t> input;
        auto hash = sha256_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
        if (got == expected) passed++; else failed++;
        print_result("SHA-256: Empty string", got, expected);
    }

    // SHA-256 Test 2: "abc"
    {
        std::vector<uint8_t> input = {'a', 'b', 'c'};
        auto hash = sha256_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad";
        if (got == expected) passed++; else failed++;
        print_result("SHA-256: 'abc'", got, expected);
    }

    // SHA-256 Test 3: Long string
    {
        std::string input_str = "abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq";
        std::vector<uint8_t> input(input_str.begin(), input_str.end());
        auto hash = sha256_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1";
        if (got == expected) passed++; else failed++;
        print_result("SHA-256: Long string", got, expected);
    }

    // SHA-256 Test 4: One million 'a's
    {
        std::vector<uint8_t> input(1000000, 'a');
        auto hash = sha256_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0";
        if (got == expected) passed++; else failed++;
        print_result("SHA-256: One million 'a's", got, expected);
    }

    std::cout << std::endl << "========== SHA-384 NIST FIPS 180-4 Test Vectors ==========" << std::endl;
    
    // SHA-384 Test 1: Empty string
    {
        std::vector<uint8_t> input;
        auto hash = sha384_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "38b060a751ac96384cd9327eb1b1e36a21fdb71114be07434c0cc7bf63f6e1da274edebfe76f65fbd51ad2f14898b95b";
        if (got == expected) passed++; else failed++;
        print_result("SHA-384: Empty string", got, expected);
    }

    // SHA-384 Test 2: "abc"
    {
        std::vector<uint8_t> input = {'a', 'b', 'c'};
        auto hash = sha384_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "cb00753f45a35e8bb5a03d699ac65007272c32ab0eded1631a8b605a43ff5bed8086072ba1e7cc2358baeca134c825a7";
        if (got == expected) passed++; else failed++;
        print_result("SHA-384: 'abc'", got, expected);
    }

    // SHA-384 Test 3: Long string
    {
        std::string input_str = "abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu";
        std::vector<uint8_t> input(input_str.begin(), input_str.end());
        auto hash = sha384_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "09330c33f71147e83d192fc782cd1b4753111b173b3b05d22fa08086e3b0f712fcc7c71a557e2db966c3e9fa91746039";
        if (got == expected) passed++; else failed++;
        print_result("SHA-384: Long string", got, expected);
    }

    // SHA-384 Test 4: One million 'a's
    {
        std::vector<uint8_t> input(1000000, 'a');
        auto hash = sha384_hash(input);
        std::string got = to_hex(hash);
        std::string expected = "9d0e1809716474cb086e834e310a4a1ced149e9c00f248527972cec5704c2a5b07b8b3dc38ecc4ebae97ddd87f3d8985";
        if (got == expected) passed++; else failed++;
        print_result("SHA-384: One million 'a's", got, expected);
    }

    std::cout << std::endl << "========== Summary ==========" << std::endl;
    std::cout << "Passed: " << passed << std::endl;
    std::cout << "Failed: " << failed << std::endl;
    std::cout << "Total:  " << (passed + failed) << std::endl;

    return (failed == 0) ? 0 : 1;
}
