from __future__ import annotations

import enum
import re
import sys
from collections import Counter, defaultdict
from collections.abc import Collection
from enum import Enum, IntEnum
from functools import lru_cache
from pathlib import Path
from re import Pattern
from typing import TYPE_CHECKING, Any, NamedTuple, Optional, cast

from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, ValidationError, model_validator
from pydantic.dataclasses import dataclass

from codeflash.models.test_type import TestType

if TYPE_CHECKING:
    from collections.abc import Iterator

    import libcst as cst
    from rich.tree import Tree


@dataclass(frozen=True)
class AIServiceRefinerRequest:
    """Request model for code refinement API.

    Supports multi-language optimization refinement with optional multi-file context.
    """

    optimization_id: str
    original_source_code: str
    read_only_dependency_code: str
    original_code_runtime: int
    optimized_source_code: str
    optimized_explanation: str
    optimized_code_runtime: int
    speedup: str
    trace_id: str
    original_line_profiler_results: str
    optimized_line_profiler_results: str
    function_references: str | None = None
    call_sequence: int | None = None
    # Multi-language support
    language: str = "python"  # 'python', 'javascript', 'typescript'
    language_version: str | None = None  # e.g., '3.11.0' for Python, 'ES2022' for JS
    # Multi-file context support
    additional_context_files: dict[str, str] | None = None  # {filepath: content} for imported modules


# this should be possible to auto serialize
@dataclass(frozen=True)
class AdaptiveOptimizedCandidate:
    optimization_id: str
    source_code: str
    # TODO: introduce repair explanation for code repair candidates to help the llm understand the full process
    explanation: str
    source: OptimizedCandidateSource
    speedup: str


@dataclass(frozen=True)
class AIServiceAdaptiveOptimizeRequest:
    trace_id: str
    original_source_code: str
    candidates: list[AdaptiveOptimizedCandidate]


class TestDiffScope(str, Enum):
    RETURN_VALUE = "return_value"
    STDOUT = "stdout"
    DID_PASS = "did_pass"  # noqa: S105


@dataclass
class TestDiff:
    scope: TestDiffScope
    original_pass: bool
    candidate_pass: bool

    original_value: str | None = None
    candidate_value: str | None = None
    test_src_code: Optional[str] = None
    candidate_pytest_error: Optional[str] = None
    original_pytest_error: Optional[str] = None


@dataclass(frozen=True)
class AIServiceCodeRepairRequest:
    optimization_id: str
    original_source_code: str
    modified_source_code: str
    trace_id: str
    test_diffs: list[TestDiff]
    language: str = "python"


class OptimizationReviewResult(NamedTuple):
    """Result from the optimization review API."""

    review: str  # "high", "medium", "low", or ""
    explanation: str


class FunctionRepairInfo(NamedTuple):
    function_name: str
    reason: str


class TestFileReview(NamedTuple):
    test_index: int
    functions_to_repair: list[FunctionRepairInfo]


# If the method spam is in the class Ham, which is at the top level of the module eggs in the package foo, the fully
# qualified name of the method is foo.eggs.Ham.spam, its qualified name is Ham.spam, and its name is spam. The full name
# of the module is foo.eggs.


class ValidCode(BaseModel):
    model_config = ConfigDict(frozen=True)

    source_code: str
    normalized_code: str


# TODO COVER FIX
class CoverReturnCode(IntEnum):
    DID_NOT_RUN = -1
    NO_DIFFERENCES = 0
    COUNTER_EXAMPLES = 1
    ERROR = 2


@dataclass(frozen=True)
class FunctionSource:
    file_path: Path
    qualified_name: str
    fully_qualified_name: str
    only_function_name: str
    source_code: str
    definition_type: str | None = None  # e.g. "function", "class"; None for non-Python languages

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, FunctionSource):
            return False
        return (
            self.file_path == other.file_path
            and self.qualified_name == other.qualified_name
            and self.fully_qualified_name == other.fully_qualified_name
            and self.only_function_name == other.only_function_name
            and self.source_code == other.source_code
        )

    def __hash__(self) -> int:
        return hash(
            (self.file_path, self.qualified_name, self.fully_qualified_name, self.only_function_name, self.source_code)
        )


class BestOptimization(BaseModel):
    candidate: OptimizedCandidate
    explanation_v2: Optional[str] = None
    helper_functions: list[FunctionSource]
    code_context: CodeOptimizationContext
    runtime: int
    replay_performance_gain: Optional[dict[BenchmarkKey, float]] = None
    winning_behavior_test_results: TestResults
    winning_benchmarking_test_results: TestResults
    winning_replay_benchmarking_test_results: Optional[TestResults] = None
    line_profiler_test_results: dict[Any, Any]
    async_throughput: Optional[int] = None
    concurrency_metrics: Optional[ConcurrencyMetrics] = None


@dataclass(frozen=True)
class BenchmarkKey:
    module_path: str
    function_name: str

    def __str__(self) -> str:
        return f"{self.module_path}::{self.function_name}"


@dataclass
class ConcurrencyMetrics:
    sequential_time_ns: int
    concurrent_time_ns: int
    concurrency_factor: int
    concurrency_ratio: float  # sequential_time / concurrent_time


@dataclass
class BenchmarkDetail:
    benchmark_name: str
    test_function: str
    original_timing: str
    expected_new_timing: str
    speedup_percent: float

    def to_string(self) -> str:
        return (
            f"Original timing for {self.benchmark_name}::{self.test_function}: {self.original_timing}\n"
            f"Expected new timing for {self.benchmark_name}::{self.test_function}: {self.expected_new_timing}\n"
            f"Benchmark speedup for {self.benchmark_name}::{self.test_function}: {self.speedup_percent:.2f}%\n"
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "benchmark_name": self.benchmark_name,
            "test_function": self.test_function,
            "original_timing": self.original_timing,
            "expected_new_timing": self.expected_new_timing,
            "speedup_percent": self.speedup_percent,
        }


@dataclass
class ProcessedBenchmarkInfo:
    benchmark_details: list[BenchmarkDetail]

    def to_string(self) -> str:
        if not self.benchmark_details:
            return ""

        result = "Benchmark Performance Details:\n"
        for detail in self.benchmark_details:
            result += detail.to_string() + "\n"
        return result

    def to_dict(self) -> dict[str, list[dict[str, Any]]]:
        return {"benchmark_details": [detail.to_dict() for detail in self.benchmark_details]}


class CodeString(BaseModel):
    code: str
    file_path: Optional[Path] = None
    language: str = "python"  # Language for validation

    @model_validator(mode="after")
    def validate_code_syntax(self) -> CodeString:
        """Validate code syntax for the specified language."""
        if self.language == "python":
            from codeflash.code_utils.code_utils import validate_python_code

            validate_python_code(self.code)
        else:
            from codeflash.languages.registry import get_language_support

            lang_support = get_language_support(self.language)
            if not lang_support.validate_syntax(self.code, file_path=self.file_path):
                msg = f"Invalid {self.language.title()} code"
                raise ValueError(msg)
        return self


def get_comment_prefix(file_path: Path) -> str:
    """Get the comment prefix for a given language."""
    from codeflash.languages.registry import get_language_support

    support = get_language_support(file_path)
    return support.comment_prefix


def get_code_block_splitter(file_path: Path | None) -> str:
    if file_path is None:
        return ""
    comment_prefix = get_comment_prefix(file_path)
    return f"{comment_prefix} file: {file_path.as_posix()}"


# Pattern to match markdown code blocks with optional language tag and file path
# Matches: ```language:filepath\ncode\n``` or ```language\ncode\n```
markdown_pattern = re.compile(r"```(\w+)(?::([^\n]+))?\n(.*?)\n```", re.DOTALL)


class CodeStringsMarkdown(BaseModel):
    code_strings: list[CodeString] = []
    language: str = "python"  # Language for markdown code block tags
    _cache: dict[str, Any] = PrivateAttr(default_factory=dict)

    @property
    def flat(self) -> str:
        """Returns the combined source code module from all code blocks.

        Each block is prefixed by a file path comment to indicate its origin.
        The comment prefix is determined by the language attribute.

        Returns:
            str: The concatenated code of all blocks with file path annotations.

        !! Important !!:
        Avoid parsing the flat code with multiple files,
        parsing may result in unexpected behavior.


        """
        if self._cache.get("flat") is not None:
            return self._cache["flat"]
        self._cache["flat"] = "\n".join(
            get_code_block_splitter(block.file_path) + "\n" + block.code for block in self.code_strings
        )
        return self._cache["flat"]

    @property
    def markdown(self) -> str:
        """Returns a Markdown-formatted string containing all code blocks.

        Each block is enclosed in a triple-backtick code block with an optional
        file path suffix (e.g., ```python:filename.py or ```javascript:file.js).

        The language tag is determined by the `language` attribute.

        Returns:
            str: Markdown representation of the code blocks.

        """
        return "\n".join(
            [
                f"```{self.language}{':' + code_string.file_path.as_posix() if code_string.file_path else ''}\n{code_string.code.strip()}\n```"
                for code_string in self.code_strings
            ]
        )

    def file_to_path(self) -> dict[str, str]:
        """Return a dictionary mapping file paths to their corresponding code blocks.

        Returns:
            dict[str, str]: Mapping from file path (as string) to code.

        """
        try:
            return self._cache["file_to_path"]
        except KeyError:
            mapping = {str(code_string.file_path): code_string.code for code_string in self.code_strings}
            self._cache["file_to_path"] = mapping
            return mapping

    @staticmethod
    def parse_markdown_code(markdown_code: str, expected_language: str = "python") -> CodeStringsMarkdown:
        """Parse a Markdown string into a CodeStringsMarkdown object.

        Extracts code blocks and their associated file paths and constructs a new CodeStringsMarkdown instance.

        Args:
            markdown_code (str): The Markdown-formatted string to parse.
            expected_language (str): The expected language of code blocks (default: "python").

        Returns:
            CodeStringsMarkdown: Parsed object containing code blocks.

        """
        matches = markdown_pattern.findall(markdown_code)
        code_string_list = []
        detected_language = expected_language
        try:
            for language, file_path, code in matches:
                # Use the first detected language or the expected language
                if language:
                    detected_language = language
                if file_path:
                    path = file_path.strip()
                    code_string_list.append(CodeString(code=code, file_path=Path(path), language=detected_language))
                else:
                    # No file path specified - skip this block or create with None
                    code_string_list.append(CodeString(code=code, file_path=None, language=detected_language))
            return CodeStringsMarkdown(code_strings=code_string_list, language=detected_language)
        except ValidationError:
            # if any file is invalid, return an empty CodeStringsMarkdown for the entire context
            return CodeStringsMarkdown(language=expected_language)


class CodeOptimizationContext(BaseModel):
    testgen_context: CodeStringsMarkdown
    read_writable_code: CodeStringsMarkdown
    read_only_context_code: str = ""
    hashing_code_context: str = ""
    hashing_code_context_hash: str = ""
    helper_functions: list[FunctionSource]
    testgen_helper_fqns: list[str] = []
    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]]


class CodeContextType(str, Enum):
    READ_WRITABLE = "READ_WRITABLE"
    READ_ONLY = "READ_ONLY"
    TESTGEN = "TESTGEN"
    HASHING = "HASHING"


class OptimizedCandidateResult(BaseModel):
    max_loop_count: int
    best_test_runtime: int
    behavior_test_results: TestResults
    benchmarking_test_results: TestResults
    replay_benchmarking_test_results: Optional[dict[BenchmarkKey, TestResults]] = None
    optimization_candidate_index: int
    total_candidate_timing: int
    async_throughput: Optional[int] = None
    concurrency_metrics: Optional[ConcurrencyMetrics] = None


class GeneratedTests(BaseModel):
    generated_original_test_source: str
    instrumented_behavior_test_source: str
    instrumented_perf_test_source: str
    raw_generated_test_source: str | None = None
    behavior_file_path: Path
    perf_file_path: Path


class GeneratedTestsList(BaseModel):
    generated_tests: list[GeneratedTests]


class TestFile(BaseModel):
    instrumented_behavior_file_path: Path
    benchmarking_file_path: Optional[Path] = None
    original_file_path: Optional[Path] = None
    original_source: Optional[str] = None
    test_type: TestType
    tests_in_file: Optional[list[TestsInFile]] = None


class TestFiles(BaseModel):
    test_files: list[TestFile]

    def add(self, test_file: TestFile) -> None:
        if test_file not in self.test_files:
            self.test_files.append(test_file)
        else:
            msg = "Test file already exists in the list"
            raise ValueError(msg)

    def get_by_original_file_path(self, file_path: Path) -> TestFile | None:
        normalized = self._normalize_path_for_comparison(file_path)
        for test_file in self.test_files:
            if test_file.original_file_path is None:
                continue
            normalized_test_path = self._normalize_path_for_comparison(test_file.original_file_path)
            if normalized == normalized_test_path:
                return test_file
        return None

    def get_test_type_by_instrumented_file_path(self, file_path: Path) -> TestType | None:
        normalized = self._normalize_path_for_comparison(file_path)
        for test_file in self.test_files:
            normalized_behavior_path = self._normalize_path_for_comparison(test_file.instrumented_behavior_file_path)
            if normalized == normalized_behavior_path:
                return test_file.test_type
            if test_file.benchmarking_file_path is not None:
                normalized_benchmark_path = self._normalize_path_for_comparison(test_file.benchmarking_file_path)
                if normalized == normalized_benchmark_path:
                    return test_file.test_type

        # Fallback: try filename-only matching for JavaScript/TypeScript
        # Jest/Vitest JUnit XML may have relative paths that don't match absolute paths
        file_name = file_path.name
        for test_file in self.test_files:
            if (
                test_file.instrumented_behavior_file_path
                and test_file.instrumented_behavior_file_path.name == file_name
            ):
                return test_file.test_type
            if test_file.benchmarking_file_path and test_file.benchmarking_file_path.name == file_name:
                return test_file.test_type

        return None

    def get_test_type_by_original_file_path(self, file_path: Path) -> TestType | None:
        normalized = self._normalize_path_for_comparison(file_path)
        for test_file in self.test_files:
            if test_file.original_file_path is None:
                continue
            normalized_test_path = self._normalize_path_for_comparison(test_file.original_file_path)
            if normalized == normalized_test_path:
                return test_file.test_type
        return None

    @staticmethod
    @lru_cache(maxsize=4096)
    def _normalize_path_for_comparison(path: Path) -> str:
        """Normalize a path for cross-platform comparison.

        Resolves the path to an absolute path and handles Windows case-insensitivity.
        """
        try:
            resolved = str(path.resolve())
        except (OSError, RuntimeError):
            # If resolve fails (e.g., file doesn't exist), use absolute path
            resolved = str(path.absolute())
        # Only lowercase on Windows where filesystem is case-insensitive
        return resolved.lower() if sys.platform == "win32" else resolved

    def __iter__(self) -> Iterator[TestFile]:
        return iter(self.test_files)

    def __len__(self) -> int:
        return len(self.test_files)


class OptimizationSet(BaseModel):
    control: list[OptimizedCandidate]
    experiment: Optional[list[OptimizedCandidate]]


@dataclass
class CandidateEvaluationContext:
    """Holds tracking state during candidate evaluation in determine_best_candidate."""

    speedup_ratios: dict[str, float | None] = Field(default_factory=dict)
    optimized_runtimes: dict[str, float | None] = Field(default_factory=dict)
    is_correct: dict[str, bool] = Field(default_factory=dict)
    optimized_line_profiler_results: dict[str, str] = Field(default_factory=dict)
    ast_code_to_id: dict = Field(default_factory=dict)
    optimizations_post: dict[str, str] = Field(default_factory=dict)
    valid_optimizations: list = Field(default_factory=list)

    def record_failed_candidate(self, optimization_id: str) -> None:
        """Record results for a failed candidate."""
        self.optimized_runtimes[optimization_id] = None
        self.is_correct[optimization_id] = False
        self.speedup_ratios[optimization_id] = None

    def record_successful_candidate(self, optimization_id: str, runtime: float, speedup: float) -> None:
        """Record results for a successful candidate."""
        self.optimized_runtimes[optimization_id] = runtime
        self.is_correct[optimization_id] = True
        self.speedup_ratios[optimization_id] = speedup

    def record_line_profiler_result(self, optimization_id: str, result: str) -> None:
        """Record line profiler results for a candidate."""
        self.optimized_line_profiler_results[optimization_id] = result

    def handle_duplicate_candidate(
        self, candidate: OptimizedCandidate, normalized_code: str, original_flat_code: str
    ) -> None:
        """Handle a candidate that has been seen before."""
        past_opt_id = self.ast_code_to_id[normalized_code]["optimization_id"]

        # Copy results from the previous evaluation (use .get() in case past_opt_id was registered
        # but never benchmarked due to an unhandled exception in process_single_candidate)
        self.speedup_ratios[candidate.optimization_id] = self.speedup_ratios.get(past_opt_id)
        self.is_correct[candidate.optimization_id] = self.is_correct.get(past_opt_id)
        self.optimized_runtimes[candidate.optimization_id] = self.optimized_runtimes.get(past_opt_id)

        # Line profiler results only available for successful runs
        if past_opt_id in self.optimized_line_profiler_results:
            self.optimized_line_profiler_results[candidate.optimization_id] = self.optimized_line_profiler_results[
                past_opt_id
            ]

        self.optimizations_post[candidate.optimization_id] = self.ast_code_to_id[normalized_code][
            "shorter_source_code"
        ].markdown
        self.optimizations_post[past_opt_id] = self.ast_code_to_id[normalized_code]["shorter_source_code"].markdown

        # Update to shorter code if this candidate has a shorter diff
        from codeflash.code_utils.code_utils import diff_length

        new_diff_len = diff_length(candidate.source_code.flat, original_flat_code)
        if new_diff_len < self.ast_code_to_id[normalized_code]["diff_len"]:
            self.ast_code_to_id[normalized_code]["shorter_source_code"] = candidate.source_code
            self.ast_code_to_id[normalized_code]["diff_len"] = new_diff_len

    def register_new_candidate(
        self, normalized_code: str, candidate: OptimizedCandidate, original_flat_code: str
    ) -> None:
        """Register a new candidate that hasn't been seen before."""
        from codeflash.code_utils.code_utils import diff_length

        self.ast_code_to_id[normalized_code] = {
            "optimization_id": candidate.optimization_id,
            "shorter_source_code": candidate.source_code,
            "diff_len": diff_length(candidate.source_code.flat, original_flat_code),
        }

    def get_speedup_ratio(self, optimization_id: str) -> float | None:
        return self.speedup_ratios.get(optimization_id)

    def get_optimized_runtime(self, optimization_id: str) -> float | None:
        return self.optimized_runtimes.get(optimization_id)


@dataclass(frozen=True)
class TestsInFile:
    test_file: Path
    test_class: Optional[str]
    test_function: str
    test_type: TestType


class OptimizedCandidateSource(str, Enum):
    OPTIMIZE = "OPTIMIZE"
    OPTIMIZE_LP = "OPTIMIZE_LP"
    REFINE = "REFINE"
    REPAIR = "REPAIR"
    ADAPTIVE = "ADAPTIVE"
    JIT_REWRITE = "JIT_REWRITE"


@dataclass(frozen=True)
class OptimizedCandidate:
    source_code: CodeStringsMarkdown
    explanation: str
    optimization_id: str
    source: OptimizedCandidateSource
    parent_id: str | None = None
    model: str | None = None  # Which LLM model generated this candidate


@dataclass(frozen=True)
class FunctionCalledInTest:
    tests_in_file: TestsInFile
    position: CodePosition


@dataclass(frozen=True)
class CodePosition:
    line_no: int
    col_no: int


# Re-export FunctionParent for backward compatibility
from codeflash.models.function_types import FunctionParent  # noqa: E402


class OriginalCodeBaseline(BaseModel):
    behavior_test_results: TestResults
    benchmarking_test_results: TestResults
    replay_benchmarking_test_results: Optional[dict[BenchmarkKey, TestResults]] = None
    line_profile_results: dict
    runtime: int
    coverage_results: Optional[CoverageData]
    async_throughput: Optional[int] = None
    concurrency_metrics: Optional[ConcurrencyMetrics] = None


class CoverageStatus(Enum):
    NOT_FOUND = "Coverage Data Not Found"
    PARSED_SUCCESSFULLY = "Parsed Successfully"


@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
class CoverageData:
    """Represents the coverage data for a specific function in a source file, using one or more test files."""

    file_path: Path
    coverage: float
    function_name: str
    functions_being_tested: list[str]
    graph: dict[str, dict[str, Collection[object]]]
    code_context: CodeOptimizationContext
    main_func_coverage: FunctionCoverage
    dependent_func_coverage: Optional[FunctionCoverage]
    status: CoverageStatus
    blank_re: Pattern[str] = re.compile(r"\s*(#|$)")
    else_re: Pattern[str] = re.compile(r"\s*else\s*:\s*(#|$)")

    def build_message(self) -> str:
        if self.status == CoverageStatus.NOT_FOUND:
            return f"No coverage data found for {self.function_name}"
        return f"{self.coverage:.1f}%"

    def log_coverage(self) -> None:
        from rich.tree import Tree

        from codeflash.cli_cmds.console import console, logger
        from codeflash.code_utils.env_utils import is_end_to_end

        tree = Tree("Test Coverage Results")
        tree.add(f"Main Function: {self.main_func_coverage.name}: {self.coverage:.2f}%")
        if self.dependent_func_coverage:
            tree.add(
                f"Dependent Function: {self.dependent_func_coverage.name}: {self.dependent_func_coverage.coverage:.2f}%"
            )
        tree.add(f"Total Coverage: {self.coverage:.2f}%")
        console.print(tree)
        console.rule()

        if not self.coverage:
            logger.debug(self.graph)
        if is_end_to_end():
            console.print(self)

    @classmethod
    def create_empty(cls, file_path: Path, function_name: str, code_context: CodeOptimizationContext) -> CoverageData:
        return cls(
            file_path=file_path,
            coverage=0.0,
            function_name=function_name,
            functions_being_tested=[function_name],
            graph={
                function_name: {
                    "executed_lines": set(),
                    "unexecuted_lines": set(),
                    "executed_branches": [],
                    "unexecuted_branches": [],
                }
            },
            code_context=code_context,
            main_func_coverage=FunctionCoverage(
                name=function_name,
                coverage=0.0,
                executed_lines=[],
                unexecuted_lines=[],
                executed_branches=[],
                unexecuted_branches=[],
            ),
            dependent_func_coverage=None,
            status=CoverageStatus.NOT_FOUND,
        )


@dataclass
class FunctionCoverage:
    """Represents the coverage data for a specific function in a source file."""

    name: str
    coverage: float
    executed_lines: list[int]
    unexecuted_lines: list[int]
    executed_branches: list[list[int]]
    unexecuted_branches: list[list[int]]


class TestingMode(enum.Enum):
    BEHAVIOR = "behavior"
    PERFORMANCE = "performance"
    LINE_PROFILE = "line_profile"
    CONCURRENCY = "concurrency"


# TODO this class is duplicated in codeflash_capture
class VerificationType(str, Enum):
    FUNCTION_CALL = (
        "function_call"  # Correctness verification for a test function, checks input values and output values)
    )
    INIT_STATE_FTO = "init_state_fto"  # Correctness verification for fto class instance attributes after init
    INIT_STATE_HELPER = "init_state_helper"  # Correctness verification for helper class instance attributes after init
    VOID_STATE = "void_state"  # Correctness verification for void methods (no return value)


@dataclass(frozen=True, slots=True)
class InvocationId:
    test_module_path: str  # The fully qualified name of the test module
    test_class_name: Optional[str]  # The name of the class where the test is defined
    test_function_name: Optional[str]  # The name of the test_function. Does not include the components of the file_name
    function_getting_tested: str
    iteration_id: Optional[str]

    # test_module_path:TestSuiteClass.test_function_name:function_tested:iteration_id
    def id(self) -> str:
        class_prefix = f"{self.test_class_name}." if self.test_class_name else ""
        return (
            f"{self.test_module_path}:{class_prefix}{self.test_function_name}:"
            f"{self.function_getting_tested}:{self.iteration_id}"
        )

    # TestSuiteClass.test_function_name
    def test_fn_qualified_name(self) -> str:
        # Use f-string with inline conditional to reduce string concatenation operations
        return (
            f"{self.test_class_name}.{self.test_function_name}"
            if self.test_class_name
            else str(self.test_function_name)
        )

    def find_func_in_class(self, class_node: cst.ClassDef, func_name: str) -> Optional[cst.FunctionDef]:
        import libcst as cst

        for stmt in class_node.body.body:
            if isinstance(stmt, cst.FunctionDef) and stmt.name.value == func_name:
                return stmt
        return None

    def get_src_code(self, test_path: Path) -> Optional[str]:
        import libcst as cst

        if not test_path.exists():
            return None
        try:
            test_src = test_path.read_text(encoding="utf-8")
            module_node = cst.parse_module(test_src)
        except Exception:
            # libcst can't parse non-Python files (JS/TS) — return a descriptive string
            # so the code repair API receives a non-None test_src_code.
            return (
                f"// Test: {self.test_function_name}\n"
                f"// File: {test_path.name}\n"
                f"// Testing function: {self.function_getting_tested}"
            )

        if self.test_class_name:
            for stmt in module_node.body:
                if isinstance(stmt, cst.ClassDef) and stmt.name.value == self.test_class_name:
                    func_node = self.find_func_in_class(stmt, self.test_function_name)
                    if func_node:
                        return module_node.code_for_node(func_node).strip()
            return None

        # Otherwise, look for a top level function
        for stmt in module_node.body:
            if isinstance(stmt, cst.FunctionDef) and stmt.name.value == self.test_function_name:
                return module_node.code_for_node(stmt).strip()
        return None

    @staticmethod
    def from_str_id(string_id: str, iteration_id: str | None = None) -> InvocationId:
        components = string_id.split(":")
        assert len(components) == 4
        second_components = components[1].split(".")
        if len(second_components) == 1:
            test_class_name = None
            test_function_name = second_components[0]
        else:
            test_class_name = second_components[0]
            test_function_name = second_components[1]
        return InvocationId(
            test_module_path=components[0],
            test_class_name=test_class_name,
            test_function_name=test_function_name,
            function_getting_tested=components[2],
            iteration_id=iteration_id if iteration_id else components[3],
        )


@dataclass(frozen=True, slots=True)
class FunctionTestInvocation:
    loop_index: int  # The loop index of the function invocation, starts at 1
    id: InvocationId  # The fully qualified name of the function invocation (id)
    file_name: Path  # The file where the test is defined
    did_pass: bool  # Whether the test this function invocation was part of, passed or failed
    runtime: Optional[int]  # Time in nanoseconds
    test_framework: str  # unittest or pytest
    test_type: TestType
    return_value: Optional[object]  # The return value of the function invocation
    timed_out: Optional[bool]
    verification_type: Optional[str] = VerificationType.FUNCTION_CALL
    stdout: Optional[str] = None

    @property
    def unique_invocation_loop_id(self) -> str:
        return f"{self.loop_index}:{self.id.id()}"


class TestResults(BaseModel):  # noqa: PLW1641
    # don't modify these directly, use the add method
    # also we don't support deletion of test results elements - caution is advised
    test_results: list[FunctionTestInvocation] = []
    test_result_idx: dict[str, int] = {}

    perf_stdout: Optional[str] = None
    # mapping between test function name and stdout failure message
    test_failures: Optional[dict[str, str]] = None

    def add(self, function_test_invocation: FunctionTestInvocation) -> None:
        unique_id = function_test_invocation.unique_invocation_loop_id
        test_result_idx = self.test_result_idx
        if unique_id in test_result_idx:
            from codeflash.cli_cmds.console import DEBUG_MODE, logger

            if DEBUG_MODE:
                logger.warning(f"Test result with id {unique_id} already exists. SKIPPING")
            return
        test_results = self.test_results
        test_result_idx[unique_id] = len(test_results)
        test_results.append(function_test_invocation)

    def merge(self, other: TestResults) -> None:
        original_len = len(self.test_results)
        self.test_results.extend(other.test_results)
        for k, v in other.test_result_idx.items():
            if k in self.test_result_idx:
                msg = f"Test result with id {k} already exists."
                raise ValueError(msg)
            self.test_result_idx[k] = v + original_len

    def group_by_benchmarks(
        self, benchmark_keys: list[BenchmarkKey], benchmark_replay_test_dir: Path, project_root: Path
    ) -> dict[BenchmarkKey, TestResults]:
        """Group TestResults by benchmark for calculating improvements for each benchmark."""
        from codeflash.code_utils.code_utils import module_name_from_file_path

        test_results_by_benchmark = defaultdict(TestResults)
        benchmark_module_path = {}
        for benchmark_key in benchmark_keys:
            benchmark_module_path[benchmark_key] = module_name_from_file_path(
                benchmark_replay_test_dir.resolve()
                / f"test_{benchmark_key.module_path.replace('.', '_')}__replay_test_",
                project_root,
                traverse_up=True,
            )
        for test_result in self.test_results:
            if test_result.test_type == TestType.REPLAY_TEST:
                for benchmark_key, module_path in benchmark_module_path.items():
                    if test_result.id.test_module_path.startswith(module_path):
                        test_results_by_benchmark[benchmark_key].add(test_result)

        return test_results_by_benchmark

    def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None:
        try:
            return self.test_results[self.test_result_idx[unique_invocation_loop_id]]
        except (IndexError, KeyError):
            return None

    def get_all_unique_invocation_loop_ids(self) -> set[str]:
        return {test_result.unique_invocation_loop_id for test_result in self.test_results}

    def number_of_loops(self) -> int:
        if not self.test_results:
            return 0
        return max(test_result.loop_index for test_result in self.test_results)

    def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]:
        report: dict[TestType, dict[str, int]] = {tt: {"passed": 0, "failed": 0} for tt in TestType}
        for test_result in self.test_results:
            if test_result.loop_index != 1:
                continue
            if test_result.did_pass:
                report[test_result.test_type]["passed"] += 1
            else:
                report[test_result.test_type]["failed"] += 1
        return report

    @staticmethod
    def report_to_string(report: dict[TestType, dict[str, int]]) -> str:
        return " ".join(
            [
                f"{test_type.to_name()}- (Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']})"
                for test_type in TestType
            ]
        )

    @staticmethod
    def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree:
        from rich.tree import Tree

        from codeflash.lsp.helpers import is_LSP_enabled

        tree = Tree(title)

        if is_LSP_enabled():
            from codeflash.cli_cmds.console import lsp_log
            from codeflash.lsp.helpers import report_to_markdown_table
            from codeflash.lsp.lsp_message import LspMarkdownMessage

            # Build markdown table
            markdown = report_to_markdown_table(report, title)
            lsp_log(LspMarkdownMessage(markdown=markdown))
            return tree

        for test_type in TestType:
            if test_type is TestType.INIT_STATE_TEST:
                continue
            tree.add(
                f"{test_type.to_name()} - Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']}"
            )
        return tree

    def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]:
        from codeflash.cli_cmds.console import logger

        # Efficient single traversal, directly accumulating into a dict.
        # can track mins here and only sums can be return in total_passed_runtime
        by_id: dict[InvocationId, list[int]] = {}
        for result in self.test_results:
            if result.did_pass:
                if result.runtime is not None:
                    by_id.setdefault(result.id, []).append(result.runtime)
                else:
                    msg = (
                        f"Ignoring test case that passed but had no runtime -> {result.id}, "
                        f"Loop # {result.loop_index}, Test Type: {result.test_type}, "
                        f"Verification Type: {result.verification_type}"
                    )
                    logger.debug(msg)
        return by_id

    def total_passed_runtime(self) -> int:
        """Calculate the sum of runtimes of all test cases that passed.

        A testcase runtime is the minimum value of all looped execution runtimes.

        :return: The runtime in nanoseconds.
        """
        # TODO this doesn't look at the intersection of tests of baseline and original
        return sum(
            [min(usable_runtime_data) for _, usable_runtime_data in self.usable_runtime_data_by_test_case().items()]
        )

    def effective_loop_count(self) -> int:
        """Calculate the effective number of complete loops.

        For consistent behavior across Python and JavaScript tests, this returns
        the maximum loop_index seen across all test results. This represents
        the number of timing iterations that were performed.

        Note: For JavaScript tests without the loop-runner, each test case may have
        different iteration counts due to internal looping in capturePerf. We use
        max() to report the highest iteration count achieved.

        :return: The effective loop count, or 0 if no test results.
        """
        if not self.test_results:
            return 0
        # Get all loop indices from results that have timing data
        loop_indices = {result.loop_index for result in self.test_results if result.runtime is not None}
        if not loop_indices:
            # Fallback: use all loop indices even without runtime
            loop_indices = {result.loop_index for result in self.test_results}
        return max(loop_indices) if loop_indices else 0

    def file_to_no_of_tests(self, test_functions_to_remove: list[str]) -> Counter[Path]:
        map_gen_test_file_to_no_of_tests = Counter()
        for gen_test_result in self.test_results:
            if (
                gen_test_result.test_type == TestType.GENERATED_REGRESSION
                and gen_test_result.id.test_function_name not in test_functions_to_remove
            ):
                map_gen_test_file_to_no_of_tests[gen_test_result.file_name] += 1
        return map_gen_test_file_to_no_of_tests

    def __iter__(self) -> Iterator[FunctionTestInvocation]:
        return iter(self.test_results)

    def __len__(self) -> int:
        return len(self.test_results)

    def __getitem__(self, index: int) -> FunctionTestInvocation:
        return self.test_results[index]

    def __setitem__(self, index: int, value: FunctionTestInvocation) -> None:
        self.test_results[index] = value

    def __contains__(self, value: FunctionTestInvocation) -> bool:
        return value in self.test_results

    def __bool__(self) -> bool:
        return bool(self.test_results)

    def __eq__(self, other: object) -> bool:
        from codeflash.verification.comparator import comparator

        # Unordered comparison
        if type(self) is not type(other):
            return False
        if len(self) != len(other):
            return False
        original_recursion_limit = sys.getrecursionlimit()
        cast("TestResults", other)
        for test_result in self:
            other_test_result = other.get_by_unique_invocation_loop_id(test_result.unique_invocation_loop_id)
            if other_test_result is None:
                return False

            if original_recursion_limit < 5000:
                sys.setrecursionlimit(5000)
            if (
                test_result.file_name != other_test_result.file_name
                or test_result.did_pass != other_test_result.did_pass
                or test_result.runtime != other_test_result.runtime
                or test_result.test_framework != other_test_result.test_framework
                or test_result.test_type != other_test_result.test_type
                or not comparator(test_result.return_value, other_test_result.return_value)
            ):
                sys.setrecursionlimit(original_recursion_limit)
                return False
        sys.setrecursionlimit(original_recursion_limit)
        return True
