"""Tests for orchestration.providers.cost — cost tracking and budget enforcement."""

from __future__ import annotations

import pytest

from orxaq.orchestration.providers.cost import (
    BudgetConfig,
    CostAwareRouter,
    CostTracker,
    UsageRecord,
)

# ── Fixtures ──────────────────────────────────────────────────


@pytest.fixture()
def tracker() -> CostTracker:
    return CostTracker()


@pytest.fixture()
def budget_tracker() -> CostTracker:
    return CostTracker(budget=BudgetConfig(hard_limit_usd=1.0, soft_limit_usd=0.5))


@pytest.fixture()
def router(tracker: CostTracker) -> CostAwareRouter:
    return CostAwareRouter(tracker)


# ── CostTracker: Pricing ─────────────────────────────────────


class TestCostTrackerPricing:
    def test_get_pricing_known_model(self, tracker: CostTracker) -> None:
        assert tracker.get_pricing("openai", "gpt-4") == (0.03, 0.06)

    def test_get_pricing_default_fallback(self, tracker: CostTracker) -> None:
        assert tracker.get_pricing("openai", "unknown-model") == (0.01, 0.03)

    def test_get_pricing_unknown_provider(self, tracker: CostTracker) -> None:
        assert tracker.get_pricing("unknown", "model") == (0.0, 0.0)

    def test_set_pricing_override(self, tracker: CostTracker) -> None:
        tracker.set_pricing("openai", "gpt-4", 0.001, 0.002)
        assert tracker.get_pricing("openai", "gpt-4") == (0.001, 0.002)

    def test_set_pricing_new_provider(self, tracker: CostTracker) -> None:
        tracker.set_pricing("custom", "model-x", 0.01, 0.02)
        assert tracker.get_pricing("custom", "model-x") == (0.01, 0.02)

    def test_vllm_is_free(self, tracker: CostTracker) -> None:
        assert tracker.get_pricing("vllm", "anything") == (0.0, 0.0)


# ── CostTracker: Compute Cost ────────────────────────────────


class TestComputeCost:
    def test_known_model_cost(self, tracker: CostTracker) -> None:
        # gpt-3.5-turbo: input=0.0005, output=0.0015 per 1K
        cost = tracker.compute_cost("openai", "gpt-3.5-turbo", 1000, 1000)
        assert cost == pytest.approx(0.002)

    def test_zero_tokens(self, tracker: CostTracker) -> None:
        assert tracker.compute_cost("openai", "gpt-4", 0, 0) == 0.0

    def test_vllm_zero_cost(self, tracker: CostTracker) -> None:
        assert tracker.compute_cost("vllm", "default", 10000, 10000) == 0.0

    def test_anthropic_opus_cost(self, tracker: CostTracker) -> None:
        # opus: input=0.015, output=0.075 per 1K
        cost = tracker.compute_cost("anthropic", "claude-3-opus", 2000, 500)
        expected = (2000 * 0.015 + 500 * 0.075) / 1000.0
        assert cost == pytest.approx(expected)


# ── CostTracker: Record ──────────────────────────────────────


class TestRecord:
    def test_record_returns_usage_record(self, tracker: CostTracker) -> None:
        rec = tracker.record("openai", "gpt-4", 100, 50, latency_ms=200.0)
        assert isinstance(rec, UsageRecord)
        assert rec.provider == "openai"
        assert rec.model == "gpt-4"
        assert rec.input_tokens == 100
        assert rec.output_tokens == 50
        assert rec.latency_ms == 200.0
        assert rec.cost_usd > 0
        assert rec.timestamp > 0

    def test_record_accumulates(self, tracker: CostTracker) -> None:
        tracker.record("openai", "gpt-4", 100, 50)
        tracker.record("anthropic", "claude-3-sonnet", 200, 100)
        assert len(tracker.records) == 2

    def test_records_property_returns_copy(self, tracker: CostTracker) -> None:
        tracker.record("vllm", "default", 100, 50)
        recs = tracker.records
        recs.clear()
        assert len(tracker.records) == 1  # original unchanged


# ── CostTracker: Budget ──────────────────────────────────────


class TestBudget:
    def test_unlimited_budget(self, tracker: CostTracker) -> None:
        budget = tracker.check_budget()
        assert budget["hard_limit_usd"] == 0.0
        assert budget["within_budget"] is True
        assert budget["warning"] is False

    def test_within_hard_limit(self, budget_tracker: CostTracker) -> None:
        budget_tracker.record("vllm", "default", 100, 50)  # $0 cost
        budget = budget_tracker.check_budget()
        assert budget["within_budget"] is True

    def test_over_hard_limit(self, budget_tracker: CostTracker) -> None:
        # Record enough to exceed $1.0 hard limit
        for _ in range(20):
            budget_tracker.record("openai", "gpt-4", 5000, 5000)
        budget = budget_tracker.check_budget()
        assert budget["within_budget"] is False

    def test_soft_limit_warning(self, budget_tracker: CostTracker) -> None:
        # Record enough to exceed $0.5 soft but under $1.0 hard
        for _ in range(7):
            budget_tracker.record("openai", "gpt-4", 1000, 1000)
        budget = budget_tracker.check_budget()
        assert budget["warning"] is True
        assert budget["within_budget"] is True

    def test_is_within_budget_unlimited(self, tracker: CostTracker) -> None:
        assert tracker.is_within_budget() is True

    def test_is_within_budget_over_limit(self, budget_tracker: CostTracker) -> None:
        for _ in range(20):
            budget_tracker.record("openai", "gpt-4", 5000, 5000)
        assert budget_tracker.is_within_budget() is False

    def test_records_in_window_count(self, budget_tracker: CostTracker) -> None:
        budget_tracker.record("openai", "gpt-4", 100, 50)
        budget_tracker.record("openai", "gpt-4", 200, 100)
        budget = budget_tracker.check_budget()
        assert budget["records_in_window"] == 2


# ── CostTracker: Summary ─────────────────────────────────────


class TestSummary:
    def test_empty_summary(self, tracker: CostTracker) -> None:
        s = tracker.summary()
        assert s["total_cost_usd"] == 0.0
        assert s["total_calls"] == 0
        assert s["by_provider"] == {}

    def test_summary_with_records(self, tracker: CostTracker) -> None:
        tracker.record("openai", "gpt-4", 1000, 500)
        tracker.record("anthropic", "claude-3-sonnet", 2000, 1000)
        s = tracker.summary()
        assert s["total_calls"] == 2
        assert s["total_cost_usd"] > 0
        assert "openai" in s["by_provider"]
        assert "anthropic" in s["by_provider"]
        assert s["by_provider"]["openai"]["calls"] == 1
        assert s["by_provider"]["anthropic"]["calls"] == 1

    def test_summary_model_breakdown(self, tracker: CostTracker) -> None:
        tracker.record("openai", "gpt-4", 100, 50)
        tracker.record("openai", "gpt-3.5-turbo", 200, 100)
        s = tracker.summary()
        models = s["by_provider"]["openai"]["models"]
        assert "gpt-4" in models
        assert "gpt-3.5-turbo" in models

    def test_summary_token_totals(self, tracker: CostTracker) -> None:
        tracker.record("vllm", "default", 500, 300)
        s = tracker.summary()
        assert s["total_input_tokens"] == 500
        assert s["total_output_tokens"] == 300


# ── CostTracker: Reset ───────────────────────────────────────


class TestReset:
    def test_reset_clears_records(self, tracker: CostTracker) -> None:
        tracker.record("openai", "gpt-4", 100, 50)
        tracker.reset()
        assert len(tracker.records) == 0
        assert tracker.summary()["total_calls"] == 0


# ── CostAwareRouter ──────────────────────────────────────────


class TestCostAwareRouter:
    def test_rank_by_cost_cheapest_first(self, router: CostAwareRouter) -> None:
        ranked = router.rank_by_cost(["openai", "anthropic", "vllm"])
        assert ranked[0] == "vllm"  # free
        assert ranked[-1] == "openai"  # most expensive

    def test_should_skip_no_budget(self, router: CostAwareRouter) -> None:
        # No hard limit configured → never skip
        assert router.should_skip("openai") is False

    def test_should_skip_over_budget(self) -> None:
        tracker = CostTracker(budget=BudgetConfig(hard_limit_usd=0.01))
        # Spend enough to exceed limit
        for _ in range(5):
            tracker.record("openai", "gpt-4", 1000, 1000)
        car = CostAwareRouter(tracker)
        assert car.should_skip("openai", estimated_tokens=1000) is True

    def test_should_skip_vllm_free_under_budget(self) -> None:
        tracker = CostTracker(budget=BudgetConfig(hard_limit_usd=100.0))
        tracker.record("openai", "gpt-4", 100, 50)  # small spend
        car = CostAwareRouter(tracker)
        # vLLM adds $0 cost so should not exceed budget
        assert car.should_skip("vllm", estimated_tokens=10000) is False

    def test_select_filters_and_ranks(self, router: CostAwareRouter) -> None:
        selected = router.select(["openai", "anthropic", "vllm"])
        assert selected[0] == "vllm"  # cheapest first
        assert len(selected) == 3

    def test_select_filters_over_budget(self) -> None:
        tracker = CostTracker(budget=BudgetConfig(hard_limit_usd=0.01))
        for _ in range(5):
            tracker.record("openai", "gpt-4", 1000, 1000)
        car = CostAwareRouter(tracker)
        selected = car.select(["openai", "anthropic", "vllm"])
        # vllm should survive (free), expensive providers filtered
        assert "vllm" in selected

    def test_select_returns_fallback_when_all_over_budget(self) -> None:
        tracker = CostTracker(budget=BudgetConfig(hard_limit_usd=0.001))
        for _ in range(5):
            tracker.record("openai", "gpt-4", 1000, 1000)
        car = CostAwareRouter(tracker, provider_costs={"a": 10.0, "b": 20.0})
        selected = car.select(["a", "b"])
        # Should return first candidate as fallback
        assert len(selected) == 1
        assert selected[0] == "a"
