"""Project-level optimization orchestrator for Python.

Composes function discovery, global ranking, module preparation,
and per-function optimization into a complete run.  The
orchestrator owns the outer loop; per-function work is delegated
to a caller-supplied callable (*optimize_fn*).

This keeps the orchestrator decoupled from AI clients, test
runners, and other infrastructure — callers wire those in via
the callable.  Different callers (CLI, subagent, tests) can
provide different per-function strategies.
"""

from __future__ import annotations

import logging
import sys
from pathlib import Path
from typing import TYPE_CHECKING

import attrs

if sys.version_info >= (3, 11):
    from typing import Self
else:
    from typing_extensions import Self

if TYPE_CHECKING:
    import ast
    from collections.abc import Awaitable, Callable

    import git

    from codeflash_core import (
        BenchmarkDetail,
        Candidate,
        OptimizationReviewResult,
        PlatformClient,
    )

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

log = logging.getLogger(__name__)


@attrs.frozen
class FunctionInput:
    """Everything the per-function optimizer needs for one function.

    Built by :meth:`PythonOptimizer.run` and passed to the
    caller-supplied *optimize_fn*.
    """

    function: FunctionToOptimize
    module_path: Path
    source_code: str
    normalized_code: str
    module_ast: ast.Module
    validated_code: dict[Path, ValidCode]
    function_benchmark_timings: dict[BenchmarkKey, int] = attrs.Factory(
        dict,
    )
    total_benchmark_timings: dict[BenchmarkKey, int] = attrs.Factory(
        dict,
    )


@attrs.frozen
class PrData:
    """Data needed to create a PR for a single function optimization.

    Populated by :class:`PythonFunctionOptimizer` and consumed by
    :meth:`PythonOptimizer.run` when PR creation is enabled.
    """

    function_trace_id: str
    existing_tests_source: str = ""
    generated_tests_source: str = ""
    replay_tests: str = ""
    concolic_tests: str = ""
    coverage_message: str = ""
    speedup_x: str = ""
    speedup_pct: str = ""
    best_runtime_ns: int = 0
    original_runtime_ns: int = 0
    loop_count: int = 0
    report_table: dict[str, dict[str, int]] = attrs.Factory(dict)
    benchmark_details: tuple[BenchmarkDetail, ...] | None = None
    original_async_throughput: int | None = None
    best_async_throughput: int | None = None


@attrs.frozen
class FunctionResult:
    """Outcome of optimizing a single function."""

    function: FunctionToOptimize
    module_path: Path
    success: bool
    message: str = ""
    best_candidate: Candidate | None = None
    explanation: str = ""
    review: OptimizationReviewResult | None = None
    pr_data: PrData | None = None

    def to_dict(self) -> dict[str, object]:
        """Serialize to a plain dictionary."""
        candidate = None
        if self.best_candidate is not None:
            candidate = {
                "code": self.best_candidate.code,
                "explanation": self.best_candidate.explanation,
                "candidate_id": self.best_candidate.candidate_id,
            }
        review = None
        if self.review is not None:
            review = {
                "review": self.review.review,
                "explanation": self.review.explanation,
            }
        return {
            "function": self.function.to_dict(),
            "module_path": str(self.module_path),
            "success": self.success,
            "message": self.message,
            "best_candidate": candidate,
            "explanation": self.explanation,
            "review": review,
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> Self:
        """Restore from a serialized dictionary."""
        from codeflash_core import (  # noqa: PLC0415
            Candidate,
            OptimizationReviewResult,
        )

        from .._model import FunctionToOptimize  # noqa: PLC0415

        fn_data = data["function"]
        fn = FunctionToOptimize.from_dict(fn_data)  # type: ignore[arg-type]
        candidate = None
        cand_data = data.get("best_candidate")
        if cand_data is not None:
            candidate = Candidate(**cand_data)  # type: ignore[arg-type]
        review = None
        review_data = data.get("review")
        if review_data is not None:
            review = OptimizationReviewResult(**review_data)  # type: ignore[arg-type]
        return cls(
            function=fn,
            module_path=Path(str(data["module_path"])),
            success=bool(data["success"]),
            message=str(data.get("message", "")),
            best_candidate=candidate,
            explanation=str(data.get("explanation", "")),
            review=review,
        )


@attrs.define
class PythonOptimizer:
    """Project-level optimization orchestrator for Python.

    Composes function discovery, global ranking, module
    preparation, and per-function optimization into a complete
    run.  The optimizer owns the outer loop; per-function work
    is delegated to a caller-supplied *optimize_fn*.

    Usage::

        cfg = PythonConfiguration(...)
        state = PythonState(cfg=cfg)
        plugin = PythonPlugin(configuration=cfg, state=state)
        optimizer = PythonOptimizer(
            plugin=plugin,
            project_root=project_root,
        )
        results = optimizer.run(
            file_to_funcs=discovered,
            optimize_fn=my_per_function_optimizer,
        )
    """

    plugin: PythonPlugin
    project_root: Path
    no_pr: bool = False
    git_remote: str = "origin"
    platform_client: PlatformClient | None = None
    git_repo: git.Repo | None = None

    async def run(  # noqa: PLR0913
        self,
        file_to_funcs: dict[Path, list[FunctionToOptimize]],
        *,
        optimize_fn: Callable[[FunctionInput], Awaitable[FunctionResult]],
        trace_file: Path | None = None,
        call_graph: ReferenceGraph | None = None,
        tests_root: Path | None = None,
        function_benchmark_timings: dict[str, dict[BenchmarkKey, int]]
        | None = None,
        total_benchmark_timings: dict[BenchmarkKey, int] | None = None,
    ) -> list[FunctionResult]:
        """Run the project-level optimization loop.

        1. Clean up leftover instrumented test files.
        2. Rank all functions globally by optimization impact.
        3. Prepare each module once (parse, normalize, validate).
        4. Call *optimize_fn* for each ranked function.

        *function_benchmark_timings* maps qualified function names
        (with module) to per-benchmark timing dicts.
        *total_benchmark_timings* holds the global median ns per
        benchmark key.  Both are collected by
        :class:`~codeflash_python.benchmarking.CodeFlashBenchmarkPlugin`
        and forwarded to the per-function optimizer via
        :class:`FunctionInput`.

        Returns one :class:`FunctionResult` per function processed.
        """
        from ._orchestrator import (  # noqa: PLC0415
            prepare_module_for_optimization,
            rank_functions_globally,
        )

        # 1. Global ranking.
        ranked = rank_functions_globally(
            file_to_funcs,
            trace_file_path=trace_file,
            call_graph=call_graph,
        )
        if not ranked:
            log.info("No functions to optimize")
            return []

        log.info("Optimizing %d functions", len(ranked))

        # 3 & 4. Prepare modules (cached) and optimize functions.
        prepared: dict[Path, tuple[dict[Path, ValidCode], ast.Module]] = {}
        results: list[FunctionResult] = []

        for i, (module_path, func) in enumerate(ranked):
            log.info(
                "Optimizing function %d of %d: %s",
                i + 1,
                len(ranked),
                func.qualified_name,
            )

            # Prepare module once per file.
            if module_path not in prepared:
                prep = prepare_module_for_optimization(
                    module_path,
                    self.project_root,
                )
                if prep is None:
                    log.warning(
                        "Skipping %s — module preparation failed",
                        module_path,
                    )
                    results.append(
                        FunctionResult(
                            function=func,
                            module_path=module_path,
                            success=False,
                            message="Module preparation failed",
                        ),
                    )
                    continue
                prepared[module_path] = prep

            validated_code, module_ast = prepared[module_path]
            module_valid = validated_code[module_path]

            # Extract per-function benchmark timings.
            # The all-functions dict is keyed by qualified name
            # with module; we look up this function's entry and
            # pass it (plus the global totals) through to the
            # per-function optimizer.
            func_timings: dict[BenchmarkKey, int] = {}
            func_total_timings: dict[BenchmarkKey, int] = {}
            if function_benchmark_timings and total_benchmark_timings:
                qname = func.qualified_name_with_modules_from_root(
                    self.project_root,
                )
                if qname in function_benchmark_timings:
                    func_timings = function_benchmark_timings[qname]
                    func_total_timings = total_benchmark_timings

            # Delegate to per-function optimizer.
            result = await optimize_fn(
                FunctionInput(
                    function=func,
                    module_path=module_path,
                    source_code=module_valid.source_code,
                    normalized_code=module_valid.normalized_code,
                    module_ast=module_ast,
                    validated_code=validated_code,
                    function_benchmark_timings=func_timings,
                    total_benchmark_timings=func_total_timings,
                ),
            )
            results.append(result)

            # Apply winning candidate's code to disk.
            if result.success and result.best_candidate is not None:
                updated = self._apply_candidate(
                    result,
                    module_valid.source_code,
                )
                # Create PR if enabled and review is not "low".
                if updated is not None:
                    self._maybe_create_pr(
                        result,
                        module_valid.source_code,
                        updated,
                    )

        optimized = sum(1 for r in results if r.success)
        log.info(
            "Optimization complete: %d of %d functions improved",
            optimized,
            len(results),
        )
        return results

    @staticmethod
    def _apply_candidate(
        result: FunctionResult,
        original_source: str,
    ) -> str | None:
        """Write the winning candidate's code to disk.

        Returns the updated source string on success, or *None*
        on failure.
        """
        from ..codegen._replacement import (  # noqa: PLC0415
            replace_functions_in_file,
        )

        candidate = result.best_candidate
        if candidate is None:
            return None

        try:
            updated = replace_functions_in_file(
                source_code=original_source,
                original_function_names=[
                    result.function.function_name,
                ],
                optimized_code=candidate.code,
                preexisting_objects=set(),
            )
        except Exception:  # noqa: BLE001
            log.warning(
                "Failed to apply optimized code for %s",
                result.function.qualified_name,
            )
            return None

        result.module_path.write_text(updated, encoding="utf-8")
        log.info(
            "Applied optimized code for %s",
            result.function.qualified_name,
        )
        return updated

    def _maybe_create_pr(
        self,
        result: FunctionResult,
        original_source: str,
        new_source: str,
    ) -> None:
        """Create a PR for this function if conditions are met."""
        if self.no_pr:
            return
        if self.platform_client is None or self.git_repo is None:
            return
        review = result.review
        if review is not None and review.review.lower() == "low":
            return
        pr_data = result.pr_data
        if pr_data is None:
            return

        from codeflash_core import PrComment  # noqa: PLC0415

        from ..codegen._create_pr import (  # noqa: PLC0415
            check_create_pr,
        )

        rel_path = (
            result.module_path.resolve()
            .relative_to(self.project_root.resolve())
            .as_posix()
        )
        original_code = {rel_path: original_source}
        new_code = {rel_path: new_source}

        pr_comment = PrComment(
            optimization_explanation=result.explanation,
            best_runtime=pr_data.best_runtime_ns,
            original_runtime=pr_data.original_runtime_ns,
            function_name=result.function.qualified_name,
            relative_file_path=rel_path,
            speedup_x=pr_data.speedup_x,
            speedup_pct=pr_data.speedup_pct,
            loop_count=pr_data.loop_count,
            report_table=pr_data.report_table,
            benchmark_details=pr_data.benchmark_details,
            original_async_throughput=(pr_data.original_async_throughput),
            best_async_throughput=pr_data.best_async_throughput,
        )

        try:
            check_create_pr(
                platform_client=self.platform_client,
                git_repo=self.git_repo,
                original_code=original_code,
                new_code=new_code,
                pr_comment=pr_comment,
                existing_tests=pr_data.existing_tests_source,
                generated_tests=pr_data.generated_tests_source,
                function_trace_id=pr_data.function_trace_id,
                coverage_message=pr_data.coverage_message,
                replay_tests=pr_data.replay_tests,
                concolic_tests=pr_data.concolic_tests,
                optimization_review=(
                    review.review if review is not None else ""
                ),
                git_remote=self.git_remote,
            )
        except Exception:  # noqa: BLE001
            log.warning(
                "PR creation failed for %s",
                result.function.qualified_name,
                exc_info=True,
            )
