"""Tests for the @track decorator."""

import warnings
from pathlib import Path
from types import SimpleNamespace
from typing import Optional

import pytest

import agentbudget.tracker as tracker_module
from agentbudget.decorators import track
from agentbudget.exceptions import BudgetExceededError, BudgetExceededWarning


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def make_openai_response(
    model: str = "gpt-4o",
    prompt_tokens: int = 100,
    completion_tokens: int = 50,
) -> SimpleNamespace:
    """SimpleNamespace mimicking an OpenAI ChatCompletion (detected via model prefix)."""
    return SimpleNamespace(
        model=model,
        usage=SimpleNamespace(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens),
    )


def make_anthropic_response(
    model: str = "claude-3-5-sonnet-20241022",
    input_tokens: int = 100,
    output_tokens: int = 50,
) -> SimpleNamespace:
    """SimpleNamespace mimicking an Anthropic Message (detected via usage.input_tokens)."""
    return SimpleNamespace(
        model=model,
        usage=SimpleNamespace(input_tokens=input_tokens, output_tokens=output_tokens),
    )


# ---------------------------------------------------------------------------
# 1. test_track_happy_path
# ---------------------------------------------------------------------------

class TestTrackHappyPath:
    def test_track_happy_path(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """@track wraps function, returns the original response, and tracks cost."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response(model="gpt-4o", prompt_tokens=100, completion_tokens=50)

        @track(label="happy-path", verbose=False)
        def call_llm() -> SimpleNamespace:
            return mock_response

        result = call_llm()

        # The decorator must return the original response object unchanged
        assert result is mock_response

        # The session must have been persisted to storage
        import json
        assert (tmp_path / "sessions.json").exists()
        with open(tmp_path / "sessions.json") as f:
            data = json.load(f)
        assert len(data) == 1
        session = data[0]
        assert session["total_tokens"] == 100 + 50
        # gpt-4o: $2.50/M input, $10.00/M output
        # 100 input => $0.000250, 50 output => $0.000500 => total $0.000750
        expected_cost = (100 / 1_000_000) * 2.50 + (50 / 1_000_000) * 10.00
        assert abs(session["total_cost"] - expected_cost) < 1e-9

    def test_track_happy_path_anthropic(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """@track correctly handles Anthropic responses as well."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_anthropic_response(
            model="claude-3-5-sonnet-20241022",
            input_tokens=200,
            output_tokens=80,
        )

        @track(label="anthropic-step", verbose=False)
        def call_claude() -> SimpleNamespace:
            return mock_response

        result = call_claude()

        assert result is mock_response

        import json
        with open(tmp_path / "sessions.json") as f:
            data = json.load(f)
        assert data[0]["total_tokens"] == 200 + 80


# ---------------------------------------------------------------------------
# 2. test_track_default_label
# ---------------------------------------------------------------------------

class TestTrackDefaultLabel:
    def test_track_default_label_uses_func_name(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """When label is empty, the step label defaults to the function name."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track(verbose=False)  # no label argument
        def my_llm_call() -> SimpleNamespace:
            return mock_response

        my_llm_call()

        import json
        with open(tmp_path / "sessions.json") as f:
            data = json.load(f)
        step = data[0]["steps"][0]
        assert step["label"] == "my_llm_call"

    def test_track_empty_string_label_uses_func_name(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """Explicitly passing label='' also falls back to the function name."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track(label="", verbose=False)
        def generate_text() -> SimpleNamespace:
            return mock_response

        generate_text()

        import json
        with open(tmp_path / "sessions.json") as f:
            data = json.load(f)
        step = data[0]["steps"][0]
        assert step["label"] == "generate_text"


# ---------------------------------------------------------------------------
# 3. test_track_custom_label
# ---------------------------------------------------------------------------

class TestTrackCustomLabel:
    def test_track_custom_label_appears_in_persisted_data(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """A custom label is stored in the step data."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track(label="my-custom-step", verbose=False)
        def call_api() -> SimpleNamespace:
            return mock_response

        call_api()

        import json
        with open(tmp_path / "sessions.json") as f:
            data = json.load(f)
        step = data[0]["steps"][0]
        assert step["label"] == "my-custom-step"

    def test_track_custom_label_appears_in_summary_output(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
    ) -> None:
        """Custom label is visible in the printed summary when verbose=True."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track(label="summarize-document", verbose=True)
        def call_api() -> SimpleNamespace:
            return mock_response

        call_api()

        out = capsys.readouterr().out
        assert "summarize-document" in out


# ---------------------------------------------------------------------------
# 4. test_track_returns_none_raises
# ---------------------------------------------------------------------------

class TestTrackReturnsNoneRaises:
    def test_returns_none_raises_value_error(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """A function returning None must raise ValueError with a clear message."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        @track(verbose=False)
        def bad_llm_call() -> Optional[SimpleNamespace]:
            return None

        with pytest.raises(ValueError, match="returned None"):
            bad_llm_call()

    def test_returns_none_error_message_contains_function_name(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """The ValueError message should identify the offending function by name."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        @track(verbose=False)
        def my_broken_function() -> Optional[SimpleNamespace]:
            return None

        with pytest.raises(ValueError, match="my_broken_function"):
            my_broken_function()

    def test_returns_none_does_not_persist_session(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """When the function returns None the session should not be saved."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        @track(verbose=False)
        def none_returning() -> Optional[SimpleNamespace]:
            return None

        with pytest.raises(ValueError):
            none_returning()

        # No session was tracked, so nothing is persisted
        assert not (tmp_path / "sessions.json").exists()


# ---------------------------------------------------------------------------
# 5. test_track_budget_warning
# ---------------------------------------------------------------------------

class TestTrackBudgetWarning:
    def test_budget_warning_emitted_when_exceeded(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """budget=0.00001 with a large response triggers BudgetExceededWarning."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        # gpt-4o: $2.50/M input tokens — 1M tokens will far exceed $0.00001
        mock_response = make_openai_response(
            model="gpt-4o", prompt_tokens=1_000_000, completion_tokens=0
        )

        @track(budget=0.00001, strict=False, verbose=False)
        def expensive_call() -> SimpleNamespace:
            return mock_response

        with pytest.warns(BudgetExceededWarning, match="Budget of"):
            result = expensive_call()

        # Despite the warning the response is still returned
        assert result is mock_response

    def test_budget_warning_contains_dollar_amount(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """The BudgetExceededWarning message must reference the budget value."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response(
            model="gpt-4o", prompt_tokens=1_000_000, completion_tokens=0
        )

        @track(budget=0.00001, strict=False, verbose=False)
        def call_api() -> SimpleNamespace:
            return mock_response

        with pytest.warns(BudgetExceededWarning) as record:
            call_api()

        assert "$0.00" in str(record[0].message)

    def test_no_warning_when_within_budget(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """No BudgetExceededWarning when cost stays below budget."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        # gpt-4o-mini with 1 token each is essentially free
        mock_response = make_openai_response(model="gpt-4o-mini", prompt_tokens=1, completion_tokens=1)

        @track(budget=100.0, strict=False, verbose=False)
        def cheap_call() -> SimpleNamespace:
            return mock_response

        with warnings.catch_warnings():
            warnings.simplefilter("error", BudgetExceededWarning)
            cheap_call()  # must not raise


# ---------------------------------------------------------------------------
# 6. test_track_budget_strict
# ---------------------------------------------------------------------------

class TestTrackBudgetStrict:
    def test_budget_strict_raises_budget_exceeded_error(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """budget=0.00001, strict=True must raise BudgetExceededError."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response(
            model="gpt-4o", prompt_tokens=1_000_000, completion_tokens=0
        )

        @track(budget=0.00001, strict=True, verbose=False)
        def very_expensive_call() -> SimpleNamespace:
            return mock_response

        with pytest.raises(BudgetExceededError, match="Budget of"):
            very_expensive_call()

    def test_budget_strict_error_message_contains_budget(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """BudgetExceededError message must include the configured budget."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response(
            model="gpt-4o", prompt_tokens=1_000_000, completion_tokens=0
        )

        @track(budget=0.00001, strict=True, verbose=False)
        def expensive_call() -> SimpleNamespace:
            return mock_response

        with pytest.raises(BudgetExceededError) as exc_info:
            expensive_call()

        assert "$0.00" in str(exc_info.value)

    def test_strict_false_does_not_raise_error(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """With strict=False the BudgetExceededError must never be raised."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response(
            model="gpt-4o", prompt_tokens=1_000_000, completion_tokens=0
        )

        @track(budget=0.00001, strict=False, verbose=False)
        def call_api() -> SimpleNamespace:
            return mock_response

        # Only a warning, not an error
        with pytest.warns(BudgetExceededWarning):
            result = call_api()
        assert result is mock_response


# ---------------------------------------------------------------------------
# 7. test_track_verbose_false
# ---------------------------------------------------------------------------

class TestTrackVerboseFalse:
    def test_verbose_false_suppresses_all_output(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
    ) -> None:
        """verbose=False must produce no stdout output."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track(label="quiet", verbose=False)
        def call_api() -> SimpleNamespace:
            return mock_response

        call_api()

        out, err = capsys.readouterr()
        assert out == ""

    def test_verbose_true_produces_summary_output(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
    ) -> None:
        """verbose=True (default) must print session summary to stdout."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track(label="loud", verbose=True)
        def call_api() -> SimpleNamespace:
            return mock_response

        call_api()

        out = capsys.readouterr().out
        assert "Session:" in out
        assert "TOTAL" in out

    def test_verbose_defaults_to_true(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
    ) -> None:
        """When verbose is not specified the default is True — output must appear."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track()  # no verbose arg at all
        def call_api() -> SimpleNamespace:
            return mock_response

        call_api()

        out = capsys.readouterr().out
        assert "Session:" in out


# ---------------------------------------------------------------------------
# 8. test_track_functools_wraps
# ---------------------------------------------------------------------------

class TestTrackFunctoolsWraps:
    def test_wrapper_preserves_function_name(self) -> None:
        """The decorated function must retain the original __name__."""

        @track(verbose=False)
        def my_special_function() -> SimpleNamespace:
            """My docstring."""
            return make_openai_response()

        assert my_special_function.__name__ == "my_special_function"

    def test_wrapper_preserves_docstring(self) -> None:
        """The decorated function must retain the original __doc__."""

        @track(verbose=False)
        def documented_function() -> SimpleNamespace:
            """This is the original docstring."""
            return make_openai_response()

        assert documented_function.__doc__ == "This is the original docstring."

    def test_wrapper_preserves_none_docstring(self) -> None:
        """Functions without a docstring should have __doc__ == None after wrapping."""

        @track(verbose=False)
        def undocumented() -> SimpleNamespace:
            return make_openai_response()

        assert undocumented.__doc__ is None

    def test_wrapper_preserves_name_with_custom_label(self) -> None:
        """Custom label must not affect __name__ preservation."""

        @track(label="override", verbose=False)
        def original_name() -> SimpleNamespace:
            return make_openai_response()

        assert original_name.__name__ == "original_name"


# ---------------------------------------------------------------------------
# Additional edge cases
# ---------------------------------------------------------------------------

class TestTrackEdgeCases:
    def test_track_multiple_calls_each_gets_own_session(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """Each invocation of the decorated function must create a separate session."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track(verbose=False)
        def repeated_call() -> SimpleNamespace:
            return mock_response

        repeated_call()
        repeated_call()
        repeated_call()

        import json
        with open(tmp_path / "sessions.json") as f:
            data = json.load(f)
        assert len(data) == 3

    def test_track_session_name_is_function_name(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """The CostSession name is always the function name, regardless of label."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track(label="some-label", verbose=False)
        def my_pipeline_step() -> SimpleNamespace:
            return mock_response

        my_pipeline_step()

        import json
        with open(tmp_path / "sessions.json") as f:
            data = json.load(f)
        assert data[0]["name"] == "my_pipeline_step"

    def test_track_passes_args_and_kwargs_to_wrapped_function(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """The decorator must pass all positional and keyword arguments to the original function."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        received: dict = {}

        @track(verbose=False)
        def function_with_args(prompt: str, temperature: float = 0.7) -> SimpleNamespace:
            received["prompt"] = prompt
            received["temperature"] = temperature
            return make_openai_response()

        function_with_args("hello world", temperature=0.5)

        assert received["prompt"] == "hello world"
        assert received["temperature"] == 0.5

    def test_track_no_budget_never_warns(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """Without a budget, no BudgetExceededWarning is ever emitted."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        # Even a very expensive call should not warn without a budget
        mock_response = make_openai_response(
            model="gpt-4o", prompt_tokens=10_000_000, completion_tokens=10_000_000
        )

        @track(verbose=False)
        def very_expensive() -> SimpleNamespace:
            return mock_response

        with warnings.catch_warnings():
            warnings.simplefilter("error", BudgetExceededWarning)
            very_expensive()  # must not raise

    def test_track_decorator_can_be_applied_multiple_times_independently(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """Two different functions can both use @track independently."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        mock_response = make_openai_response()

        @track(label="step-alpha", verbose=False)
        def alpha() -> SimpleNamespace:
            return mock_response

        @track(label="step-beta", verbose=False)
        def beta() -> SimpleNamespace:
            return mock_response

        alpha()
        beta()

        import json
        with open(tmp_path / "sessions.json") as f:
            data = json.load(f)
        assert len(data) == 2
        step_labels = [session["steps"][0]["label"] for session in data]
        assert "step-alpha" in step_labels
        assert "step-beta" in step_labels
