import json
import os
import re
import subprocess

# File paths
SOLUTIONS_FILE = "solutions.json"
PROBLEMS_FILE = "problems.json"
TEMP_VERILOG_FILE = "temp.v"
TEMP_TESTBENCH_FILE = "testbench.v"
TCL_SCRIPT_FILE = "run_testbench.tcl"

# --- CONFIGURATION ---
# Default fallback path if env var is missing
DEFAULT_VIVADO_PATH = "/tools/Xilinx/Vivado/2023.1/bin" 

# FIX 1: Add 'top_module' as an argument so the function can use it
def write_tcl(top_module):
    # Generate the TCL script for Vivado
    tcl_commands = f"""
    create_project temp_project ./temp_project -force -part xc7z020clg400-1
    set_property source_mgmt_mode All [current_project]
    add_files {TEMP_VERILOG_FILE}
    add_files -fileset sim_1 {TEMP_TESTBENCH_FILE}
    set_property top {top_module} [get_filesets sim_1]
    launch_simulation -simset sim_1 -mode behavioral
    run 3000ns
    close_sim
    exit
    """
    # Write the Tcl script
    with open(TCL_SCRIPT_FILE, "w", encoding="utf-8") as file:
        file.write(tcl_commands)

# Function to extract the top module name from the testbench
def extract_top_module_name(testbench_file):
    with open(testbench_file, 'r', encoding="utf-8") as file:
        for line in file:
            match = re.search(r'\s*module\s+(\w+)\s*;', line)
            if match:
                return match.group(1)  # Extract module name
    return None  # Return None if no module found

def run_functional_correctness():
    # Load JSON files
    if not os.path.exists(SOLUTIONS_FILE):
        print(f"Error: {SOLUTIONS_FILE} not found.")
        return

    with open(SOLUTIONS_FILE, "r", encoding="utf-8") as file:
        solutions_data = json.load(file)

    with open(PROBLEMS_FILE, "r", encoding="utf-8") as file:
        problems_data = json.load(file)

    # Map module names to their testbenches
    module_testbenches = {}
    for category, problems in problems_data.items():
        for problem in problems:
            module_name = problem.get("module")
            testbench_code = problem.get("Testbench")
            if module_name and testbench_code:
                module_testbenches[module_name] = testbench_code

    # FIX 2: Handle Linux Path correctly (No .bat)
    vivado_path_env = os.environ.get("vivado")
    
    if vivado_path_env:
        vivado_bin = os.path.join(vivado_path_env, "vivado")
    elif os.path.exists(os.path.join(DEFAULT_VIVADO_PATH, "vivado")):
        vivado_bin = os.path.join(DEFAULT_VIVADO_PATH, "vivado")
    else:
        vivado_bin = "vivado" # Hope it's in PATH
    print(vivado_bin)
    # Iterate over solutions and test them
    for model, categories in solutions_data.items():
        for category, modules in categories.items():
            for module_entry in modules:
                module_name = module_entry["module"]

                if module_name not in module_testbenches:
                    print(f"Skipping {module_name}: No testbench found.")
                    continue

                testbench_code = module_testbenches[module_name]
                solutions = module_entry["solutions"]

                # Iterate over all solutions
                for idx, solution_entry in enumerate(solutions):
                    verilog_code = solution_entry["solution"]

                    # Write the Verilog design to a file
                    with open(TEMP_VERILOG_FILE, "w", encoding="utf-8") as f:
                        f.write(verilog_code)

                    # Write the testbench to a file
                    with open(TEMP_TESTBENCH_FILE, "w", encoding="utf-8") as f:
                        f.write(testbench_code)

                    # Extract the top module name
                    top_module = extract_top_module_name(TEMP_TESTBENCH_FILE)
                    if not top_module:
                        print(f"Error: Could not extract top module from {module_name}. Skipping...")
                        solution_entry["pass"] = "Error: Could not extract top module."
                        continue

                    print(f"Testing {module_name} (Solution {idx+1})...")

                    # FIX 3: Pass the variable to the function
                    write_tcl(top_module)

                    # Run Vivado in batch mode
                    try:
                        process = subprocess.run([vivado_bin, "-mode", "batch", "-source", TCL_SCRIPT_FILE], capture_output=True, text=True)
                        output_log = process.stdout + "\n" + process.stderr
                    except FileNotFoundError:
                        print("CRITICAL ERROR: Vivado executable not found. Check path.")
                        return

                    test_passed = "All tests passed" in output_log

                    # Determine pass/fail status
                    if test_passed:
                        solution_entry["pass"] = "true"
                        print(f"   ✅ PASS")
                    else:
                        solution_entry["pass"] = "false" # Keep it simple for now
                        print(f"   ❌ FAIL")

                    # Save results incrementally
                    with open(SOLUTIONS_FILE, "w", encoding="utf-8") as file:
                        json.dump(solutions_data, file, indent=4)

    print("All tests completed.")

if __name__ == "__main__":
    run_functional_correctness()