"""Tests for CostSession context manager."""

import warnings
import pytest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import agentbudget.tracker as tracker_module
from agentbudget.session import CostSession
from agentbudget.exceptions import BudgetExceededWarning, BudgetExceededError


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

def make_openai_response(
    model: str = "gpt-4o",
    prompt_tokens: int = 800,
    completion_tokens: int = 443,
) -> SimpleNamespace:
    """SimpleNamespace mimicking an OpenAI ChatCompletion (model-prefix detected)."""
    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 (input_tokens fingerprint)."""
    return SimpleNamespace(
        model=model,
        usage=SimpleNamespace(input_tokens=input_tokens, output_tokens=output_tokens),
    )


# ---------------------------------------------------------------------------
# __enter__ / __exit__ basics
# ---------------------------------------------------------------------------

class TestCostSessionContextManager:
    def test_enter_returns_session(self) -> None:
        session = CostSession(name="test-session", verbose=False)
        result = session.__enter__()
        assert result is session

    def test_with_block_returns_session(self) -> None:
        with CostSession(name="test-session", verbose=False) as session:
            assert isinstance(session, CostSession)

    def test_empty_session_exit_does_nothing(self, capsys: pytest.CaptureFixture, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """Empty session: __exit__ should not write to storage or print."""
        # Redirect storage so we can confirm nothing is written
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        with CostSession(name="empty", verbose=True) as session:
            pass  # no track() calls

        # Storage file must not exist — nothing was persisted
        assert not (tmp_path / "sessions.json").exists()
        out = capsys.readouterr().out
        assert out == ""


# ---------------------------------------------------------------------------
# track() with OpenAI response
# ---------------------------------------------------------------------------

class TestCostSessionTrackOpenAI:
    def test_track_openai_accumulates_cost(self) -> None:
        """gpt-4o: input=$2.50/M, output=$10/M.
        800 input → $0.002000, 443 output → $0.004430  => total $0.006430
        """
        response = make_openai_response(model="gpt-4o", prompt_tokens=800, completion_tokens=443)
        with CostSession(name="test", verbose=False) as session:
            session.track(response, label="gen")

        assert session.total_tokens == 800 + 443
        expected_cost = (800 / 1_000_000) * 2.50 + (443 / 1_000_000) * 10.00
        assert abs(session.total_cost - expected_cost) < 1e-9

    def test_track_openai_step_added(self) -> None:
        response = make_openai_response()
        with CostSession(name="test", verbose=False) as session:
            session.track(response, label="step-a")

        assert len(session._steps) == 1
        assert session._steps[0]["label"] == "step-a"
        assert session._steps[0]["provider"] == "openai"
        assert session._steps[0]["model"] == "gpt-4o"

    def test_track_multiple_openai_steps(self) -> None:
        r1 = make_openai_response(prompt_tokens=100, completion_tokens=50)
        r2 = make_openai_response(prompt_tokens=200, completion_tokens=100)
        with CostSession(name="test", verbose=False) as session:
            session.track(r1, label="step-1")
            session.track(r2, label="step-2")

        assert session.total_tokens == (100 + 50) + (200 + 100)
        assert len(session._steps) == 2


# ---------------------------------------------------------------------------
# track() with Anthropic response
# ---------------------------------------------------------------------------

class TestCostSessionTrackAnthropic:
    def test_track_anthropic_accumulates_cost(self) -> None:
        """claude-3-5-sonnet-20241022: input=$3/M, output=$15/M.
        100 input → $0.000300, 50 output → $0.000750  => total $0.001050
        """
        response = make_anthropic_response(
            model="claude-3-5-sonnet-20241022",
            input_tokens=100,
            output_tokens=50,
        )
        with CostSession(name="test", verbose=False) as session:
            session.track(response, label="gen")

        assert session.total_tokens == 150
        expected_cost = (100 / 1_000_000) * 3.00 + (50 / 1_000_000) * 15.00
        assert abs(session.total_cost - expected_cost) < 1e-9

    def test_track_anthropic_step_added(self) -> None:
        response = make_anthropic_response()
        with CostSession(name="test", verbose=False) as session:
            session.track(response, label="summarize")

        assert len(session._steps) == 1
        assert session._steps[0]["label"] == "summarize"
        assert session._steps[0]["provider"] == "anthropic"


# ---------------------------------------------------------------------------
# Default label auto-increment
# ---------------------------------------------------------------------------

class TestCostSessionDefaultLabel:
    def test_empty_label_defaults_to_step_1(self) -> None:
        response = make_openai_response()
        with CostSession(name="test", verbose=False) as session:
            session.track(response)  # no label

        assert session._steps[0]["label"] == "step-1"

    def test_two_empty_labels_default_to_step_1_step_2(self) -> None:
        r1 = make_openai_response()
        r2 = make_openai_response()
        with CostSession(name="test", verbose=False) as session:
            session.track(r1)
            session.track(r2)

        assert session._steps[0]["label"] == "step-1"
        assert session._steps[1]["label"] == "step-2"


# ---------------------------------------------------------------------------
# Budget enforcement
# ---------------------------------------------------------------------------

class TestCostSessionBudget:
    def test_budget_warning_when_exceeded_strict_false(self) -> None:
        """By default (strict=False) exceeding budget emits BudgetExceededWarning."""
        response = make_openai_response(
            model="gpt-4o", prompt_tokens=1_000_000, completion_tokens=1_000_000
        )
        with CostSession(name="test", budget=0.001, strict=False, verbose=False) as session:
            with pytest.warns(BudgetExceededWarning, match="Budget of"):
                session.track(response)

    def test_budget_error_when_exceeded_strict_true(self) -> None:
        """With strict=True, exceeding budget raises BudgetExceededError."""
        response = make_openai_response(
            model="gpt-4o", prompt_tokens=1_000_000, completion_tokens=1_000_000
        )
        with pytest.raises(BudgetExceededError, match="Budget of"):
            with CostSession(name="test", budget=0.001, strict=True, verbose=False) as session:
                session.track(response)

    def test_no_warning_when_within_budget(self) -> None:
        response = make_openai_response(model="gpt-4o-mini", prompt_tokens=1, completion_tokens=1)
        with CostSession(name="test", budget=100.0, strict=False, verbose=False) as session:
            with warnings.catch_warnings():
                warnings.simplefilter("error", BudgetExceededWarning)
                session.track(response)  # should not raise

    def test_budget_message_contains_current_cost(self) -> None:
        """Warning/error message should reference the budget value."""
        response = make_openai_response(
            model="gpt-4o", prompt_tokens=1_000_000, completion_tokens=1_000_000
        )
        with CostSession(name="test", budget=0.001, strict=False, verbose=False) as session:
            with pytest.warns(BudgetExceededWarning) as record:
                session.track(response)

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

    def test_no_budget_no_warning(self) -> None:
        """Without budget set, no warning is ever emitted."""
        response = make_openai_response(model="gpt-4o", prompt_tokens=5_000_000, completion_tokens=5_000_000)
        with CostSession(name="test", verbose=False) as session:
            with warnings.catch_warnings():
                warnings.simplefilter("error", BudgetExceededWarning)
                session.track(response)  # should not raise


# ---------------------------------------------------------------------------
# verbose=False suppresses output
# ---------------------------------------------------------------------------

class TestCostSessionVerbose:
    def test_verbose_false_suppresses_stdout(self, capsys: pytest.CaptureFixture) -> None:
        response = make_openai_response()
        with CostSession(name="quiet-session", verbose=False) as session:
            session.track(response)

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

    def test_verbose_true_prints_summary(self, capsys: pytest.CaptureFixture, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        response = make_openai_response()
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        with CostSession(name="loud-session", verbose=True) as session:
            session.track(response)

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


# ---------------------------------------------------------------------------
# summary() format
# ---------------------------------------------------------------------------

class TestCostSessionSummary:
    def test_summary_contains_session_name(self) -> None:
        response = make_openai_response()
        session = CostSession(name="my-pipeline", verbose=False)
        session.track(response, label="gen")
        result = session.summary()

        assert "Session:" in result
        assert "my-pipeline" in result

    def test_summary_contains_total(self) -> None:
        response = make_openai_response()
        session = CostSession(name="test", verbose=False)
        session.track(response, label="gen")
        result = session.summary()

        assert "TOTAL" in result

    def test_summary_contains_separator(self) -> None:
        response = make_openai_response()
        session = CostSession(name="test", verbose=False)
        session.track(response, label="gen")
        result = session.summary()

        assert "─" in result

    def test_summary_shows_step_label(self) -> None:
        response = make_openai_response()
        session = CostSession(name="test", verbose=False)
        session.track(response, label="my-custom-label")
        result = session.summary()

        assert "my-custom-label" in result

    def test_summary_shows_model(self) -> None:
        response = make_openai_response(model="gpt-4o")
        session = CostSession(name="test", verbose=False)
        session.track(response, label="gen")
        result = session.summary()

        assert "gpt-4o" in result

    def test_summary_shows_budget_info_when_set(self) -> None:
        response = make_openai_response(model="gpt-4o-mini", prompt_tokens=1, completion_tokens=1)
        session = CostSession(name="test", budget=1.00, verbose=False)
        session.track(response, label="gen")
        result = session.summary()

        assert "Budget" in result

    def test_summary_multiple_steps(self) -> None:
        r1 = make_openai_response(prompt_tokens=100, completion_tokens=50)
        r2 = make_anthropic_response(input_tokens=200, output_tokens=100)
        session = CostSession(name="multi", verbose=False)
        session.track(r1, label="step-1")
        session.track(r2, label="step-2")
        result = session.summary()

        assert "step-1" in result
        assert "step-2" in result


# ---------------------------------------------------------------------------
# total_cost and total_tokens properties
# ---------------------------------------------------------------------------

class TestCostSessionProperties:
    def test_total_cost_zero_on_init(self) -> None:
        session = CostSession(name="test", verbose=False)
        assert session.total_cost == 0.0

    def test_total_tokens_zero_on_init(self) -> None:
        session = CostSession(name="test", verbose=False)
        assert session.total_tokens == 0

    def test_total_tokens_accumulates(self) -> None:
        r1 = make_openai_response(prompt_tokens=100, completion_tokens=50)
        r2 = make_openai_response(prompt_tokens=200, completion_tokens=80)
        session = CostSession(name="test", verbose=False)
        session.track(r1)
        session.track(r2)

        assert session.total_tokens == 100 + 50 + 200 + 80

    def test_total_cost_accumulates(self) -> None:
        r1 = make_openai_response(model="gpt-4o-mini", prompt_tokens=1_000_000, completion_tokens=0)
        r2 = make_openai_response(model="gpt-4o-mini", prompt_tokens=1_000_000, completion_tokens=0)
        session = CostSession(name="test", verbose=False)
        session.track(r1)
        session.track(r2)

        # gpt-4o-mini: $0.15 / M input => 2 * $0.15 = $0.30
        assert abs(session.total_cost - 0.30) < 1e-9


# ---------------------------------------------------------------------------
# __exit__ calls CostTracker.save_session
# ---------------------------------------------------------------------------

class TestCostSessionPersistence:
    def test_exit_calls_save_session_when_steps_exist(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """__exit__ should persist the session when steps exist."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        response = make_openai_response()
        with CostSession(name="persist-test", verbose=False) as session:
            session.track(response)

        # Confirm the session was written to disk
        import json
        assert (tmp_path / "sessions.json").exists()
        with open(tmp_path / "sessions.json") as f:
            data = json.load(f)
        assert len(data) == 1
        assert data[0]["name"] == "persist-test"
        assert "total_cost" in data[0]
        assert "total_tokens" in data[0]
        assert "steps" in data[0]

    def test_exit_does_not_call_save_session_when_empty(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        """__exit__ on an empty session should not create the storage file."""
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tracker_module.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        with CostSession(name="empty", verbose=False) as session:
            pass

        assert not (tmp_path / "sessions.json").exists()


# ---------------------------------------------------------------------------
# Unknown provider raises ValueError
# ---------------------------------------------------------------------------

class TestCostSessionUnknownProvider:
    def test_unknown_response_raises_value_error(self) -> None:
        response = SimpleNamespace(model="llama-2-7b")
        session = CostSession(name="test", verbose=False)
        with pytest.raises(ValueError):
            session.track(response)

    def test_explicit_unknown_provider_raises_value_error(self) -> None:
        response = make_openai_response()
        session = CostSession(name="test", verbose=False, provider="foobar")
        with pytest.raises(ValueError, match="Unknown provider"):
            session.track(response)


def make_gemini_response(
    model: str = "gemini-2.0-flash",
    prompt_token_count: int = 200,
    candidates_token_count: int = 100,
) -> SimpleNamespace:
    """SimpleNamespace mimicking a Gemini response (usage_metadata fingerprint)."""
    return SimpleNamespace(
        model=model,
        usage_metadata=SimpleNamespace(
            prompt_token_count=prompt_token_count,
            candidates_token_count=candidates_token_count,
        ),
    )


class TestCostSessionTrackGemini:
    def test_gemini_response_tracked(self) -> None:
        response = make_gemini_response(model="gemini-2.0-flash", prompt_token_count=500, candidates_token_count=250)
        with CostSession(name="test", verbose=False) as session:
            session.track(response, label="gen")
        assert len(session._steps) == 1
        assert session._steps[0]["model"] == "gemini-2.0-flash"
        assert session._steps[0]["provider"] == "gemini"
        assert session._steps[0]["input_tokens"] == 500
        assert session._steps[0]["output_tokens"] == 250

    def test_gemini_explicit_provider(self) -> None:
        response = make_gemini_response()
        session = CostSession(name="test", verbose=False, provider="gemini")
        session.track(response)
        assert session._steps[0]["provider"] == "gemini"

    def test_model_override_replaces_response_model(self) -> None:
        """model= kwarg overrides the model extracted from the response."""
        response = SimpleNamespace(
            usage_metadata=SimpleNamespace(prompt_token_count=100, candidates_token_count=50)
        )  # no model attr — would fallback to "gemini"
        session = CostSession(name="test", verbose=False)
        session.track(response, model="gemini-2.5-pro")
        assert session._steps[0]["model"] == "gemini-2.5-pro"

    def test_model_override_works_for_openai_too(self) -> None:
        """model= override also works for non-Gemini providers."""
        response = make_openai_response(model="gpt-4o")
        session = CostSession(name="test", verbose=False)
        session.track(response, model="gpt-4o-mini")
        assert session._steps[0]["model"] == "gpt-4o-mini"
