from unittest.mock import MagicMock

from agentcost.client import EventClient
from agentcost.event import CostEvent
from agentcost.pricing import PricingDB
from agentcost.tracker import Tracker, RunTracker, _active_run


class TestRunTracker:
    def _make_tracker(self):
        client = MagicMock(spec=EventClient)
        pricing = PricingDB()
        return RunTracker(
            client=client,
            pricing=pricing,
            agent_id="test-agent",
            run_id="run-123",
            metadata={"task": "test"},
        ), client

    def test_record_with_explicit_cost(self):
        tracker, client = self._make_tracker()
        event = tracker.record("stripe_api", cost=0.02)
        assert event.cost_usd == 0.02
        assert event.component_name == "stripe_api"
        assert event.event_type == "tool_call"
        client.enqueue.assert_called_once()

    def test_record_resolves_cost_from_pricing(self):
        tracker, client = self._make_tracker()
        event = tracker.record("stripe_api")
        assert event.cost_usd == 0.02  # default pricing

    def test_record_llm(self):
        tracker, client = self._make_tracker()
        event = tracker.record_llm("gpt-4o", tokens_in=500, tokens_out=200)
        assert event.event_type == "llm_call"
        assert event.tokens_in == 500
        assert event.tokens_out == 200
        assert event.cost_usd > 0

    def test_record_llm_with_explicit_cost(self):
        tracker, client = self._make_tracker()
        event = tracker.record_llm("gpt-4o", tokens_in=500, tokens_out=200, cost=0.05)
        assert event.cost_usd == 0.05

    def test_total_cost(self):
        tracker, _ = self._make_tracker()
        tracker.record("stripe_api", cost=0.02)
        tracker.record("sendgrid", cost=0.001)
        assert tracker.total_cost == 0.021

    def test_events_list(self):
        tracker, _ = self._make_tracker()
        tracker.record("a", cost=0.01)
        tracker.record("b", cost=0.02)
        events = tracker.events
        assert len(events) == 2
        # Should be a copy
        events.clear()
        assert len(tracker.events) == 2

    def test_metadata_merging(self):
        tracker, _ = self._make_tracker()
        event = tracker.record("tool", cost=0.01, metadata={"extra": "data"})
        assert event.metadata["task"] == "test"
        assert event.metadata["extra"] == "data"


class TestTracker:
    def test_run_context_manager(self):
        client = MagicMock(spec=EventClient)
        pricing = PricingDB()
        tracker = Tracker(
            client=client, pricing=pricing, agent_id="my-agent"
        )

        with tracker.run(task="refund") as run:
            run.record("stripe_api", cost=0.02)
            assert run.total_cost == 0.02

        # Should have enqueued the event + run summary
        assert client.enqueue.call_count == 2
        summary_event = client.enqueue.call_args_list[1][0][0]
        assert summary_event.event_type == "run_end"
        assert summary_event.cost_usd == 0.02

    def test_run_sets_active_context(self):
        client = MagicMock(spec=EventClient)
        pricing = PricingDB()
        tracker = Tracker(
            client=client, pricing=pricing, agent_id="my-agent"
        )

        assert _active_run.get(None) is None
        with tracker.run(task="test") as run:
            assert _active_run.get(None) is run
        assert _active_run.get(None) is None

    def test_run_summary_includes_duration(self):
        client = MagicMock(spec=EventClient)
        pricing = PricingDB()
        tracker = Tracker(
            client=client, pricing=pricing, agent_id="my-agent"
        )

        with tracker.run(task="test") as run:
            pass

        summary = client.enqueue.call_args_list[0][0][0]
        assert summary.duration_ms is not None
        assert summary.duration_ms >= 0
