# ruff: noqa: SLF001
from __future__ import annotations

import ast
import enum
import hashlib
import os
import pickle
import re
import sqlite3
import subprocess
import unittest
from collections import defaultdict
from functools import lru_cache
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Optional, final

if TYPE_CHECKING:
    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from pydantic.dataclasses import dataclass
from rich.panel import Panel
from rich.text import Text

from codeflash.cli_cmds.console import console, logger, test_files_progress_bar
from codeflash.code_utils.code_utils import (
    ImportErrorPattern,
    custom_addopts,
    get_run_tmp_file,
    module_name_from_file_path,
)
from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE, codeflash_cache_db
from codeflash.code_utils.shell_utils import get_cross_platform_subprocess_run_args
from codeflash.models.models import CodePosition, FunctionCalledInTest, TestsInFile, TestType

if TYPE_CHECKING:
    from codeflash.verification.verification_utils import TestConfig


def existing_unit_test_count(
    func: FunctionToOptimize, project_root: Path, function_to_tests: dict[str, set[FunctionCalledInTest]]
) -> int:
    key = f"{module_name_from_file_path_cached(func.file_path, project_root)}.{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)


@final
class PytestExitCode(enum.IntEnum):  # don't need to import entire pytest just for this
    #: Tests passed.
    OK = 0
    #: Tests failed.
    TESTS_FAILED = 1
    #: pytest was interrupted.
    INTERRUPTED = 2
    #: An internal error got in the way.
    INTERNAL_ERROR = 3
    #: pytest was misused.
    USAGE_ERROR = 4
    #: pytest couldn't find tests.
    NO_TESTS_COLLECTED = 5


@dataclass(frozen=True)
class TestFunction:
    function_name: str
    test_class: Optional[str]
    parameters: Optional[str]
    test_type: TestType


ERROR_PATTERN = re.compile(r"={3,}\s*ERRORS\s*={3,}\n([\s\S]*?)(?:={3,}|$)")
PYTEST_PARAMETERIZED_TEST_NAME_REGEX = re.compile(r"[\[\]]")
UNITTEST_PARAMETERIZED_TEST_NAME_REGEX = re.compile(r"^test_\w+_\d+(?:_\w+)*")
UNITTEST_STRIP_NUMBERED_SUFFIX_REGEX = re.compile(r"_\d+(?:_\w+)*$")
FUNCTION_NAME_REGEX = re.compile(r"([^.]+)\.([a-zA-Z0-9_]+)$")


class TestsCache:
    SCHEMA_VERSION = 1  # Increment this when schema changes

    def __init__(self, project_root_path: Path) -> None:
        self.project_root_path = project_root_path.resolve().as_posix()
        self.connection = sqlite3.connect(codeflash_cache_db)
        self.cur = self.connection.cursor()

        self.cur.execute(
            """
            CREATE TABLE IF NOT EXISTS schema_version(
                version INTEGER PRIMARY KEY
            )
            """
        )

        self.cur.execute("SELECT version FROM schema_version")
        result = self.cur.fetchone()
        current_version = result[0] if result else None

        if current_version != self.SCHEMA_VERSION:
            logger.debug(
                f"Schema version mismatch (current: {current_version}, expected: {self.SCHEMA_VERSION}). Recreating tables."
            )
            self.cur.execute("DROP TABLE IF EXISTS discovered_tests")
            self.cur.execute("DROP INDEX IF EXISTS idx_discovered_tests_project_file_path_hash")
            self.cur.execute("DELETE FROM schema_version")
            self.cur.execute("INSERT INTO schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,))
            self.connection.commit()

        self.cur.execute(
            """
            CREATE TABLE IF NOT EXISTS discovered_tests(
                project_root_path TEXT,
                file_path TEXT,
                file_hash TEXT,
                qualified_name_with_modules_from_root TEXT,
                function_name TEXT,
                test_class TEXT,
                test_function TEXT,
                test_type TEXT,
                line_number INTEGER,
                col_number INTEGER
            )
            """
        )
        self.cur.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_discovered_tests_project_file_path_hash
            ON discovered_tests (project_root_path, file_path, file_hash)
            """
        )

        self.memory_cache = {}
        self.pending_rows: list[tuple[str, str, str, str, str, str, int | TestType, int, int]] = []
        self.writes_enabled = True

    def insert_test(
        self,
        file_path: str,
        file_hash: str,
        qualified_name_with_modules_from_root: str,
        function_name: str,
        test_class: str,
        test_function: str,
        test_type: TestType,
        line_number: int,
        col_number: int,
    ) -> None:
        test_type_value = test_type.value if hasattr(test_type, "value") else test_type
        self.pending_rows.append(
            (
                file_path,
                file_hash,
                qualified_name_with_modules_from_root,
                function_name,
                test_class,
                test_function,
                test_type_value,
                line_number,
                col_number,
            )
        )

    def flush(self) -> None:
        if not self.pending_rows:
            return
        if not self.writes_enabled:
            self.pending_rows.clear()
            return
        try:
            self.cur.executemany(
                "INSERT INTO discovered_tests VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                [(self.project_root_path, *row) for row in self.pending_rows],
            )
            self.connection.commit()
        except sqlite3.OperationalError as e:
            logger.debug(f"Failed to persist discovered test cache, disabling cache writes: {e}")
            self.writes_enabled = False
        finally:
            self.pending_rows.clear()

    def get_function_to_test_map_for_file(
        self, file_path: str, file_hash: str
    ) -> dict[str, set[FunctionCalledInTest]] | None:
        cache_key = (self.project_root_path, file_path, file_hash)
        if cache_key in self.memory_cache:
            return self.memory_cache[cache_key]

        self.cur.execute(
            "SELECT * FROM discovered_tests WHERE project_root_path = ? AND file_path = ? AND file_hash = ?",
            (self.project_root_path, file_path, file_hash),
        )
        rows = self.cur.fetchall()
        if not rows:
            return None

        function_to_test_map = defaultdict(set)

        for row in rows:
            qualified_name_with_modules_from_root = row[3]
            function_called_in_test = FunctionCalledInTest(
                tests_in_file=TestsInFile(
                    test_file=Path(row[1]), test_class=row[5], test_function=row[6], test_type=TestType(int(row[7]))
                ),
                position=CodePosition(line_no=row[8], col_no=row[9]),
            )
            function_to_test_map[qualified_name_with_modules_from_root].add(function_called_in_test)

        result = dict(function_to_test_map)
        self.memory_cache[cache_key] = result
        return result

    @staticmethod
    def compute_file_hash(path: Path) -> str:
        h = hashlib.sha256(usedforsecurity=False)
        with path.open("rb", buffering=0) as f:
            buf = bytearray(8192)
            mv = memoryview(buf)
            while True:
                n = f.readinto(mv)
                if n == 0:
                    break
                h.update(mv[:n])
        return h.hexdigest()

    def close(self) -> None:
        self.flush()
        self.cur.close()
        self.connection.close()


class ImportAnalyzer(ast.NodeVisitor):
    """AST-based analyzer to check if any qualified names from function_names_to_find are imported or used in a test file."""

    def __init__(self, function_names_to_find: set[str]) -> None:
        self.function_names_to_find = function_names_to_find
        self.found_any_target_function: bool = False
        self.found_qualified_name = None
        self.imported_modules: set[str] = set()
        self.has_dynamic_imports: bool = False
        self.wildcard_modules: set[str] = set()
        # Track aliases: alias_name -> original_name
        self.alias_mapping: dict[str, str] = {}
        # Track instances: variable_name -> class_name
        self.instance_mapping: dict[str, str] = {}

        # Precompute function_names for prefix search
        # For prefix match, store mapping from prefix-root to candidates for O(1) matching
        self._exact_names = function_names_to_find
        self._prefix_roots: dict[str, list[str]] = {}
        # Precompute sets for faster lookup during visit_Attribute()
        self._dot_names: set[str] = set()
        self._dot_methods: dict[str, set[str]] = {}
        self._class_method_to_target: dict[tuple[str, str], str] = {}

        # Optimize prefix-roots and dot_methods construction
        add_dot_methods = self._dot_methods.setdefault
        add_prefix_roots = self._prefix_roots.setdefault
        dot_names_add = self._dot_names.add
        class_method_to_target = self._class_method_to_target
        for name in function_names_to_find:
            if "." in name:
                root, method = name.rsplit(".", 1)
                dot_names_add(name)
                add_dot_methods(method, set()).add(root)
                class_method_to_target[(root, method)] = name
                root_prefix = name.split(".", 1)[0]
                add_prefix_roots(root_prefix, []).append(name)

    def visit_Import(self, node: ast.Import) -> None:
        """Handle 'import module' statements."""
        if self.found_any_target_function:
            return

        for alias in node.names:
            module_name = alias.asname if alias.asname else alias.name
            self.imported_modules.add(module_name)

            # Check for dynamic import modules
            if alias.name == "importlib":
                self.has_dynamic_imports = True

            # Check if module itself is a target qualified name
            if module_name in self.function_names_to_find:
                self.found_any_target_function = True
                self.found_qualified_name = module_name
                return
            # Check if any target qualified name starts with this module
            for target_func in self.function_names_to_find:
                if target_func.startswith(f"{module_name}."):
                    self.found_any_target_function = True
                    self.found_qualified_name = target_func
                    return

    def visit_Assign(self, node: ast.Assign) -> None:
        """Track variable assignments, especially class instantiations."""
        if self.found_any_target_function:
            return

        # Check if the assignment is a class instantiation
        value = node.value
        if isinstance(value, ast.Call) and isinstance(value.func, ast.Name):
            class_name = value.func.id
            if class_name in self.imported_modules:
                # Map the variable to the actual class name (handling aliases)
                original_class = self.alias_mapping.get(class_name, class_name)
                # Use list comprehension for direct assignment to instance_mapping, reducing loop overhead
                targets = node.targets
                instance_mapping = self.instance_mapping
                # since ast.Name nodes are heavily used, avoid local lookup for isinstance
                # and reuse locals for faster attribute access
                for target in targets:
                    if isinstance(target, ast.Name):
                        instance_mapping[target.id] = original_class

        # Continue visiting child nodes
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        """Handle 'from module import name' statements."""
        if self.found_any_target_function:
            return

        mod = node.module
        if not mod:
            return

        fnames = self._exact_names
        proots = self._prefix_roots

        for alias in node.names:
            aname = alias.name
            if aname == "*":
                self.wildcard_modules.add(mod)
                continue

            imported_name = alias.asname if alias.asname else aname
            self.imported_modules.add(imported_name)

            if alias.asname:
                self.alias_mapping[imported_name] = aname

            # Fast check for dynamic import
            if mod == "importlib" and aname == "import_module":
                self.has_dynamic_imports = True

            qname = f"{mod}.{aname}"

            # Fast exact match check
            if aname in fnames:
                self.found_any_target_function = True
                self.found_qualified_name = aname
                return
            if qname in fnames:
                self.found_any_target_function = True
                self.found_qualified_name = qname
                return

            # Check if any target function is a method of the imported class/module
            # Be conservative except when an alias is used (which requires exact method matching)
            for target_func in fnames:
                if "." in target_func:
                    class_name, _method_name = target_func.split(".", 1)
                    if aname == class_name and not alias.asname:
                        self.found_any_target_function = True
                        self.found_qualified_name = target_func
                        return
                        # If an alias is used, track it for later method access detection
                        # The actual method usage will be detected in visit_Attribute

            prefix = qname + "."
            # Only bother if one of the targets startswith the prefix-root
            candidates = proots.get(qname, ())
            for target_func in candidates:
                if target_func.startswith(prefix):
                    self.found_any_target_function = True
                    self.found_qualified_name = target_func
                    return

    def visit_Attribute(self, node: ast.Attribute) -> None:
        """Handle attribute access like module.function_name."""
        if self.found_any_target_function:
            return

        # Check if this is accessing a target function through an imported module

        node_value = node.value
        node_attr = node.attr

        # Check if this is accessing a target function through an imported module

        # Accessing a target function through an imported module (fast path for imported modules)
        val_id = getattr(node_value, "id", None)
        if val_id is not None and val_id in self.imported_modules:
            if node_attr in self.function_names_to_find:
                self.found_any_target_function = True
                self.found_qualified_name = node_attr
                return
            # Methods via imported modules using precomputed _dot_methods and _class_method_to_target
            roots_possible = self._dot_methods.get(node_attr)
            if roots_possible:
                imported_name = val_id
                original_name = self.alias_mapping.get(imported_name, imported_name)
                if original_name in roots_possible:
                    self.found_any_target_function = True
                    self.found_qualified_name = self._class_method_to_target[(original_name, node_attr)]
                    return
                # Also check if the imported name itself (without resolving alias) matches
                # This handles cases where the class itself is the target
                if imported_name in roots_possible:
                    self.found_any_target_function = True
                    self.found_qualified_name = self._class_method_to_target.get(
                        (imported_name, node_attr), f"{imported_name}.{node_attr}"
                    )
                    return

        # Methods on instance variables (tighten type check and lookup, important for larger ASTs)
        if val_id is not None and val_id in self.instance_mapping:
            class_name = self.instance_mapping[val_id]
            roots_possible = self._dot_methods.get(node_attr)
            if roots_possible and class_name in roots_possible:
                self.found_any_target_function = True
                self.found_qualified_name = self._class_method_to_target[(class_name, node_attr)]
                return

        # Check for dynamic import match
        if self.has_dynamic_imports and node_attr in self.function_names_to_find:
            self.found_any_target_function = True
            self.found_qualified_name = node_attr
            return

        # Replace self.generic_visit with base class impl directly: removes an attribute lookup
        if not self.found_any_target_function:
            ast.NodeVisitor.generic_visit(self, node)

    def visit_Call(self, node: ast.Call) -> None:
        """Handle function calls, particularly __import__."""
        if self.found_any_target_function:
            return

        # Check if this is a __import__ call
        if isinstance(node.func, ast.Name) and node.func.id == "__import__":
            self.has_dynamic_imports = True
            # When __import__ is used, any target function could potentially be imported
            # Be conservative and assume it might import target functions

        self.generic_visit(node)

    def visit_Name(self, node: ast.Name) -> None:
        """Handle direct name usage like target_function()."""
        if self.found_any_target_function:
            return

        # Check for __import__ usage
        if node.id == "__import__":
            self.has_dynamic_imports = True

        # Check if this is a direct usage of a target function name
        # This catches cases like: result = target_function()
        if node.id in self.function_names_to_find:
            self.found_any_target_function = True
            self.found_qualified_name = node.id
            return

        # Check if this name could come from a wildcard import
        for wildcard_module in self.wildcard_modules:
            for target_func in self.function_names_to_find:
                # Check if target_func is from this wildcard module and name matches
                if target_func.startswith(f"{wildcard_module}.") and target_func.endswith(f".{node.id}"):
                    self.found_any_target_function = True
                    self.found_qualified_name = target_func
                    return

        self.generic_visit(node)

    def generic_visit(self, node: ast.AST) -> None:
        """Override generic_visit to stop traversal if a target function is found."""
        if self.found_any_target_function:
            return
        # Direct base call improves run speed (avoids extra method resolution)
        self._fast_generic_visit(node)

    def _fast_generic_visit(self, node: ast.AST) -> None:
        """Faster generic_visit: Inline traversal, avoiding method resolution overhead.

        Short-circuits (returns) if found_any_target_function is True.
        """
        # This logic is derived from ast.NodeVisitor.generic_visit, but with optimizations.
        if self.found_any_target_function:
            return

        # Local bindings for improved lookup speed (10-15% faster for inner loop)
        visit_cache = type(self).__dict__
        node_fields = node._fields

        # Use manual stack for iterative traversal, replacing recursion
        stack = [(node_fields, node)]
        append = stack.append
        pop = stack.pop

        while stack:
            fields, curr_node = pop()
            for field in fields:
                value = getattr(curr_node, field, None)
                if isinstance(value, list):
                    for item in value:
                        if self.found_any_target_function:
                            return
                        if isinstance(item, ast.AST):
                            # Method resolution: fast dict lookup first, then getattr fallback
                            meth = visit_cache.get("visit_" + item.__class__.__name__)
                            if meth is not None:
                                meth(self, item)
                            else:
                                append((item._fields, item))
                    continue
                if isinstance(value, ast.AST):
                    if self.found_any_target_function:
                        return
                    meth = visit_cache.get("visit_" + value.__class__.__name__)
                    if meth is not None:
                        meth(self, value)
                    else:
                        append((value._fields, value))


def analyze_imports_in_test_file(test_file_path: Path | str, target_functions: set[str]) -> bool:
    """Analyze a test file to see if it imports any of the target functions."""
    try:
        with Path(test_file_path).open("r", encoding="utf-8") as f:
            source_code = f.read()
        tree = ast.parse(source_code, filename=str(test_file_path))
        analyzer = ImportAnalyzer(target_functions)
        analyzer.visit(tree)
    except (SyntaxError, FileNotFoundError) as e:
        logger.debug(f"Failed to analyze imports in {test_file_path}: {e}")
        return True

    if analyzer.found_any_target_function:
        # logger.debug(f"Test file {test_file_path} imports target function: {analyzer.found_qualified_name}")
        return True

    # Be conservative with dynamic imports - if __import__ is used and a target function
    # is referenced, we should process the file
    if analyzer.has_dynamic_imports:
        # Check if any target function name appears as a string literal or direct usage
        for target_func in target_functions:
            if target_func in source_code:
                # logger.debug(f"Test file {test_file_path} has dynamic imports and references {target_func}")
                return True

    # logger.debug(f"Test file {test_file_path} does not import any target functions.")
    return False


def filter_test_files_by_imports(
    file_to_test_map: dict[Path, list[TestsInFile]], target_functions: set[str]
) -> dict[Path, list[TestsInFile]]:
    """Filter test files based on import analysis to reduce Jedi processing.

    Args:
        file_to_test_map: Original mapping of test files to test functions
        target_functions: Set of function names we're optimizing

    Returns:
        Filtered mapping of test files to test functions

    """
    if not target_functions:
        return file_to_test_map

    # logger.debug(f"Target functions for import filtering: {target_functions}")

    filtered_map = {}
    for test_file, test_functions in file_to_test_map.items():
        should_process = analyze_imports_in_test_file(test_file, target_functions)
        if should_process:
            filtered_map[test_file] = test_functions

    logger.debug(
        f"analyzed {len(file_to_test_map)} test files for imports, filtered down to {len(filtered_map)} relevant files"
    )
    return filtered_map


def _detect_language_from_functions(file_to_funcs: dict[Path, list[FunctionToOptimize]] | None) -> str | None:
    """Detect language from the functions to optimize.

    Args:
        file_to_funcs: Dictionary mapping file paths to functions.

    Returns:
        Language string (e.g., "python", "javascript") or None if not determinable.

    """
    if not file_to_funcs:
        return None

    for funcs in file_to_funcs.values():
        if funcs:
            return funcs[0].language
    return None


def discover_tests_for_language(
    cfg: TestConfig, language: str, file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] | None
) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
    """Discover tests using language-specific support.

    Args:
        cfg: Test configuration.
        language: Language identifier (e.g., "javascript").
        file_to_funcs_to_optimize: Dictionary mapping file paths to functions.

    Returns:
        Tuple of (function_to_tests_map, num_tests, num_replay_tests).

    """
    from codeflash.languages import get_language_support
    from codeflash.languages.base import Language

    try:
        lang_support = get_language_support(Language(language))
    except Exception:
        logger.warning(f"Unsupported language {language}, returning empty test map")
        return {}, 0, 0

    # Collect all functions and build a mapping from simple qualified_name to full qualified_name_with_modules
    all_functions: list[FunctionToOptimize] = []
    simple_to_full_name: dict[str, str] = {}
    if file_to_funcs_to_optimize:
        for funcs in file_to_funcs_to_optimize.values():
            for func in funcs:
                all_functions.append(func)
                # Map simple qualified_name to full qualified_name_with_modules_from_root
                simple_to_full_name[func.qualified_name] = func.qualified_name_with_modules_from_root(
                    cfg.project_root_path
                )

    # Use language support to discover tests
    test_map = lang_support.discover_tests(cfg.tests_root, all_functions)

    # Convert TestInfo back to FunctionCalledInTest format
    # Use the full qualified name (with modules) as the key for consistency with Python
    function_to_tests: dict[str, set[FunctionCalledInTest]] = defaultdict(set)
    num_unit_tests = 0
    num_replay_tests = 0

    for qualified_name, test_infos in test_map.items():
        # Convert simple qualified_name to full qualified_name_with_modules
        full_qualified_name = simple_to_full_name.get(qualified_name, qualified_name)
        for test_info in test_infos:
            is_replay = getattr(test_info, "is_replay", False)
            test_type = TestType.REPLAY_TEST if is_replay else TestType.EXISTING_UNIT_TEST
            function_to_tests[full_qualified_name].add(
                FunctionCalledInTest(
                    tests_in_file=TestsInFile(
                        test_file=test_info.test_file,
                        test_class=test_info.test_class,
                        test_function=test_info.test_name,
                        test_type=test_type,
                    ),
                    position=CodePosition(line_no=0, col_no=0),
                )
            )
            if is_replay:
                num_replay_tests += 1
            else:
                num_unit_tests += 1

    return dict(function_to_tests), num_unit_tests, num_replay_tests


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]:
    from codeflash.languages.current import current_language_support

    # Detect language from functions being optimized
    language = _detect_language_from_functions(file_to_funcs_to_optimize)

    # Route to language-specific test discovery for non-Python languages
    if current_language_support().test_result_serialization_format != "pickle":
        current_language_support().adjust_test_config_for_discovery(cfg)
        return discover_tests_for_language(cfg, language, file_to_funcs_to_optimize)

    # Existing Python logic
    framework_strategies: dict[str, Callable] = {"pytest": discover_tests_pytest, "unittest": discover_tests_unittest}
    strategy = framework_strategies.get(cfg.test_framework)
    if not strategy:
        error_message = f"Unsupported test framework: {cfg.test_framework}"
        raise ValueError(error_message)

    # Extract all functions to optimize for import filtering
    functions_to_optimize = None
    if file_to_funcs_to_optimize:
        functions_to_optimize = [func for funcs_list in file_to_funcs_to_optimize.values() for func in funcs_list]
    function_to_tests, num_discovered_tests, num_discovered_replay_tests = strategy(
        cfg, discover_only_these_tests, functions_to_optimize
    )
    return function_to_tests, num_discovered_tests, num_discovered_replay_tests


def discover_tests_pytest(
    cfg: TestConfig,
    discover_only_these_tests: list[Path] | None = None,
    functions_to_optimize: list[FunctionToOptimize] | None = None,
) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
    tests_root = cfg.tests_root
    project_root = cfg.project_root_path

    tmp_pickle_path = get_run_tmp_file("collected_tests.pkl")
    with custom_addopts():
        run_kwargs = get_cross_platform_subprocess_run_args(
            cwd=project_root, check=False, text=True, capture_output=True
        )
        result = subprocess.run(  # noqa: PLW1510
            [
                SAFE_SYS_EXECUTABLE,
                Path(__file__).parent / "pytest_new_process_discovery.py",
                str(project_root),
                str(tests_root),
                str(tmp_pickle_path),
            ],
            **run_kwargs,
        )
    try:
        with tmp_pickle_path.open(mode="rb") as f:
            exitcode, tests, pytest_rootdir = pickle.load(f)
    except Exception as e:
        tests, pytest_rootdir = [], None
        logger.exception(f"Failed to discover tests: {e}")
        exitcode = -1
    finally:
        tmp_pickle_path.unlink(missing_ok=True)
    if exitcode != 0:
        if exitcode == 2 and "ERROR collecting" in result.stdout:
            # Pattern matches "===== ERRORS =====" (any number of =) and captures everything after
            match = ERROR_PATTERN.search(result.stdout)
            error_section = match.group(1) if match else result.stdout

            logger.warning(
                f"Failed to collect tests. Pytest Exit code: {exitcode}={PytestExitCode(exitcode).name}\n {error_section}"
            )
            if "ModuleNotFoundError" in result.stdout:
                match = ImportErrorPattern.search(result.stdout)
                if match:
                    error_message = match.group()
                    panel = Panel(Text.from_markup(f"⚠️  {error_message} ", style="bold red"), expand=False)
                    console.print(panel)

        elif 0 <= exitcode <= 5:
            logger.warning(f"Failed to collect tests. Pytest Exit code: {exitcode}={PytestExitCode(exitcode).name}")
        else:
            logger.warning(f"Failed to collect tests. Pytest Exit code: {exitcode}")
        console.rule()
    else:
        logger.debug(f"Pytest collection exit code: {exitcode}")
    if pytest_rootdir is not None:
        cfg.tests_project_rootdir = Path(pytest_rootdir)
    if discover_only_these_tests:
        resolved_discover_only = {p.resolve() for p in discover_only_these_tests}
    else:
        resolved_discover_only = None
    file_to_test_map: dict[Path, list[FunctionCalledInTest]] = 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)
    # Within these test files, find the project functions they are referring to and return their names/locations
    return process_test_files(file_to_test_map, cfg, functions_to_optimize)


def discover_tests_unittest(
    cfg: TestConfig,
    discover_only_these_tests: list[Path] | None = None,
    functions_to_optimize: list[FunctionToOptimize] | None = None,
) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
    tests_root: Path = cfg.tests_root
    loader: unittest.TestLoader = unittest.TestLoader()
    tests: unittest.TestSuite = loader.discover(str(tests_root))
    file_to_test_map: defaultdict[Path, list[TestsInFile]] = defaultdict(list)

    def get_test_details(_test: unittest.TestCase) -> TestsInFile | None:
        _test_function, _test_module, _test_suite_name = (
            _test._testMethodName,
            _test.__class__.__module__,
            _test.__class__.__qualname__,
        )

        _test_module_path = Path(_test_module.replace(".", os.sep)).with_suffix(".py")
        _test_module_path = tests_root / _test_module_path
        if not _test_module_path.exists() or (
            discover_only_these_tests and _test_module_path not in discover_only_these_tests
        ):
            return None
        if "__replay_test" in str(_test_module_path):
            test_type = TestType.REPLAY_TEST
        elif "test_concolic_coverage" in str(_test_module_path):
            test_type = TestType.CONCOLIC_COVERAGE_TEST
        else:
            test_type = TestType.EXISTING_UNIT_TEST
        return TestsInFile(
            test_file=_test_module_path, test_function=_test_function, test_type=test_type, test_class=_test_suite_name
        )

    for _test_suite in tests._tests:
        for test_suite_2 in _test_suite._tests:
            if not hasattr(test_suite_2, "_tests"):
                logger.warning(f"Didn't find tests for {test_suite_2}")
                continue

            for test in test_suite_2._tests:
                # some test suites are nested, so we need to go deeper
                if not hasattr(test, "_testMethodName") and hasattr(test, "_tests"):
                    for test_2 in test._tests:
                        if not hasattr(test_2, "_testMethodName"):
                            logger.warning(f"Didn't find tests for {test_2}")  # it goes deeper?
                            continue
                        details = get_test_details(test_2)
                        if details is not None:
                            file_to_test_map[details.test_file].append(details)
                else:
                    details = get_test_details(test)
                    if details is not None:
                        file_to_test_map[details.test_file].append(details)
    return process_test_files(file_to_test_map, cfg, functions_to_optimize)


def discover_parameters_unittest(function_name: str) -> tuple[bool, str, str | None]:
    function_parts = function_name.split("_")
    if len(function_parts) > 1 and function_parts[-1].isdigit():
        return True, "_".join(function_parts[:-1]), function_parts[-1]

    return False, function_name, None


def process_test_files(
    file_to_test_map: dict[Path, list[TestsInFile]],
    cfg: TestConfig,
    functions_to_optimize: list[FunctionToOptimize] | None = None,
) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
    import jedi

    project_root_path = cfg.project_root_path
    test_framework = cfg.test_framework

    if functions_to_optimize:
        target_function_names = {func.qualified_name for func in functions_to_optimize}
        file_to_test_map = filter_test_files_by_imports(file_to_test_map, target_function_names)

    function_to_test_map = defaultdict(set)
    num_discovered_tests = 0
    num_discovered_replay_tests = 0
    functions_to_optimize_by_name: dict[str, list[FunctionToOptimize]] = defaultdict(list)
    if functions_to_optimize:
        for function_to_optimize in functions_to_optimize:
            functions_to_optimize_by_name[function_to_optimize.function_name].append(function_to_optimize)

    # Set up sys_path for Jedi to resolve imports correctly
    import sys

    jedi_sys_path = list(sys.path)
    # Add project root and its parent to sys_path so modules can be imported
    if str(project_root_path) not in jedi_sys_path:
        jedi_sys_path.insert(0, str(project_root_path))
    parent_path = project_root_path.parent
    if str(parent_path) not in jedi_sys_path:
        jedi_sys_path.insert(0, str(parent_path))

    jedi_project = jedi.Project(path=project_root_path, sys_path=jedi_sys_path)

    tests_cache = TestsCache(project_root_path)
    logger.info("!lsp|Discovering tests and processing unit tests")
    console.rule()
    with test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
        progress,
        task_id,
    ):
        for test_file, functions in file_to_test_map.items():
            file_hash = TestsCache.compute_file_hash(test_file)

            cached_function_to_test_map = tests_cache.get_function_to_test_map_for_file(str(test_file), file_hash)

            if cfg.use_cache and cached_function_to_test_map:
                for qualified_name, test_set in cached_function_to_test_map.items():
                    function_to_test_map[qualified_name].update(test_set)

                    for function_called_in_test in test_set:
                        if function_called_in_test.tests_in_file.test_type == TestType.REPLAY_TEST:
                            num_discovered_replay_tests += 1
                        num_discovered_tests += 1

                progress.advance(task_id)
                continue
            try:
                script = jedi.Script(path=test_file, project=jedi_project)
                test_functions = set()

                all_names = script.get_names(all_scopes=True, references=True)
                all_names_top = script.get_names(all_scopes=True)
                all_defs = [name for name in all_names if name.is_definition()]

                top_level_functions = {name.name: name for name in all_names_top if name.type == "function"}
                top_level_classes = {name.name: name for name in all_names_top if name.type == "class"}

            except Exception as e:
                logger.debug(f"Failed to get jedi script for {test_file}: {e}")
                progress.advance(task_id)
                continue

            if test_framework == "pytest":
                for function in functions:
                    if "[" in function.test_function:
                        function_name = PYTEST_PARAMETERIZED_TEST_NAME_REGEX.split(function.test_function)[0]
                        parameters = PYTEST_PARAMETERIZED_TEST_NAME_REGEX.split(function.test_function)[1]
                        if function_name in top_level_functions:
                            test_functions.add(
                                TestFunction(function_name, function.test_class, parameters, function.test_type)
                            )
                    elif function.test_function in top_level_functions:
                        test_functions.add(
                            TestFunction(function.test_function, function.test_class, None, function.test_type)
                        )
                    elif UNITTEST_PARAMETERIZED_TEST_NAME_REGEX.match(function.test_function):
                        base_name = UNITTEST_STRIP_NUMBERED_SUFFIX_REGEX.sub("", function.test_function)
                        if base_name in top_level_functions:
                            test_functions.add(
                                TestFunction(
                                    function_name=base_name,
                                    test_class=function.test_class,
                                    parameters=function.test_function,
                                    test_type=function.test_type,
                                )
                            )

            elif test_framework == "unittest":
                functions_to_search = [elem.test_function for elem in functions]
                test_suites = {elem.test_class for elem in functions}

                matching_names = test_suites & top_level_classes.keys()
                for matched_name in matching_names:
                    for def_name in all_defs:
                        if (
                            def_name.type == "function"
                            and def_name.full_name is not None
                            and f".{matched_name}." in def_name.full_name
                        ):
                            for function in functions_to_search:
                                (is_parameterized, new_function, parameters) = discover_parameters_unittest(function)

                                if is_parameterized and new_function == def_name.name:
                                    test_functions.add(
                                        TestFunction(
                                            function_name=def_name.name,
                                            test_class=matched_name,
                                            parameters=parameters,
                                            test_type=functions[0].test_type,
                                        )
                                    )
                                elif function == def_name.name:
                                    test_functions.add(
                                        TestFunction(
                                            function_name=def_name.name,
                                            test_class=matched_name,
                                            parameters=None,
                                            test_type=functions[0].test_type,
                                        )
                                    )

            test_functions_by_name = defaultdict(list)
            for func in test_functions:
                test_functions_by_name[func.function_name].append(func)

            test_function_names_set = set(test_functions_by_name.keys())
            relevant_names = []
            for name in all_names:
                if name.full_name is None:
                    continue
                match = FUNCTION_NAME_REGEX.search(name.full_name)
                if match and match.group(1) in test_function_names_set:
                    relevant_names.append((name, match.group(1)))

            for name, scope in relevant_names:
                try:
                    definition = name.goto(follow_imports=True, follow_builtin_imports=False)
                except Exception as e:
                    logger.debug(str(e))
                    continue
                try:
                    if not definition or definition[0].type != "function":
                        # Fallback: Try to match against functions_to_optimize when Jedi can't resolve
                        # This handles cases where Jedi fails with pytest fixtures
                        if functions_to_optimize_by_name and name.name:
                            for func_to_opt in functions_to_optimize_by_name.get(name.name, []):
                                qualified_name_with_modules = func_to_opt.qualified_name_with_modules_from_root(
                                    project_root_path
                                )

                                # Only add if this test actually tests the function we're optimizing
                                for test_func in test_functions_by_name[scope]:
                                    if test_func.parameters is not None:
                                        if test_framework == "pytest":
                                            scope_test_function = f"{test_func.function_name}[{test_func.parameters}]"
                                        else:  # unittest
                                            scope_test_function = f"{test_func.function_name}_{test_func.parameters}"
                                    else:
                                        scope_test_function = test_func.function_name

                                    function_to_test_map[qualified_name_with_modules].add(
                                        FunctionCalledInTest(
                                            tests_in_file=TestsInFile(
                                                test_file=test_file,
                                                test_class=test_func.test_class,
                                                test_function=scope_test_function,
                                                test_type=test_func.test_type,
                                            ),
                                            position=CodePosition(line_no=name.line, col_no=name.column),
                                        )
                                    )
                                    tests_cache.insert_test(
                                        file_path=str(test_file),
                                        file_hash=file_hash,
                                        qualified_name_with_modules_from_root=qualified_name_with_modules,
                                        function_name=scope,
                                        test_class=test_func.test_class or "",
                                        test_function=scope_test_function,
                                        test_type=test_func.test_type,
                                        line_number=name.line,
                                        col_number=name.column,
                                    )

                                    if test_func.test_type == TestType.REPLAY_TEST:
                                        num_discovered_replay_tests += 1

                                    num_discovered_tests += 1
                        continue
                    definition_obj = definition[0]
                    definition_path = str(definition_obj.module_path)

                    project_root_str = str(project_root_path)
                    if (
                        definition_path.startswith(project_root_str + os.sep)
                        and definition_obj.module_name != name.module_name
                        and definition_obj.full_name is not None
                    ):
                        # Pre-compute common values outside the inner loop
                        module_prefix = definition_obj.module_name + "."
                        full_name_without_module_prefix = definition_obj.full_name.replace(module_prefix, "", 1)
                        qualified_name_with_modules_from_root = f"{module_name_from_file_path(definition_obj.module_path, project_root_path)}.{full_name_without_module_prefix}"

                        for test_func in test_functions_by_name[scope]:
                            if test_func.parameters is not None:
                                if test_framework == "pytest":
                                    scope_test_function = f"{test_func.function_name}[{test_func.parameters}]"
                                else:  # unittest
                                    scope_test_function = f"{test_func.function_name}_{test_func.parameters}"
                            else:
                                scope_test_function = test_func.function_name

                            function_to_test_map[qualified_name_with_modules_from_root].add(
                                FunctionCalledInTest(
                                    tests_in_file=TestsInFile(
                                        test_file=test_file,
                                        test_class=test_func.test_class,
                                        test_function=scope_test_function,
                                        test_type=test_func.test_type,
                                    ),
                                    position=CodePosition(line_no=name.line, col_no=name.column),
                                )
                            )
                            tests_cache.insert_test(
                                file_path=str(test_file),
                                file_hash=file_hash,
                                qualified_name_with_modules_from_root=qualified_name_with_modules_from_root,
                                function_name=scope,
                                test_class=test_func.test_class or "",
                                test_function=scope_test_function,
                                test_type=test_func.test_type,
                                line_number=name.line,
                                col_number=name.column,
                            )

                            if test_func.test_type == TestType.REPLAY_TEST:
                                num_discovered_replay_tests += 1

                            num_discovered_tests += 1
                except Exception as e:
                    logger.debug(str(e))
                    continue

            tests_cache.flush()
            progress.advance(task_id)

    tests_cache.close()

    return dict(function_to_test_map), num_discovered_tests, num_discovered_replay_tests


# Cache module name resolution to avoid repeated Path.resolve()/relative_to() calls
@lru_cache(maxsize=128)
def module_name_from_file_path_cached(file_path: Path, project_root: Path) -> str:
    return module_name_from_file_path(file_path, project_root)
