import os
import tempfile
from pathlib import Path

from codeflash.code_utils.code_utils import ImportErrorPattern
from codeflash.languages import current_language_support
from codeflash.models.models import TestFile, TestFiles, TestType
from codeflash.verification.parse_test_output import parse_test_xml
from codeflash.verification.verification_utils import TestConfig


def test_unittest_runner():
    code = """import time
import gc
import unittest
def sorter(arr):
    arr.sort()
    return arr
class TestUnittestRunnerSorter(unittest.TestCase):
    def test_sort(self):
        arr = [5, 4, 3, 2, 1, 0]
        gc.disable()
        counter = time.perf_counter_ns()
        output = sorter(arr)
        duration = time.perf_counter_ns() - counter
        gc.enable()
        print(f"#####test_sorter__unit_test_0:TestUnittestRunnerSorter.test_sort:sorter:0#####{duration}^^^^^")
"""
    cur_dir_path = Path(__file__).resolve().parent
    config = TestConfig(
        tests_root=cur_dir_path,
        project_root_path=cur_dir_path,
        test_framework="unittest",
        tests_project_rootdir=cur_dir_path.parent,
    )

    test_env = os.environ.copy()
    test_env["CODEFLASH_TEST_ITERATION"] = "0"
    test_env["CODEFLASH_TRACER_DISABLE"] = "1"
    if "PYTHONPATH" not in test_env:
        test_env["PYTHONPATH"] = str(config.project_root_path)
    else:
        test_env["PYTHONPATH"] += os.pathsep + str(config.project_root_path)

    with tempfile.TemporaryDirectory(dir=cur_dir_path) as temp_dir:
        test_file_path = Path(temp_dir) / "test_xx.py"
        test_files = TestFiles(
            test_files=[TestFile(instrumented_behavior_file_path=test_file_path, test_type=TestType.EXISTING_UNIT_TEST)]
        )
        test_file_path.write_text(code, encoding="utf-8")
        result_file, process, _, _ = current_language_support().run_behavioral_tests(
            test_paths=test_files, test_env=test_env, cwd=Path(config.project_root_path)
        )
        results = parse_test_xml(result_file, test_files, config, process)
    assert results[0].did_pass, "Test did not pass as expected"
    result_file.unlink(missing_ok=True)


def test_pytest_runner():
    code = """
def sorter(arr):
    arr.sort()
    return arr

def test_sort():
    arr = [5, 4, 3, 2, 1, 0]
    output = sorter(arr)
    assert output == [0, 1, 2, 3, 4, 5]
"""
    cur_dir_path = Path(__file__).resolve().parent
    config = TestConfig(
        tests_root=cur_dir_path,
        project_root_path=cur_dir_path,
        test_framework="pytest",
        tests_project_rootdir=cur_dir_path.parent,
    )

    test_env = os.environ.copy()
    test_env["CODEFLASH_TEST_ITERATION"] = "0"
    test_env["CODEFLASH_TRACER_DISABLE"] = "1"
    if "PYTHONPATH" not in test_env:
        test_env["PYTHONPATH"] = str(config.project_root_path)
    else:
        test_env["PYTHONPATH"] += os.pathsep + str(config.project_root_path)

    with tempfile.TemporaryDirectory(dir=cur_dir_path) as temp_dir:
        test_file_path = Path(temp_dir) / "test_xx.py"
        test_files = TestFiles(
            test_files=[TestFile(instrumented_behavior_file_path=test_file_path, test_type=TestType.EXISTING_UNIT_TEST)]
        )
        test_file_path.write_text(code, encoding="utf-8")
        result_file, process, _, _ = current_language_support().run_behavioral_tests(
            test_paths=test_files, test_env=test_env, cwd=Path(config.project_root_path), timeout=1
        )
        results = parse_test_xml(
            test_xml_file_path=result_file, test_files=test_files, test_config=config, run_result=process
        )
    assert results[0].did_pass, "Test did not pass as expected"
    result_file.unlink(missing_ok=True)

    code = """import torch_does_not_exist
def sorter(arr):
    print(torch.ones(1))
    arr.sort()
    return arr

def test_sort():
    arr = [5, 4, 3, 2, 1, 0]
    output = sorter(arr)
    assert output == [0, 1, 2, 3, 4, 5]
    """
    cur_dir_path = Path(__file__).resolve().parent
    config = TestConfig(
        tests_root=cur_dir_path,
        project_root_path=cur_dir_path,
        test_framework="pytest",
        tests_project_rootdir=cur_dir_path.parent,
    )

    test_env = os.environ.copy()
    test_env["CODEFLASH_TEST_ITERATION"] = "0"
    test_env["CODEFLASH_TRACER_DISABLE"] = "1"
    if "PYTHONPATH" not in test_env:
        test_env["PYTHONPATH"] = str(config.project_root_path)
    else:
        test_env["PYTHONPATH"] += os.pathsep + str(config.project_root_path)

    with tempfile.TemporaryDirectory(dir=cur_dir_path) as temp_dir:
        test_file_path = Path(temp_dir) / "test_xx.py"
        test_files = TestFiles(
            test_files=[TestFile(instrumented_behavior_file_path=test_file_path, test_type=TestType.EXISTING_UNIT_TEST)]
        )
        test_file_path.write_text(code, encoding="utf-8")
        result_file, process, _, _ = current_language_support().run_behavioral_tests(
            test_paths=test_files, test_env=test_env, cwd=Path(config.project_root_path), timeout=1
        )
        results = parse_test_xml(
            test_xml_file_path=result_file, test_files=test_files, test_config=config, run_result=process
        )
    match = ImportErrorPattern.search(process.stdout).group()
    assert match == "ModuleNotFoundError: No module named 'torch_does_not_exist'"
    result_file.unlink(missing_ok=True)