"""Tests for OptimizationSession and optimize_function."""

from __future__ import annotations

import textwrap
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from codeflash_python.api import (
    OptimizationConfig,
    OptimizationSession,
    optimize_function,
)


@pytest.fixture(name="config")
def _config(tmp_path: Path) -> OptimizationConfig:
    return OptimizationConfig(
        project_root=tmp_path,
        module_root=tmp_path / "src",
    )


@pytest.fixture(name="sample_file")
def _sample_file(tmp_path: Path) -> Path:
    p = tmp_path / "sample.py"
    p.write_text(
        textwrap.dedent("""\
            def add(a, b):
                return a + b

            def multiply(x, y):
                return x * y
        """),
        encoding="utf-8",
    )
    return p


class TestOptimizationSession:
    """OptimizationSession construction and lifecycle."""

    def test_context_manager(self, config: OptimizationConfig) -> None:
        """Session can be used as a context manager."""
        with OptimizationSession(config) as session:
            assert session.config is config

    def test_close_without_ai_client(
        self,
        config: OptimizationConfig,
    ) -> None:
        """Closing without ever using ai_client is a no-op."""
        session = OptimizationSession(config)
        session.close()

    def test_close_with_ai_client(
        self,
        config: OptimizationConfig,
    ) -> None:
        """Closing after using ai_client calls close on it."""
        session = OptimizationSession(config)
        mock_client = MagicMock()
        session._ai_client = mock_client
        session.close()
        mock_client.close.assert_called_once()
        assert session._ai_client is None

    def test_context_manager_calls_close(
        self,
        config: OptimizationConfig,
    ) -> None:
        """Exiting the context manager calls close."""
        with OptimizationSession(config) as session:
            mock_client = MagicMock()
            session._ai_client = mock_client
        mock_client.close.assert_called_once()

    def test_ai_client_lazy_creation(
        self,
        config: OptimizationConfig,
    ) -> None:
        """AI client is not created until accessed."""
        session = OptimizationSession(config)
        assert session._ai_client is None

    def test_ai_client_created_on_access(
        self,
        config: OptimizationConfig,
    ) -> None:
        """AI client is created on first property access."""
        mock_cls = MagicMock()
        with patch(
            "codeflash_core.AIClient",
            mock_cls,
        ):
            session = OptimizationSession(config)
            client = session.ai_client
            assert client is mock_cls.return_value
            mock_cls.assert_called_once_with(
                api_key=config.api_key,
                timeout=config.ai_timeout,
            )

    def test_discover_functions(
        self,
        config: OptimizationConfig,
        sample_file: Path,
    ) -> None:
        """discover_functions finds functions in a file."""
        session = OptimizationSession(config)
        functions = session.discover_functions(sample_file)
        names = {f.function_name for f in functions}
        assert "add" in names
        assert "multiply" in names

    def test_discover_functions_empty_file(
        self,
        config: OptimizationConfig,
        tmp_path: Path,
    ) -> None:
        """discover_functions returns empty list for empty file."""
        empty = tmp_path / "empty.py"
        empty.write_text("", encoding="utf-8")
        session = OptimizationSession(config)
        assert [] == session.discover_functions(empty)


class TestExperimentLoopStubs:
    """Experiment loop methods raise NotImplementedError."""

    def test_profile_raises(
        self,
        config: OptimizationConfig,
    ) -> None:
        """profile raises NotImplementedError."""
        session = OptimizationSession(config)
        with pytest.raises(NotImplementedError):
            session.profile()

    def test_build_targets_raises(
        self,
        config: OptimizationConfig,
    ) -> None:
        """build_targets raises NotImplementedError."""
        session = OptimizationSession(config)
        with pytest.raises(NotImplementedError):
            session.build_targets()

    def test_measure_raises(
        self,
        config: OptimizationConfig,
    ) -> None:
        """measure raises NotImplementedError."""
        session = OptimizationSession(config)
        with pytest.raises(NotImplementedError):
            session.measure()

    def test_evaluate_raises(
        self,
        config: OptimizationConfig,
    ) -> None:
        """evaluate raises NotImplementedError."""
        session = OptimizationSession(config)
        with pytest.raises(NotImplementedError):
            session.evaluate()


class TestOptimizeFunction:
    """optimize_function facade."""

    def test_returns_none_for_missing_function(
        self,
        config: OptimizationConfig,
        sample_file: Path,
    ) -> None:
        """Returns None when function_name is not found."""
        assert optimize_function(config, sample_file, "nonexistent") is None

    def test_returns_none_for_empty_file(
        self,
        config: OptimizationConfig,
        tmp_path: Path,
    ) -> None:
        """Returns None when file has no functions."""
        empty = tmp_path / "empty.py"
        empty.write_text("x = 1\n", encoding="utf-8")
        assert optimize_function(config, empty, "foo") is None
