"""Optimization session for programmatic API access."""

from __future__ import annotations

import logging
import sys
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:
    from pathlib import Path

    from codeflash_core import AIClient, Candidate

    from .._model import FunctionToOptimize
    from ..context.models import CodeOptimizationContext
    from ._config import OptimizationConfig

log = logging.getLogger(__name__)


@attrs.define
class OptimizationSession:
    """A stateful session for step-by-step optimization.

    Wraps the internal pipeline building blocks behind a simple
    interface suitable for programmatic use by LLM agents or MCP
    tools.  Use as a context manager to ensure the AI client is
    closed properly::

        with OptimizationSession(config) as session:
            funcs = session.discover_functions(Path("module.py"))
            ctx = session.extract_context(funcs[0])
            candidates = session.generate_candidates(funcs[0], ctx)
    """

    config: OptimizationConfig
    _ai_client: AIClient | None = attrs.field(
        default=None,
        init=False,
        repr=False,
    )

    def __enter__(self) -> Self:
        """Enter the context manager."""
        return self

    def __exit__(self, *exc_info: object) -> None:
        """Exit the context manager."""
        self.close()

    @property
    def ai_client(self) -> AIClient:
        """Lazily create and return the AI client."""
        if self._ai_client is None:
            from codeflash_core import AIClient  # noqa: PLC0415

            self._ai_client = AIClient(
                api_key=self.config.api_key,
                timeout=self.config.ai_timeout,
            )
        return self._ai_client

    def discover_functions(
        self,
        file_path: Path,
    ) -> list[FunctionToOptimize]:
        """Discover optimizable functions in a Python file."""
        from ..analysis._discovery import (  # noqa: PLC0415
            discover_functions,
        )

        source = file_path.read_text(encoding="utf-8")
        return discover_functions(source, file_path)

    def extract_context(
        self,
        function: FunctionToOptimize,
    ) -> CodeOptimizationContext:
        """Extract optimization context for a function."""
        from ..context.pipeline import (  # noqa: PLC0415
            get_code_optimization_context,
        )

        return get_code_optimization_context(
            function,
            self.config.project_root,
        )

    def generate_candidates(
        self,
        function: FunctionToOptimize,
        context: CodeOptimizationContext,
    ) -> list[Candidate]:
        """Generate optimization candidates from the AI service."""
        import platform  # noqa: PLC0415

        from codeflash_core import OptimizationRequest  # noqa: PLC0415

        request = OptimizationRequest(
            source_code=context.read_writable,
            language="python",
            language_version=platform.python_version(),
            context_code=context.read_only,
        )
        return self.ai_client.get_candidates(
            request,
            n_candidates=self.config.n_candidates,
        )

    def apply(
        self,
        function: FunctionToOptimize,
        new_source: str,
        original_source: str,
    ) -> str:
        """Replace a function in source code with optimized code."""
        from ..codegen._replacement import (  # noqa: PLC0415
            replace_function_source,
        )

        return replace_function_source(
            original_source,
            function,
            new_source,
        )

    def close(self) -> None:
        """Close the AI client and release resources."""
        if self._ai_client is not None:
            self._ai_client.close()
            self._ai_client = None

    # -- Agent experiment loop (future) ---------------------------------

    def profile(self, **kwargs: object) -> None:
        """Profile code for optimization targets."""
        msg = "Experiment loop not yet implemented"
        raise NotImplementedError(msg)

    def build_targets(self, **kwargs: object) -> None:
        """Build optimization targets from profiling data."""
        msg = "Experiment loop not yet implemented"
        raise NotImplementedError(msg)

    def measure(self, **kwargs: object) -> None:
        """Measure performance of optimized code."""
        msg = "Experiment loop not yet implemented"
        raise NotImplementedError(msg)

    def evaluate(self, **kwargs: object) -> None:
        """Evaluate whether to keep or discard optimizations."""
        msg = "Experiment loop not yet implemented"
        raise NotImplementedError(msg)


def optimize_function(
    config: OptimizationConfig,
    file_path: Path,
    function_name: str,
) -> Candidate | None:
    """One-shot optimization of a single function.

    Discovers functions in *file_path*, extracts context for the one
    matching *function_name*, generates AI candidates, and returns
    the first candidate — or *None* if the function is not found or
    no candidates are generated.
    """
    with OptimizationSession(config) as session:
        functions = session.discover_functions(file_path)
        match = next(
            (f for f in functions if f.function_name == function_name),
            None,
        )
        if match is None:
            return None

        context = session.extract_context(match)
        candidates = session.generate_candidates(match, context)
        if not candidates:
            return None
        return candidates[0]
