"""Tests for the budget engine."""

from __future__ import annotations

import time

from agent_scout.events import EventBus, EventType
from agent_scout.llm.budget import (
    BudgetConfig,
    BudgetEngine,
    BudgetExceededError,
)


def _make_engine(
    task_limit: float = 1.0,
    daily_limit: float = 10.0,
    monthly_limit: float = 50.0,
    enforce: bool = True,
    bus: EventBus | None = None,
) -> BudgetEngine:
    config = BudgetConfig(
        task_limit=task_limit,
        daily_limit=daily_limit,
        monthly_limit=monthly_limit,
        enforce=enforce,
    )
    return BudgetEngine(config=config, event_bus=bus)


# --- Basic tracking ---


def test_record_cost():
    engine = _make_engine()
    engine.record_cost("claude", 100, 50, 0.01, task_id="t1")
    assert engine.get_task_cost("t1") == 0.01


def test_get_daily_cost():
    engine = _make_engine()
    engine.record_cost("claude", 100, 50, 0.05)
    engine.record_cost("claude", 200, 100, 0.10)
    assert abs(engine.get_daily_cost() - 0.15) < 1e-9


def test_get_monthly_cost():
    engine = _make_engine()
    engine.record_cost("claude", 100, 50, 1.0)
    assert engine.get_monthly_cost() == 1.0


def test_get_total_tokens():
    engine = _make_engine()
    engine.record_cost("claude", 100, 50, 0.01, task_id="t1")
    engine.record_cost("claude", 200, 100, 0.02, task_id="t1")
    tokens = engine.get_total_tokens("t1")
    assert tokens["input_tokens"] == 300
    assert tokens["output_tokens"] == 150


def test_get_total_tokens_all():
    engine = _make_engine()
    engine.record_cost("claude", 100, 50, 0.01, task_id="t1")
    engine.record_cost("gpt-4", 200, 100, 0.02, task_id="t2")
    tokens = engine.get_total_tokens()
    assert tokens["input_tokens"] == 300
    assert tokens["output_tokens"] == 150


def test_get_summary():
    engine = _make_engine()
    engine.record_cost("claude", 100, 50, 0.05, task_id="t1")
    summary = engine.get_summary("t1")
    assert summary["task_cost"] == 0.05
    assert summary["daily_cost"] == 0.05
    assert summary["monthly_cost"] == 0.05
    assert summary["total_calls"] == 1
    assert summary["input_tokens"] == 100
    assert summary["output_tokens"] == 50


def test_reset():
    engine = _make_engine()
    engine.record_cost("claude", 100, 50, 0.05)
    engine.reset()
    assert engine.get_daily_cost() == 0.0


# --- Budget enforcement ---


async def test_check_budget_task_exceeded():
    engine = _make_engine(task_limit=0.10, enforce=True)
    engine.record_cost("claude", 100, 50, 0.15, task_id="t1")
    try:
        await engine.check_budget(task_id="t1")
        raise AssertionError("Should have raised BudgetExceededError")
    except BudgetExceededError as e:
        assert e.limit_type == "task"
        assert e.current == 0.15


async def test_check_budget_daily_exceeded():
    engine = _make_engine(daily_limit=0.10, enforce=True)
    engine.record_cost("claude", 100, 50, 0.15)
    try:
        await engine.check_budget()
        raise AssertionError("Should have raised BudgetExceededError")
    except BudgetExceededError as e:
        assert e.limit_type == "daily"


async def test_check_budget_monthly_exceeded():
    engine = _make_engine(monthly_limit=0.10, enforce=True)
    engine.record_cost("claude", 100, 50, 0.15)
    try:
        await engine.check_budget()
        raise AssertionError("Should have raised BudgetExceededError")
    except BudgetExceededError as e:
        assert e.limit_type == "monthly"


async def test_check_budget_no_enforce():
    engine = _make_engine(task_limit=0.01, enforce=False)
    engine.record_cost("claude", 100, 50, 0.10, task_id="t1")
    # Should NOT raise even though limit is exceeded
    await engine.check_budget(task_id="t1")


async def test_check_budget_ok():
    engine = _make_engine(task_limit=1.0)
    engine.record_cost("claude", 100, 50, 0.01, task_id="t1")
    await engine.check_budget(task_id="t1")  # Should not raise


# --- Event emission ---


async def test_budget_warning_event():
    bus = EventBus()
    received = []

    async def handler(event):
        received.append(event)

    bus.subscribe(handler)

    engine = _make_engine(task_limit=1.0, bus=bus)
    # 85% of limit → should trigger warning
    engine.record_cost("claude", 100, 50, 0.85, task_id="t1")
    await engine.check_budget(task_id="t1")

    warnings = [e for e in received if e.event_type == EventType.BUDGET_WARNING]
    assert len(warnings) >= 1
    assert warnings[0].data["limit_type"] == "task"


async def test_budget_exceeded_event():
    bus = EventBus()
    received = []

    async def handler(event):
        received.append(event)

    bus.subscribe(handler)

    engine = _make_engine(task_limit=0.10, enforce=False, bus=bus)
    engine.record_cost("claude", 100, 50, 0.15, task_id="t1")
    await engine.check_budget(task_id="t1")

    exceeded = [e for e in received if e.event_type == EventType.BUDGET_EXCEEDED]
    assert len(exceeded) >= 1
    assert exceeded[0].data["limit_type"] == "task"


async def test_on_cost_event():
    """Test automatic recording from COST_UPDATE events."""
    from agent_scout.events import Event

    bus = EventBus()
    engine = BudgetEngine(event_bus=bus)

    await bus.emit(
        Event(
            event_type=EventType.COST_UPDATE,
            data={
                "model": "claude",
                "input_tokens": 100,
                "output_tokens": 50,
                "cost_usd": 0.05,
            },
            task_id="t1",
        )
    )

    assert engine.get_task_cost("t1") == 0.05


# --- Old records filtering ---


def test_daily_cost_excludes_old_records():
    engine = _make_engine()
    # Add a record that's 2 days old
    from agent_scout.llm.budget import CostRecord

    old = CostRecord(
        model="claude",
        input_tokens=100,
        output_tokens=50,
        cost_usd=1.0,
        timestamp=time.time() - 200000,
    )
    engine._records.append(old)
    engine.record_cost("claude", 100, 50, 0.05)

    assert engine.get_daily_cost() == 0.05
