"""Benchmark instrumentation, validation, replay test generation, and comparison.

Combines decorator injection (libcst), error extraction from pytest
output, trace-based replay test generation, result validation/formatting,
and cross-branch comparison into one module.
"""

from __future__ import annotations

import logging
import re
import sqlite3
import textwrap
from pathlib import Path
from typing import TYPE_CHECKING, Any

import attrs
import libcst as cst

from codeflash_core import BenchmarkDetail, humanize_runtime, performance_gain

from ..analysis._discovery import inspect_top_level_functions_or_methods
from ..analysis._formatter import sort_imports
from .models import (
    ProcessedBenchmarkInfo,
    get_function_alias,
    get_unique_test_name,
)

log = logging.getLogger(__name__)

if TYPE_CHECKING:
    from collections.abc import Generator

    from .._model import FunctionToOptimize
    from .models import BenchmarkKey


class AddDecoratorTransformer(cst.CSTTransformer):
    """Add ``@codeflash_trace`` decorator to target functions."""

    def __init__(self, target_functions: set[tuple[str, str]]) -> None:
        """Initialize with *(class_name, function_name)* pairs."""
        super().__init__()
        self.target_functions = target_functions
        self.added_codeflash_trace = False
        self.class_name = ""
        self.function_name = ""
        self.decorator = cst.Decorator(
            decorator=cst.Name(value="codeflash_trace")
        )

    def leave_ClassDef(
        self,
        original_node: cst.ClassDef,
        updated_node: cst.ClassDef,
    ) -> cst.ClassDef:
        """Reset the class scope when leaving a class."""
        if self.class_name == original_node.name.value:
            self.class_name = ""
        return updated_node

    def visit_ClassDef(self, node: cst.ClassDef) -> bool | None:
        """Track class scope; skip nested classes."""
        if self.class_name:
            return False
        self.class_name = node.name.value
        return None

    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool | None:
        """Track function scope; skip nested functions."""
        if self.function_name:
            return False
        self.function_name = node.name.value
        return None

    def leave_FunctionDef(
        self,
        original_node: cst.FunctionDef,
        updated_node: cst.FunctionDef,
    ) -> cst.FunctionDef:
        """Attach ``@codeflash_trace`` to matched functions."""
        if self.function_name == original_node.name.value:
            self.function_name = ""
        if (
            self.class_name,
            original_node.name.value,
        ) in self.target_functions:
            updated_decorators = [
                *list(updated_node.decorators),
                self.decorator,
            ]
            self.added_codeflash_trace = True
            return updated_node.with_changes(decorators=updated_decorators)
        return updated_node

    def leave_Module(
        self,
        original_node: cst.Module,
        updated_node: cst.Module,
    ) -> cst.Module:
        """Insert the ``codeflash_trace`` import if decorators were added."""
        if not self.added_codeflash_trace:
            return updated_node
        import_stmt = cst.SimpleStatementLine(
            body=[
                cst.ImportFrom(
                    module=cst.Attribute(
                        value=cst.Attribute(
                            value=cst.Name(value="codeflash_python"),
                            attr=cst.Name(value="benchmarking"),
                        ),
                        attr=cst.Name(value="_benchmark_tracing"),
                    ),
                    names=[
                        cst.ImportAlias(name=cst.Name(value="codeflash_trace"))
                    ],
                )
            ]
        )
        new_body = [import_stmt, *list(updated_node.body)]
        return updated_node.with_changes(body=new_body)


def add_codeflash_decorator_to_code(
    code: str,
    functions_to_optimize: list[FunctionToOptimize],
) -> str:
    """Return *code* with ``@codeflash_trace`` added to the target functions."""
    target_functions: set[tuple[str, str]] = set()
    for fto in functions_to_optimize:
        class_name = ""
        if len(fto.parents) == 1 and fto.parents[0].type == "ClassDef":
            class_name = fto.parents[0].name
        target_functions.add((class_name, fto.function_name))
    transformer = AddDecoratorTransformer(target_functions=target_functions)
    module = cst.parse_module(code)
    modified_module = module.visit(transformer)
    return modified_module.code


def instrument_codeflash_trace_decorator(
    file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]],
) -> None:
    """Instrument source files with ``@codeflash_trace`` decorators in-place."""
    for (
        file_path,
        functions_to_optimize,
    ) in file_to_funcs_to_optimize.items():
        # Skip codeflash's own benchmarking and picklepatch modules
        # to avoid circular imports.
        _, sep, after = file_path.as_posix().rpartition("/codeflash/")
        if sep:
            submodule = after.partition("/")[0]
            if submodule in ("benchmarking", "picklepatch"):
                continue
        original_code = file_path.read_text(encoding="utf-8")
        new_code = add_codeflash_decorator_to_code(
            original_code, functions_to_optimize
        )
        modified_code = sort_imports(new_code)
        file_path.write_text(modified_code, encoding="utf-8")


def extract_benchmark_errors(output: str) -> str:
    """Extract error sections from pytest benchmark output."""
    if "ERROR collecting" in output:
        error_pattern = (
            r"={3,}\s*ERRORS\s*={3,}\n"
            r"([\s\S]*?)(?:={3,}|$)"
        )
        match = re.search(error_pattern, output)
        return match.group(1) if match else output
    if "FAILURES" in output:
        error_pattern = (
            r"={3,}\s*FAILURES\s*={3,}\n"
            r"([\s\S]*?)(?:={3,}|$)"
        )
        match = re.search(error_pattern, output)
        return match.group(1) if match else output
    return output


def get_next_arg_and_return(
    trace_file: str,
    benchmark_function_name: str,
    function_name: str,
    file_path: str,
    class_name: str | None = None,
    num_to_get: int = 256,
) -> Generator[tuple[Any, Any]]:
    """Yield ``(pickled_args, pickled_kwargs)`` rows from a trace database."""
    db = sqlite3.connect(trace_file)
    cur = db.cursor()
    normalized_file_path = Path(file_path).as_posix()
    if class_name is not None:
        cursor = cur.execute(
            "SELECT * FROM benchmark_function_timings"
            " WHERE benchmark_function_name = ?"
            " AND function_name = ?"
            " AND file_path = ?"
            " AND class_name = ? LIMIT ?",
            (
                benchmark_function_name,
                function_name,
                normalized_file_path,
                class_name,
                num_to_get,
            ),
        )
    else:
        cursor = cur.execute(
            "SELECT * FROM benchmark_function_timings"
            " WHERE benchmark_function_name = ?"
            " AND function_name = ?"
            " AND file_path = ?"
            " AND class_name = '' LIMIT ?",
            (
                benchmark_function_name,
                function_name,
                normalized_file_path,
                num_to_get,
            ),
        )
    try:
        while (val := cursor.fetchone()) is not None:
            yield val[9], val[10]
    finally:
        db.close()


def create_trace_replay_test_code(
    trace_file: str,
    functions_data: list[dict[str, Any]],
    max_run_count: int = 256,
) -> str:
    """Create a replay test for functions based on trace data."""
    imports = (
        "from codeflash_python.runtime._picklepatch.pickle_patcher"
        " import PicklePatcher as pickle\n"
        "from codeflash_python.benchmarking._benchmarking"
        " import get_next_arg_and_return\n"
    )

    function_imports = []
    for func in functions_data:
        module_name: str = func["module_name"]
        function_name: str = func["function_name"]
        class_name: str = func.get("class_name", "")
        if class_name:
            function_imports.append(
                f"from {module_name} import {class_name}"
                f" as {get_function_alias(module_name, class_name)}"
            )
        else:
            function_imports.append(
                f"from {module_name} import {function_name}"
                f" as {get_function_alias(module_name, function_name)}"
            )

    imports += "\n".join(function_imports)

    functions_to_optimize = sorted(
        {
            func["function_name"]
            for func in functions_data
            if func["function_name"] != "__init__"
        }
    )
    metadata = (
        f"functions = {functions_to_optimize}\n"
        f'trace_file_path = r"{trace_file}"\n'
    )

    test_function_body = textwrap.dedent("""\
        for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", num_to_get={max_run_count}):
            args = pickle.loads(args_pkl)
            kwargs = pickle.loads(kwargs_pkl)
            ret = {function_name}(*args, **kwargs)
            """)

    test_method_body = textwrap.dedent("""\
        for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}):
            args = pickle.loads(args_pkl)
            kwargs = pickle.loads(kwargs_pkl){filter_variables}
            function_name = "{orig_function_name}"
            if not args:
                raise ValueError("No arguments provided for the method.")
            if function_name == "__init__":
                ret = {class_name_alias}(*args[1:], **kwargs)
            else:
                ret = {class_name_alias}{method_name}(*args, **kwargs)
            """)

    test_class_method_body = textwrap.dedent("""\
        for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}):
            args = pickle.loads(args_pkl)
            kwargs = pickle.loads(kwargs_pkl){filter_variables}
            if not args:
                raise ValueError("No arguments provided for the method.")
            ret = {class_name_alias}{method_name}(*args[1:], **kwargs)
            """)

    test_static_method_body = textwrap.dedent("""\
        for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}):
            args = pickle.loads(args_pkl)
            kwargs = pickle.loads(kwargs_pkl){filter_variables}
            ret = {class_name_alias}{method_name}(*args, **kwargs)
            """)

    test_template = ""

    for func in functions_data:
        module_name = func["module_name"]
        function_name = func["function_name"]
        class_name = func.get("class_name", "")
        fp = Path(func["file_path"]).as_posix()
        benchmark_function_name = func["benchmark_function_name"]
        function_properties = func["function_properties"]
        if not class_name:
            alias = get_function_alias(module_name, function_name)
            test_body = test_function_body.format(
                benchmark_function_name=benchmark_function_name,
                orig_function_name=function_name,
                function_name=alias,
                file_path=fp,
                max_run_count=max_run_count,
            )
        else:
            class_name_alias = get_function_alias(module_name, class_name)
            filter_variables = ""
            method_name = (
                "." + function_name if function_name != "__init__" else ""
            )
            if function_properties.is_classmethod:
                test_body = test_class_method_body.format(
                    benchmark_function_name=benchmark_function_name,
                    orig_function_name=function_name,
                    file_path=fp,
                    class_name_alias=class_name_alias,
                    class_name=class_name,
                    method_name=method_name,
                    max_run_count=max_run_count,
                    filter_variables=filter_variables,
                )
            elif function_properties.is_staticmethod:
                test_body = test_static_method_body.format(
                    benchmark_function_name=benchmark_function_name,
                    orig_function_name=function_name,
                    file_path=fp,
                    class_name_alias=class_name_alias,
                    class_name=class_name,
                    method_name=method_name,
                    max_run_count=max_run_count,
                    filter_variables=filter_variables,
                )
            else:
                test_body = test_method_body.format(
                    benchmark_function_name=benchmark_function_name,
                    orig_function_name=function_name,
                    file_path=fp,
                    class_name_alias=class_name_alias,
                    class_name=class_name,
                    method_name=method_name,
                    max_run_count=max_run_count,
                    filter_variables=filter_variables,
                )

        formatted_test_body = textwrap.indent(test_body, "    ")
        unique_test_name = get_unique_test_name(
            module_name,
            function_name,
            benchmark_function_name,
            class_name,
        )
        test_template += (
            f"def test_{unique_test_name}():\n{formatted_test_body}\n"
        )

    return imports + "\n" + metadata + "\n" + test_template


def _get_replay_test_file_path(
    test_dir: Path,
    function_name: str,
    iteration: int = 0,
) -> Path:
    """Return the output path for a replay test file."""
    function_name_safe = function_name.replace(".", "_")
    path = test_dir / f"test_{function_name_safe}__replay_test_{iteration}.py"
    if path.exists():
        return _get_replay_test_file_path(
            test_dir,
            function_name,
            iteration + 1,
        )
    return path


def generate_replay_test(
    trace_file_path: Path,
    output_dir: Path,
    max_run_count: int = 100,
) -> int:
    """Generate replay tests from traced function calls, grouped by benchmark."""
    count = 0
    try:
        conn = sqlite3.connect(trace_file_path.as_posix())
        cursor = conn.cursor()

        cursor.execute(
            "SELECT DISTINCT benchmark_module_path"
            " FROM benchmark_function_timings"
        )
        benchmark_files = cursor.fetchall()

        for benchmark_file in benchmark_files:
            benchmark_module_path = benchmark_file[0]
            cursor.execute(
                "SELECT DISTINCT benchmark_function_name,"
                " function_name, class_name, module_name,"
                " file_path, benchmark_line_number"
                " FROM benchmark_function_timings"
                " WHERE benchmark_module_path = ?",
                (benchmark_module_path,),
            )

            functions_data = []
            for row in cursor.fetchall():
                (
                    benchmark_function_name,
                    function_name,
                    class_name,
                    module_name,
                    file_path,
                    benchmark_line_number,
                ) = row
                functions_data.append(
                    {
                        "function_name": function_name,
                        "class_name": class_name,
                        "file_path": file_path,
                        "module_name": module_name,
                        "benchmark_function_name": benchmark_function_name,
                        "benchmark_module_path": benchmark_module_path,
                        "benchmark_line_number": benchmark_line_number,
                        "function_properties": (
                            inspect_top_level_functions_or_methods(
                                file_name=Path(file_path),
                                function_or_method_name=function_name,
                                class_name=class_name,
                            )
                        ),
                    }
                )

            if not functions_data:
                log.info(
                    "No benchmark test functions found in %s",
                    benchmark_module_path,
                )
                continue

            test_code = create_trace_replay_test_code(
                trace_file=trace_file_path.as_posix(),
                functions_data=functions_data,
                max_run_count=max_run_count,
            )
            sorted_code = sort_imports(test_code)
            output_file = _get_replay_test_file_path(
                test_dir=output_dir,
                function_name=benchmark_module_path,
            )
            output_dir.mkdir(parents=True, exist_ok=True)
            output_file.write_text(sorted_code, encoding="utf-8")
            count += 1

        conn.close()
    except Exception:
        log.exception("Error generating replay tests")

    return count


def validate_and_format_benchmark_table(
    function_benchmark_timings: dict[str, dict[BenchmarkKey, int]],
    total_benchmark_timings: dict[BenchmarkKey, int],
) -> dict[str, list[tuple[BenchmarkKey, float, float, float]]]:
    """Validate timings and return sorted benchmark result table per function."""
    function_to_result: dict[
        str,
        list[tuple[BenchmarkKey, float, float, float]],
    ] = {}
    for (
        func_path,
        test_times,
    ) in function_benchmark_timings.items():
        sorted_tests: list[tuple[BenchmarkKey, float, float, float]] = []
        for benchmark_key, func_time in test_times.items():
            total_time = total_benchmark_timings.get(benchmark_key, 0)
            if func_time > total_time:
                sorted_tests.append((benchmark_key, 0.0, 0.0, 0.0))
            elif total_time > 0:
                percentage = (func_time / total_time) * 100
                func_time_ms = func_time / 1_000_000
                total_time_ms = total_time / 1_000_000
                sorted_tests.append(
                    (
                        benchmark_key,
                        total_time_ms,
                        func_time_ms,
                        percentage,
                    )
                )
        sorted_tests.sort(key=lambda x: x[3], reverse=True)
        function_to_result[func_path] = sorted_tests
    return function_to_result


def process_benchmark_data(
    replay_performance_gain: dict[BenchmarkKey, float],
    fto_benchmark_timings: dict[BenchmarkKey, int],
    total_benchmark_timings: dict[BenchmarkKey, int],
) -> ProcessedBenchmarkInfo | None:
    """Compute expected benchmark speedups and return a summary, or ``None``."""
    if (
        not replay_performance_gain
        or not fto_benchmark_timings
        or not total_benchmark_timings
    ):
        return None
    benchmark_details: list[BenchmarkDetail] = []
    for (
        benchmark_key,
        og_benchmark_timing,
    ) in fto_benchmark_timings.items():
        total_benchmark_timing = total_benchmark_timings.get(benchmark_key, 0)
        if total_benchmark_timing == 0:
            continue
        expected_new_benchmark_timing = (
            total_benchmark_timing
            - og_benchmark_timing
            + (1 / (replay_performance_gain[benchmark_key] + 1))
            * og_benchmark_timing
        )
        benchmark_speedup_percent = (
            performance_gain(
                original_runtime_ns=total_benchmark_timing,
                optimized_runtime_ns=int(expected_new_benchmark_timing),
            )
            * 100
        )
        benchmark_details.append(
            BenchmarkDetail(
                benchmark_name=benchmark_key.module_path,
                test_function=benchmark_key.function_name,
                original_timing=humanize_runtime(int(total_benchmark_timing)),
                expected_new_timing=humanize_runtime(
                    int(expected_new_benchmark_timing)
                ),
                speedup_percent=benchmark_speedup_percent,
            )
        )
    return ProcessedBenchmarkInfo(benchmark_details=tuple(benchmark_details))


def fmt_ms(ns: int | None) -> str:
    """Format nanoseconds as a human-readable millisecond string."""
    if ns is None:
        return "-"
    ms = ns / 1_000_000
    if ms >= 1000:
        return f"{ms:,.0f}"
    if ms >= 100:
        return f"{ms:.0f}"
    if ms >= 1:
        return f"{ms:.1f}"
    return f"{ms:.2f}"


def md_speedup(before: int | None, after: int | None) -> str:
    """Return a markdown speedup indicator (ratio with coloured dot)."""
    if before is None or after is None or after == 0:
        return "-"
    ratio = before / after
    emoji = "\U0001f7e2" if ratio >= 1 else "\U0001f534"
    return f"{emoji} {ratio:.2f}x"


def md_delta(before: int | None, after: int | None) -> str:
    """Return a signed millisecond delta with percentage."""
    if before is None or after is None:
        return "-"
    delta_ms = (after - before) / 1_000_000
    pct = ((after - before) / before) * 100 if before != 0 else 0
    if delta_ms < 0:
        return f"{delta_ms:+,.0f}ms ({pct:+.0f}%)"
    return f"+{delta_ms:,.0f}ms ({pct:+.0f}%)"


def md_bar(
    before: int | None,
    after: int | None,
    width: int = 10,
) -> str:
    """Return an ASCII bar chart with percentage for markdown."""
    if before is None or after is None or before == 0:
        return "-"
    pct = ((before - after) / before) * 100
    filled = round(abs(pct) / 100 * width)
    filled = min(filled, width)
    bar = "\u2588" * filled + "\u2591" * (width - filled)
    return f"`{bar}` {pct:+.0f}%"


def pct_bar(pct: float, width: int = 10) -> str:
    """Return an ASCII bar chart for a given percentage."""
    filled = round(pct / 100 * width)
    filled = max(0, min(filled, width))
    bar = "\u2588" * filled + "\u2591" * (width - filled)
    return f"`{bar}` {pct:.1f}%"


@attrs.frozen
class CompareResult:
    """Results from a cross-branch benchmark comparison."""

    base_ref: str
    head_ref: str
    base_total_ns: dict[BenchmarkKey, int] = attrs.Factory(dict)
    head_total_ns: dict[BenchmarkKey, int] = attrs.Factory(dict)
    base_function_ns: dict[str, dict[BenchmarkKey, int]] = attrs.Factory(dict)
    head_function_ns: dict[str, dict[BenchmarkKey, int]] = attrs.Factory(dict)

    def format_markdown(self) -> str:
        """Render a full markdown comparison report."""
        lines: list[str] = []
        lines.append(
            f"## Benchmark comparison: `{self.base_ref}` vs `{self.head_ref}`"
        )
        lines.append("")

        # --- end-to-end table ---
        all_keys = sorted(
            set(self.base_total_ns) | set(self.head_total_ns),
            key=str,
        )
        if all_keys:
            lines.append("### End-to-end benchmarks")
            lines.append("")
            lines.append(
                "| Benchmark | Base (ms) |"
                " Head (ms) | Speedup | Delta |"
                " Chart |"
            )
            lines.append("|---|---:|---:|:---:|---:|---|")
            for key in all_keys:
                base_ns = self.base_total_ns.get(key)
                head_ns = self.head_total_ns.get(key)
                lines.append(
                    f"| `{key}` |"
                    f" {fmt_ms(base_ns)} |"
                    f" {fmt_ms(head_ns)} |"
                    f" {md_speedup(base_ns, head_ns)}"
                    f" |"
                    f" {md_delta(base_ns, head_ns)} |"
                    f" {md_bar(base_ns, head_ns)} |"
                )
            lines.append("")

        # --- per-function breakdown ---
        all_funcs = sorted(
            set(self.base_function_ns) | set(self.head_function_ns)
        )
        if all_funcs:
            lines.append("### Per-function breakdown")
            lines.append("")
            for func in all_funcs:
                base_map = self.base_function_ns.get(func, {})
                head_map = self.head_function_ns.get(func, {})
                func_keys = sorted(
                    set(base_map) | set(head_map),
                    key=str,
                )
                if not func_keys:
                    continue
                lines.append(f"#### `{func}`")
                lines.append("")
                lines.append(
                    "| Benchmark | Base (ms) |"
                    " Head (ms) | Share of"
                    " benchmark | Speedup |"
                )
                lines.append("|---|---:|---:|---|:---:|")
                for key in func_keys:
                    b_ns = base_map.get(key)
                    h_ns = head_map.get(key)
                    total_ns = self.base_total_ns.get(key)
                    share = "-"
                    if b_ns is not None and total_ns:
                        share_pct = b_ns / total_ns * 100
                        share = pct_bar(share_pct)
                    lines.append(
                        f"| `{key}` |"
                        f" {fmt_ms(b_ns)} |"
                        f" {fmt_ms(h_ns)} |"
                        f" {share} |"
                        f" {md_speedup(b_ns, h_ns)}"
                        f" |"
                    )
                lines.append("")

        return "\n".join(lines)
