"""Tests for cost tracking."""
import pytest

from logos_adk.agent import Agent
from logos_adk.cost import CostReport, CostTracker, DEFAULT_RATES
from logos_adk.observe import Metric, Span
from logos_adk.providers.mock import MockProvider
from logos_adk.types import Response, Usage

from datetime import datetime, UTC


def _make_span(*, model="claude-sonnet-4-20250514", input_tokens=100, output_tokens=50,
               agent_name="bot", session_id=None, name="LLM.generate"):
    now = datetime.now(UTC)
    return Span(
        name=name,
        trace_id="trace-1",
        span_id="span-1",
        parent_span_id=None,
        started_at=now,
        ended_at=now,
        duration_ms=100,
        agent_name=agent_name,
        session_id=session_id,
        attributes={"input_tokens": input_tokens, "output_tokens": output_tokens, "model": model},
    )


class TestCostReport:
    def test_dataclass_fields(self):
        r = CostReport(total_usd=0.1, input_usd=0.05, output_usd=0.05,
                        input_tokens=100, output_tokens=50, calls=1)
        assert r.total_usd == 0.1
        assert r.calls == 1

    def test_defaults(self):
        r = CostReport()
        assert r.total_usd == 0.0
        assert r.calls == 0


class TestCostTrackerSpan:
    @pytest.mark.asyncio
    async def test_tracks_single_call(self):
        tracker = CostTracker()
        span = _make_span(model="claude-sonnet-4-20250514", input_tokens=1000, output_tokens=500)
        await tracker.on_span(span)

        report = tracker.total
        assert report.input_tokens == 1000
        assert report.output_tokens == 500
        assert report.calls == 1
        assert report.total_usd > 0

    @pytest.mark.asyncio
    async def test_tracks_multiple_calls(self):
        tracker = CostTracker()
        await tracker.on_span(_make_span(input_tokens=100, output_tokens=50))
        await tracker.on_span(_make_span(input_tokens=200, output_tokens=100))

        report = tracker.total
        assert report.input_tokens == 300
        assert report.output_tokens == 150
        assert report.calls == 2

    @pytest.mark.asyncio
    async def test_ignores_non_llm_spans(self):
        tracker = CostTracker()
        await tracker.on_span(_make_span(name="Tool.search"))
        assert tracker.total.calls == 0

    @pytest.mark.asyncio
    async def test_handles_llm_stream_spans(self):
        tracker = CostTracker()
        await tracker.on_span(_make_span(name="LLM.stream", input_tokens=100, output_tokens=50))
        assert tracker.total.calls == 1

    @pytest.mark.asyncio
    async def test_custom_rates(self):
        tracker = CostTracker(rates={"custom-model": (1.0, 2.0)})
        await tracker.on_span(_make_span(model="custom-model", input_tokens=1_000_000, output_tokens=1_000_000))
        report = tracker.total
        assert abs(report.input_usd - 1.0) < 0.001
        assert abs(report.output_usd - 2.0) < 0.001

    @pytest.mark.asyncio
    async def test_unknown_model_zero_cost(self):
        tracker = CostTracker()
        await tracker.on_span(_make_span(model="unknown-model-xyz"))
        report = tracker.total
        assert report.calls == 1
        assert report.total_usd == 0.0
        assert report.input_tokens == 100

    @pytest.mark.asyncio
    async def test_report_by_agent_name(self):
        tracker = CostTracker()
        await tracker.on_span(_make_span(agent_name="analyst", input_tokens=100))
        await tracker.on_span(_make_span(agent_name="copywriter", input_tokens=200))

        report = tracker.report(agent_name="analyst")
        assert report.input_tokens == 100
        assert report.calls == 1

    @pytest.mark.asyncio
    async def test_report_by_session_id(self):
        tracker = CostTracker()
        await tracker.on_span(_make_span(session_id="sess-1", input_tokens=100))
        await tracker.on_span(_make_span(session_id="sess-2", input_tokens=200))

        report = tracker.report(session_id="sess-1")
        assert report.input_tokens == 100

    @pytest.mark.asyncio
    async def test_reset_clears(self):
        tracker = CostTracker()
        await tracker.on_span(_make_span())
        assert tracker.total.calls == 1
        tracker.reset()
        assert tracker.total.calls == 0

    @pytest.mark.asyncio
    async def test_glob_pattern_matching(self):
        tracker = CostTracker(rates={"test-*": (10.0, 20.0)})
        await tracker.on_span(_make_span(model="test-large", input_tokens=1_000_000, output_tokens=1_000_000))
        assert abs(tracker.total.input_usd - 10.0) < 0.001

    @pytest.mark.asyncio
    async def test_on_metric_is_noop(self):
        tracker = CostTracker()
        await tracker.on_metric(Metric(name="test", value=1.0))
        assert tracker.total.calls == 0

    @pytest.mark.asyncio
    async def test_flush_is_noop(self):
        tracker = CostTracker()
        await tracker.flush()  # should not raise


class TestDefaultRates:
    def test_has_claude_rates(self):
        assert any("claude" in k for k in DEFAULT_RATES)

    def test_has_gpt_rates(self):
        assert any("gpt" in k for k in DEFAULT_RATES)


class TestCostTrackerIntegration:
    @pytest.mark.asyncio
    async def test_with_agent_run(self):
        tracker = CostTracker()
        provider = MockProvider(responses=[
            Response(content="hi", tool_calls=[], usage=Usage(100, 50), raw={}),
        ])
        agent = Agent("bot", model=provider).observe(tracker)
        await agent.run("hello")
        assert tracker.total.calls == 1
        assert tracker.total.input_tokens == 100
