"""High-level pipeline orchestrator for Python optimization.

Ports the ``Optimizer`` helpers from the reference implementation,
providing functions that wire together all pipeline stages into a
complete optimization run.
"""

from __future__ import annotations

import logging
import re
import shutil
import tempfile
from collections import defaultdict
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import ast
    from collections.abc import Sequence

    from .._model import FunctionToOptimize
    from ..analysis._reference_graph import ReferenceGraph
    from ..benchmarking.models import BenchmarkKey
    from ._module_prep import ValidCode

log = logging.getLogger(__name__)

# Regex for leftover instrumented test files from previous runs.
_INSTRUMENTED_TEST_PATTERN = re.compile(
    r"(?:"
    r"test.*__perf_test_\d?\.py|"
    r"test_.*__unit_test_\d?\.py|"
    r"test_.*__perfinstrumented\.py|"
    r"test_.*__perfonlyinstrumented\.py"
    r")$"
)


def cleanup_paths(paths: Sequence[Path | None]) -> None:
    """Remove files and directories, ignoring missing or ``None`` paths."""
    for path in paths:
        if path is None or not path.exists():
            continue
        if path.is_dir():
            shutil.rmtree(path, ignore_errors=True)
        else:
            path.unlink(missing_ok=True)


def find_leftover_instrumented_test_files(
    test_root: Path,
) -> list[Path]:
    """Find instrumented test files left over from previous runs.

    Matches patterns like ``test_*__perfinstrumented.py`` and
    ``test_*__perf_test_0.py``.  Also returns stale ``tmp*``
    directories created by previous runs so that leftover
    ``__pycache__`` entries don't confuse test discovery.
    """
    stale_dirs: set[Path] = set()
    files: list[Path] = []
    for file_path in test_root.rglob("*"):
        if not file_path.is_file():
            continue
        if not _INSTRUMENTED_TEST_PATTERN.match(file_path.name):
            continue
        files.append(file_path)
        # If this file lives in a tmp* subdirectory of test_root,
        # mark the whole directory for removal.
        parent = file_path.parent
        if (
            parent != test_root
            and parent.parent == test_root
            and parent.name.startswith("tmp")
        ):
            stale_dirs.add(parent)

    # Also catch tmp dirs that only have __pycache__ left
    # (instrumented .py files were cleaned but dir wasn't).
    for entry in test_root.iterdir():
        if (
            entry.is_dir()
            and entry.name.startswith("tmp")
            and entry not in stale_dirs
        ):
            stale_dirs.add(entry)

    # Return dirs first (rmtree handles contents); then individual
    # files that live directly in test_root.
    return [
        *stale_dirs,
        *(f for f in files if f.parent not in stale_dirs),
    ]


def prepare_module_for_optimization(
    module_path: Path,
    project_root: Path,
) -> tuple[dict[Path, ValidCode], ast.Module] | None:
    """Prepare a module for optimization.

    Reads the module source and calls :func:`prepare_python_module`.
    Returns ``None`` if the module has syntax errors or invalid
    callees.
    """
    from ._module_prep import prepare_python_module  # noqa: PLC0415

    log.info("Examining file %s", module_path)
    original_code = module_path.read_text(encoding="utf8")
    return prepare_python_module(original_code, module_path, project_root)


def rank_functions_globally(
    file_to_funcs: dict[Path, list[FunctionToOptimize]],
    trace_file_path: Path | None = None,
    call_graph: ReferenceGraph | None = None,
    test_count_cache: dict[tuple[Path, str], int] | None = None,
) -> list[tuple[Path, FunctionToOptimize]]:
    """Rank all functions across all files by optimization impact.

    Tries trace-based ranking first via :class:`FunctionRanker`,
    then falls back to dependency-count ranking using *call_graph*,
    and finally returns the original discovery order.

    *test_count_cache*, when provided, is used as a secondary sort key
    so that functions with more existing unit tests are preferred among
    those with the same primary score.
    """
    all_functions: list[tuple[Path, FunctionToOptimize]] = []
    for file_path, functions in file_to_funcs.items():
        all_functions.extend((file_path, func) for func in functions)

    if not trace_file_path or not trace_file_path.exists():
        if call_graph is not None:
            return rank_by_dependency_count(
                all_functions,
                call_graph,
                test_count_cache=test_count_cache,
            )
        log.debug("No trace file available, using original function order")
        return all_functions

    try:
        from ..analysis._function_ranking import (  # noqa: PLC0415
            FunctionRanker,
        )

        ranker = FunctionRanker(trace_file_path)
        functions_only = [func for _, func in all_functions]
        ranked_functions = ranker.rank_functions(functions_only)

        # Build reverse mapping: function -> file path.
        func_to_file: dict[tuple[Path, str, int | None], Path] = {}
        for file_path, func in all_functions:
            key = (
                func.file_path,
                func.qualified_name,
                func.starting_line,
            )
            func_to_file[key] = file_path

        ranked_with_metadata: list[
            tuple[Path, FunctionToOptimize, float, int]
        ] = []
        for rank_index, func in enumerate(ranked_functions):
            key = (
                func.file_path,
                func.qualified_name,
                func.starting_line,
            )
            matched_path = func_to_file.get(key)
            if matched_path is not None:
                ranked_with_metadata.append(
                    (
                        matched_path,
                        func,
                        ranker.get_function_addressable_time(func),
                        rank_index,
                    )
                )

        if test_count_cache:
            ranked_with_metadata.sort(
                key=lambda item: (
                    -item[2],
                    -test_count_cache.get(
                        (item[0], item[1].qualified_name), 0
                    ),
                    item[3],
                )
            )

        globally_ranked = [
            (file_path, func) for file_path, func, _, _ in ranked_with_metadata
        ]

        log.info(
            "Globally ranked %d functions by addressable time "
            "(filtered %d low-importance functions)",
            len(ranked_functions),
            len(functions_only) - len(ranked_functions),
        )
    except Exception:  # noqa: BLE001
        log.warning(
            "Could not perform global ranking",
            exc_info=True,
        )
        return all_functions
    else:
        return globally_ranked


def rank_by_dependency_count(
    all_functions: list[tuple[Path, FunctionToOptimize]],
    call_graph: ReferenceGraph,
    test_count_cache: dict[tuple[Path, str], int] | None = None,
) -> list[tuple[Path, FunctionToOptimize]]:
    """Rank functions by number of callees (most complex first).

    *test_count_cache*, when provided, is used as a secondary sort key
    so that functions with more existing unit tests are preferred among
    those with the same callee count.
    """
    file_to_qns: dict[Path, set[str]] = defaultdict(set)
    for file_path, func in all_functions:
        file_to_qns[file_path].add(func.qualified_name)

    callee_counts = call_graph.count_callees_per_function(dict(file_to_qns))

    if test_count_cache:
        ranked = sorted(
            enumerate(all_functions),
            key=lambda x: (
                -callee_counts.get((x[1][0], x[1][1].qualified_name), 0),
                -test_count_cache.get((x[1][0], x[1][1].qualified_name), 0),
                x[0],
            ),
        )
    else:
        ranked = sorted(
            enumerate(all_functions),
            key=lambda x: (
                -callee_counts.get((x[1][0], x[1][1].qualified_name), 0),
                x[0],
            ),
        )
    log.debug("Ranked %d functions by dependency count", len(ranked))
    return [item for _, item in ranked]


def run_benchmarks(
    file_to_funcs: dict[Path, list[FunctionToOptimize]],
    benchmarks_root: Path,
    tests_root: Path,
    project_root: Path,
) -> tuple[
    dict[str, dict[BenchmarkKey, int]],
    dict[BenchmarkKey, int],
    Path | None,
]:
    """Run benchmarks and collect per-function timing data.

    Instruments source files with ``@codeflash_trace``, runs
    benchmark tests via pytest in a subprocess, generates replay
    tests from the trace, and extracts timing data.

    Original source files are always restored, even on failure.

    Returns ``(function_benchmark_timings, total_benchmark_timings,
    replay_tests_dir)``.  All empty when no benchmarks are found.
    """
    from ..benchmarking._benchmark_plugin import (  # noqa: PLC0415
        CodeFlashBenchmarkPlugin,
    )
    from ..benchmarking._benchmarking import (  # noqa: PLC0415
        generate_replay_test,
        instrument_codeflash_trace_decorator,
    )
    from ..testing._subprocess_runners import (  # noqa: PLC0415
        trace_benchmarks_pytest,
    )

    function_benchmark_timings: dict[str, dict[BenchmarkKey, int]] = {}
    total_benchmark_timings: dict[BenchmarkKey, int] = {}
    replay_tests_dir: Path | None = None

    # Save original source so we can restore after instrumentation.
    file_path_to_source: dict[Path, str] = {}
    for file_path in file_to_funcs:
        file_path_to_source[file_path] = file_path.read_text(
            encoding="utf-8",
        )

    try:
        instrument_codeflash_trace_decorator(file_to_funcs)

        trace_file = benchmarks_root / "benchmarks.trace"
        if trace_file.exists():
            trace_file.unlink()

        replay_tests_dir = Path(
            tempfile.mkdtemp(
                prefix="codeflash_replay_tests_",
                dir=benchmarks_root,
            ),
        )

        trace_benchmarks_pytest(
            benchmarks_root,
            tests_root,
            project_root,
            trace_file,
        )

        replay_count = generate_replay_test(
            trace_file,
            replay_tests_dir,
        )
        if replay_count == 0:
            log.info(
                "No valid benchmarks found in %s",
                benchmarks_root,
            )
        else:
            function_benchmark_timings = (
                CodeFlashBenchmarkPlugin.get_function_benchmark_timings(
                    trace_file,
                )
            )
            total_benchmark_timings = (
                CodeFlashBenchmarkPlugin.get_benchmark_timings(
                    trace_file,
                )
            )
            log.info(
                "Collected benchmark timings for %d function(s)"
                " across %d benchmark(s)",
                len(function_benchmark_timings),
                len(total_benchmark_timings),
            )
    except Exception:  # noqa: BLE001
        log.info(
            "Error while tracing existing benchmarks",
            exc_info=True,
        )
        log.info(
            "Benchmark information will not be available for this run",
        )
    finally:
        # Always restore original source code.
        for file_path, source in file_path_to_source.items():
            file_path.write_text(source, encoding="utf-8")

    return (
        function_benchmark_timings,
        total_benchmark_timings,
        replay_tests_dir,
    )
