"""Pytest test discovery orchestration."""

from __future__ import annotations

import enum
import logging
import os
import pickle
import re
import subprocess
from collections import defaultdict
from pathlib import Path
from typing import TYPE_CHECKING, final

from .linking import module_name_from_file_path, process_test_files
from .models import (
    FunctionCalledInTest,
    TestsInFile,
    TestType,
)

if TYPE_CHECKING:
    from .._model import FunctionToOptimize
    from ..testing.models import TestConfig

log = logging.getLogger(__name__)

ERROR_PATTERN = re.compile(
    r"={3,}\s*ERRORS\s*={3,}\n([\s\S]*?)(?:={3,}|$)",
)


@final
class PytestExitCode(enum.IntEnum):
    """Pytest exit codes without importing pytest."""

    OK = 0
    TESTS_FAILED = 1
    INTERRUPTED = 2
    INTERNAL_ERROR = 3
    USAGE_ERROR = 4
    NO_TESTS_COLLECTED = 5


def _count_results(
    function_to_tests: dict[str, set[FunctionCalledInTest]],
) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
    """Count discovered tests and replay tests from the result map."""
    num_tests = 0
    num_replay = 0
    for test_set in function_to_tests.values():
        for fct in test_set:
            num_tests += 1
            if fct.tests_in_file.test_type == TestType.REPLAY_TEST:
                num_replay += 1
    return function_to_tests, num_tests, num_replay


def existing_unit_test_count(
    func: FunctionToOptimize,
    project_root: Path,
    function_to_tests: dict[str, set[FunctionCalledInTest]],
) -> int:
    """Count unique existing unit tests for *func*."""
    key = (
        f"{module_name_from_file_path(func.file_path, project_root)}"
        f".{func.qualified_name}"
    )
    tests = function_to_tests.get(key, set())
    seen: set[tuple[Path, str | None, str]] = set()
    for t in tests:
        if t.tests_in_file.test_type != TestType.EXISTING_UNIT_TEST:
            continue
        tif = t.tests_in_file
        base_name = tif.test_function.split("[", 1)[0]
        seen.add((tif.test_file, tif.test_class, base_name))
    return len(seen)


def discover_unit_tests(
    cfg: TestConfig,
    discover_only_these_tests: list[Path] | None = None,
    file_to_funcs_to_optimize: (
        dict[Path, list[FunctionToOptimize]] | None
    ) = None,
) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
    """Discover unit tests and link them to the functions they call."""
    functions_to_optimize = None
    if file_to_funcs_to_optimize:
        functions_to_optimize = [
            func
            for funcs in file_to_funcs_to_optimize.values()
            for func in funcs
        ]
    return discover_tests_pytest(
        cfg,
        discover_only_these_tests,
        functions_to_optimize,
    )


def discover_tests_pytest(  # noqa: C901, PLR0912, PLR0915
    cfg: TestConfig,
    discover_only_these_tests: list[Path] | None = None,
    functions_to_optimize: (list[FunctionToOptimize] | None) = None,
) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
    """Discover pytest tests via subprocess collection."""
    from codeflash_python._compat import SAFE_SYS_EXECUTABLE  # noqa: PLC0415

    from ..runtime._codeflash_wrap_decorator import (  # noqa: PLC0415
        get_run_tmp_file,
    )
    from ..testing._pytest_config import custom_addopts  # noqa: PLC0415

    tests_root = cfg.tests_root
    project_root = Path(cfg.project_root_path)

    tmp_pickle_path = get_run_tmp_file(Path("collected_tests.pkl"))
    discovery_script = (
        Path(__file__).parent.parent / "analysis" / "_discovery_worker.py"
    )

    # Ensure PYTHONPATH covers both module_root and its parent so
    # pytest can import packages in src-layout projects regardless
    # of whether module_root is the source root (``src``) or a
    # package inside it (``src/aviary``).
    env = os.environ.copy()
    if cfg.module_root is not None:
        module_root_str = str(Path(cfg.module_root))
        parent = str(Path(cfg.module_root).parent)
        extra = (
            f"{module_root_str}{os.pathsep}{parent}"
            if module_root_str != parent
            else parent
        )
        existing = env.get("PYTHONPATH", "")
        env["PYTHONPATH"] = (
            f"{extra}{os.pathsep}{existing}" if existing else extra
        )

    with custom_addopts():
        result = subprocess.run(  # noqa: S603
            [
                SAFE_SYS_EXECUTABLE,
                str(discovery_script),
                str(project_root),
                str(tests_root),
                str(tmp_pickle_path),
            ],
            cwd=str(project_root),
            env=env,
            check=False,
            text=True,
            capture_output=True,
        )

    try:
        with tmp_pickle_path.open(mode="rb") as f:
            exitcode, tests, _pytest_rootdir = pickle.load(  # noqa: S301
                f,
            )
    except Exception:
        tests = []
        log.exception("Failed to discover tests")
        exitcode = -1
    finally:
        tmp_pickle_path.unlink(missing_ok=True)

    if exitcode != 0:
        if (
            exitcode == 2  # noqa: PLR2004
            and "ERROR collecting" in result.stdout
        ):
            match = ERROR_PATTERN.search(result.stdout)
            error_section = match.group(1) if match else result.stdout
            log.warning(
                "Failed to collect tests. Exit code: %s\n%s",
                exitcode,
                error_section,
            )
        elif 0 <= exitcode <= 5:  # noqa: PLR2004
            log.warning(
                "Failed to collect tests. Exit code: %s=%s",
                exitcode,
                PytestExitCode(exitcode).name,
            )
        else:
            log.warning(
                "Failed to collect tests. Exit code: %s",
                exitcode,
            )
    else:
        log.debug("Pytest collection exit code: %s", exitcode)

    # Build file_to_test_map from collected tests
    if discover_only_these_tests:
        resolved_discover_only: set[Path] | None = {
            p.resolve() for p in discover_only_these_tests
        }
    else:
        resolved_discover_only = None

    file_to_test_map: dict[Path, list[TestsInFile]] = defaultdict(
        list,
    )
    for test in tests:
        if "__replay_test" in test["test_file"]:
            test_type = TestType.REPLAY_TEST
        elif "test_concolic_coverage" in test["test_file"]:
            test_type = TestType.CONCOLIC_COVERAGE_TEST
        else:
            test_type = TestType.EXISTING_UNIT_TEST

        test_file_path = Path(test["test_file"]).resolve()
        test_obj = TestsInFile(
            test_file=test_file_path,
            test_class=test["test_class"],
            test_function=test["test_function"],
            test_type=test_type,
        )
        if (
            resolved_discover_only
            and test_obj.test_file not in resolved_discover_only
        ):
            continue
        file_to_test_map[test_obj.test_file].append(test_obj)

    function_to_tests = process_test_files(
        file_to_test_map,
        project_root,
        functions_to_optimize,
    )
    return _count_results(function_to_tests)
