#!/bin/python3
import argparse
import os
from glob import glob
import subprocess
from colorama import Fore, Style
import time


def check_dependencies() -> None:
    # Necessary commands list
    commands = ["astyle", "gcc", "ghc", "java", "javac", "mvn", "python3"]

    for command in commands:
        if os.system(f"which {command} > /dev/null") != 0:
            print(f"{Fore.RED}Error: {command} not found{Style.RESET_ALL}")
    return


def run_input_cases(program, args):
    verbose = args.verbose
    output = args.output
    tests_dir = args.tests_dir if args.tests_dir else "."
    time_flag = args.time

    input_files = glob(os.path.join(tests_dir, "*.in"))

    if len(input_files) < 1:
        if time_flag:
            start_time = time.perf_counter()
            result = subprocess.run(
                program,
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                universal_newlines=True,
            )
            end_time = time.perf_counter()
            elapsed_time = end_time - start_time
            print(result.stdout.strip())
            print(f"Time taken: {elapsed_time:.4f} seconds")
            return
        else:
            return os.system(program)

    for input_file in input_files:
        print(f"Running {input_file}")
        end_command = f"{program} < {input_file}"

        with open(input_file, "r") as f:
            input_data = f.read()

        out_file = input_file.replace(".in", ".out")
        expected_output = None
        if os.path.exists(out_file):
            with open(out_file, "r") as f:
                expected_output = f.read()

        if time_flag:
            start_time = time.perf_counter()
            result = subprocess.run(
                end_command,
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                universal_newlines=True,
            )
            end_time = time.perf_counter()
            elapsed_time = end_time - start_time
        else:
            result = subprocess.run(
                end_command,
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                universal_newlines=True,
            )

        program_output = result.stdout

        if verbose and not output:
            print("Input:")
            print(input_data)
            print("")
            if expected_output is not None:
                print("Expected output:")
                print(expected_output)
                print("")
            print("Program output:")
            print(program_output.strip())
            print("")

        if output:
            if expected_output is not None:
                # Compare the outputs
                if program_output == expected_output:
                    print("✅")
                else:
                    print("❌")
                    # green expected output and red program output
                    print(f"{Fore.GREEN}Expected output:{Style.RESET_ALL}")
                    print(expected_output)
                    print(f"{Fore.RED}Program output:{Style.RESET_ALL}")
                    print(program_output.strip())
            else:
                print(
                    f"{Fore.YELLOW}Expected output file not found; cannot compare outputs.{Style.RESET_ALL}"
                )
                if not verbose:
                    print("Program output:")
                    print(program_output.strip())
        else:
            if not verbose:
                print(program_output.strip())

        if time_flag:
            print("=====")
            print(f"⏳: {elapsed_time:.4f} s")

        print("-----------")
    return


def python_compile(file_name, args):
    try:
        if subprocess.run(["which", "pypy3"]).returncode == 0:
            run_input_cases("pypy3 " + file_name, args)
            return
    except Exception:
        pass
    finally:
        run_input_cases("python3 " + file_name, args)
    return


def java_compile(file_name, args):
    if args.alternative:
        os.system("mvn clean verify")
        run_input_cases("java -jar target/*.jar", args)
        return
    os.system("javac " + file_name)
    class_name = os.path.splitext(file_name)[0]
    run_input_cases("java " + class_name, args)
    return


def haskell_compile(file_name, args):
    os.system("ghc -O2 " + file_name + " -o program")
    run_input_cases("./program", args)
    return


def c_compile(file_name, args):
    # Style the code in Meijster's way
    os.system("astyle -A2s2cxgk3W3xbj " + file_name)
    os.system(
        "gcc -Wall -pedantic --std=c99 -g -o program -lm -Wno-unused-result " + file_name
    )
    run_input_cases("./program", args)
    return


def main():
    SUPPORTED_LANGS = {
        ".c": c_compile,
        ".py": python_compile,
        ".java": java_compile,
        ".hs": haskell_compile,
    }
    parser = argparse.ArgumentParser(
        prog="UCompiler", description="Compiles based on language", epilog="LOL"
    )
    parser.add_argument("filename", metavar="<file>", type=str, help="File to compile")
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        default=False,
        help="Prints input and output",
    )
    parser.add_argument(
        "-a",
        "--alternative",
        action="store_true",
        default=False,
        help="Uses alternative compiler to compile",
    )
    parser.add_argument(
        "-o",
        "--output",
        action="store_true",
        default=False,
        help="Checks diff for each .in file with .out file",
    )
    parser.add_argument(
        "-t", "--tests", metavar="<path>", type=str, help="Directory to run tests"
    )
    parser.add_argument(
        "-T",
        "--time",
        action="store_true",
        default=False,
        help="Prints time taken to run each test case",
    )

    args = parser.parse_args()
    file_name = args.filename
    file_ext = os.path.splitext(file_name)[1]
    if file_ext not in SUPPORTED_LANGS:
        print("Unsupported language")
        return

    args.tests_dir = args.tests if args.tests else "."

    SUPPORTED_LANGS[file_ext](file_name, args)

    return


if __name__ == "__main__":
    check_dependencies()
    main()