//     Richard CPU
//     Copyright (C) 2026  Richard Qin
//     This program is free software: you can redistribute it and/or modify
//     it under the terms of the GNU General Public License as published by
//     the Free Software Foundation, either version 3 of the License, or
//     (at your option) any later version.

//     This program is distributed in the hope that it will be useful,
//     but WITHOUT ANY WARRANTY; without even the implied warranty of
//     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//     GNU General Public License for more details.

//     You should have received a copy of the GNU General Public License
//     along with this program.  If not, see <https://www.gnu.org/licenses/>.

// File path: tb/unit/tb_alu.cc
// Description: ALU unit test
#include "Valu.h"
#include <cassert>
#include <cstdint>
#include <iostream>
#include <verilated.h>

enum AluOp {
    ALU_ADD    = 0,
    ALU_SUB    = 8,
    ALU_SLL    = 1,
    ALU_SLT    = 2,
    ALU_SLTU   = 3,
    ALU_XOR    = 4,
    ALU_SRL    = 5,
    ALU_SRA    = 13,
    ALU_OR     = 6,
    ALU_AND    = 7,
    ALU_MUL    = 16,
    ALU_MULH   = 17,
    ALU_MULHSU = 18,
    ALU_MULHU  = 19,
    ALU_DIV    = 20,
    ALU_DIVU   = 21,
    ALU_REM    = 22,
    ALU_REMU   = 23
};

void test_alu(Valu* alu, uint8_t op, uint64_t rs1, uint64_t rs2, uint64_t expected, const char* test_name) {
    alu->alu_op = op;
    alu->operand_a = rs1;
    alu->operand_b = rs2;

    alu->eval();
    
    if (alu->alu_result == expected) {
        std::cout << "\033[32m";
        std::cout << "[PASS] " << test_name << " | " 
                  << rs1 << " op " << rs2 << " = " << expected << "\033[0m" << std::endl;
    } else {
        std::cout << "\033[31m";
        std::cerr << "[FAIL] " << test_name << " | " 
                  << "Expected: " << expected << ", Got: " << alu->alu_result << "\033[0m" << std::endl;
        assert(false);
    }
}

int main(int argc, char** argv) {
    Verilated::commandArgs(argc, argv);
    Valu* alu = new Valu;
    
    std::cout << "----------------------------------" << std::endl;
    std::cout << "        ALU Unit Test" << std::endl;
    std::cout << "----------------------------------" << std::endl;
    
    test_alu(alu, ALU_ADD, 100, 50, 150, "ADD_Positive");
    test_alu(alu, ALU_SUB, 100, 50, 50,  "SUB_Positive");

    test_alu(alu, ALU_ADD, (uint64_t)-5, 10, 5, "ADD_Negative");

    test_alu(alu, ALU_AND, 0xFF00, 0x0F0F, 0x0F00, "AND_Mask");
    test_alu(alu, ALU_XOR, 0xAAAA, 0xFFFF, 0x5555, "XOR_Invert");

    test_alu(alu, ALU_SLL, 1, 10, 1024, "SLL_By_10");

    uint64_t minus_one = (uint64_t)-1;
    uint64_t five = 5;

    test_alu(alu, ALU_SLTU, minus_one, five, 0, "SLTU (Unsigned: -1 < 5 is False)");
    test_alu(alu, ALU_SLT,  minus_one, five, 1, "SLT  (Signed: -1 < 5 is True)");


    // Missing Basic Operations
    test_alu(alu, ALU_OR,  0xAAAA, 0x5555, 0xFFFF, "OR_All_Bits");
    test_alu(alu, ALU_SRL, 0xF000000000000000ULL, 4, 0x0F00000000000000ULL, "SRL_By_4");
    
    // Arithmetic Right Shift (SRA)
    test_alu(alu, ALU_SRA, 0xF000000000000000ULL, 4, 0xFF00000000000000ULL, "SRA_Negative_By_4");
    test_alu(alu, ALU_SRA, 0x7000000000000000ULL, 4, 0x0700000000000000ULL, "SRA_Positive_By_4");

    // Shift Edge Cases
    test_alu(alu, ALU_SLL, 1,  0,  1, "SLL_By_0");
    test_alu(alu, ALU_SLL, 1, 63,  0x8000000000000000ULL, "SLL_By_63");
    test_alu(alu, ALU_SRL, 0x8000000000000000ULL, 63, 1, "SRL_By_63");
    test_alu(alu, ALU_SRA, 0x8000000000000000ULL, 63, 0xFFFFFFFFFFFFFFFFULL, "SRA_By_63_SignExt");

    // All 0 / All 1 Cases
    test_alu(alu, ALU_ADD, 0, 0, 0, "ADD_Zero");
    test_alu(alu, ALU_ADD, 0xFFFFFFFFFFFFFFFFULL, 1, 0, "ADD_Overflow");
    test_alu(alu, ALU_XOR, 0xFFFFFFFFFFFFFFFFULL, 0xFFFFFFFFFFFFFFFFULL, 0, "XOR_Same");

    // SLT / SLTU Same Value
    test_alu(alu, ALU_SLT,  100, 100, 0, "SLT_Same_Value");
    test_alu(alu, ALU_SLTU, 100, 100, 0, "SLTU_Same_Value");
    
    // SLT Edge Cases
    test_alu(alu, ALU_SLT, 0x8000000000000000ULL, 0x7FFFFFFFFFFFFFFFULL, 1, "SLT_Min_max");
    test_alu(alu, ALU_SLT, 0x7FFFFFFFFFFFFFFFULL, 0x8000000000000000ULL, 0, "SLT_Max_Min");

    // M extension
    test_alu(alu, ALU_MUL, 7, 6, 42, "MUL_Positive");
    test_alu(alu, ALU_MUL, 0xFFFFFFFFFFFFFFFFULL, 2, 0xFFFFFFFFFFFFFFFEULL, "MUL_NegativeLow");
    test_alu(alu, ALU_MULH, 0xFFFFFFFFFFFFFFFFULL, 2, 0xFFFFFFFFFFFFFFFFULL, "MULH_Negative");
    test_alu(alu, ALU_MULHU, 0xFFFFFFFFFFFFFFFFULL, 2, 1, "MULHU_UnsignedHigh");
    test_alu(alu, ALU_DIV, (uint64_t)-21, 4, (uint64_t)-5, "DIV_SignedTruncTowardZero");
    test_alu(alu, ALU_DIVU, 21, 4, 5, "DIVU");
    test_alu(alu, ALU_REM, (uint64_t)-21, 4, (uint64_t)-1, "REM_Signed");
    test_alu(alu, ALU_REMU, 21, 4, 1, "REMU");
    test_alu(alu, ALU_DIV, 123, 0, 0xFFFFFFFFFFFFFFFFULL, "DIV_ByZero");
    test_alu(alu, ALU_REM, 123, 0, 123, "REM_ByZero");
    test_alu(alu, ALU_DIV, 0x8000000000000000ULL, 0xFFFFFFFFFFFFFFFFULL,
             0x8000000000000000ULL, "DIV_Overflow");
    test_alu(alu, ALU_REM, 0x8000000000000000ULL, 0xFFFFFFFFFFFFFFFFULL, 0, "REM_Overflow");
    std::cout << "----------------------------------" << std::endl;
    std::cout << "\033[32mAll tests passed!\033[0m" << std::endl;
    std::cout << "----------------------------------" << std::endl;

    delete alu;
    return 0;
}