"""Unit tests for UsageService (U7 — usage dashboard aggregations)."""

from __future__ import annotations

import csv
import io
import json

import pytest

from agentkit.llm.protocol import TokenUsage
from agentkit.llm.providers.usage_store import InMemoryUsageStore
from agentkit.server.admin.usage_service import (
    UsageService,
    get_usage_service,
    set_usage_service,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def store() -> InMemoryUsageStore:
    return InMemoryUsageStore()


@pytest.fixture
def service() -> UsageService:
    return UsageService()


@pytest.fixture(autouse=True)
def _reset_singleton():
    """Reset the UsageService singleton before and after each test."""
    set_usage_service(None)
    yield
    set_usage_service(None)


def _populate_store(store: InMemoryUsageStore) -> None:
    """Populate ``store`` with a mix of records for testing."""
    # User u1 in dept d1, gpt-4o, 100 tokens, $0.05
    store.record(
        "agent1",
        "gpt-4o",
        TokenUsage(prompt_tokens=60, completion_tokens=40),
        cost=0.05,
        latency_ms=200,
        user_id="u1",
        department_id="d1",
    )
    # User u1 in dept d1, claude, 200 tokens, $0.10
    store.record(
        "agent1",
        "claude",
        TokenUsage(prompt_tokens=120, completion_tokens=80),
        cost=0.10,
        latency_ms=300,
        user_id="u1",
        department_id="d1",
    )
    # User u2 in dept d2, gpt-4o, 50 tokens, $0.02
    store.record(
        "agent2",
        "gpt-4o",
        TokenUsage(prompt_tokens=30, completion_tokens=20),
        cost=0.02,
        latency_ms=100,
        user_id="u2",
        department_id="d2",
    )
    # User u3 in dept d1, gpt-4o, 500 tokens, $0.50 (top user)
    store.record(
        "agent3",
        "gpt-4o",
        TokenUsage(prompt_tokens=300, completion_tokens=200),
        cost=0.50,
        latency_ms=400,
        user_id="u3",
        department_id="d1",
    )


# ---------------------------------------------------------------------------
# get_usage_summary
# ---------------------------------------------------------------------------


class TestGetUsageSummary:
    async def test_summary_aggregates_all(self, service: UsageService, store: InMemoryUsageStore):
        _populate_store(store)
        result = await service.get_usage_summary(store)
        assert result["total_tokens"] == 850
        assert abs(result["total_cost"] - 0.67) < 1e-6
        assert result["total_requests"] == 4
        # by_model: gpt-4o (3 records, 650 tokens), claude (1, 200)
        assert "gpt-4o" in result["by_model"]
        assert "claude" in result["by_model"]
        assert result["by_model"]["gpt-4o"]["count"] == 3
        assert result["by_model"]["gpt-4o"]["total_tokens"] == 650
        # by_user: u1 (2 records, 300 tokens), u2 (1, 50), u3 (1, 500)
        assert result["by_user"]["u1"]["total_tokens"] == 300
        assert result["by_user"]["u2"]["total_tokens"] == 50
        assert result["by_user"]["u3"]["total_tokens"] == 500
        # by_department: d1 (3 records, 800 tokens), d2 (1, 50)
        assert result["by_department"]["d1"]["total_tokens"] == 800
        assert result["by_department"]["d2"]["total_tokens"] == 50

    async def test_summary_with_department_filter(
        self, service: UsageService, store: InMemoryUsageStore
    ):
        _populate_store(store)
        result = await service.get_usage_summary(store, department_id="d1")
        assert result["total_tokens"] == 800
        assert result["total_requests"] == 3
        # Only u1 and u3 are in d1.
        assert "u1" in result["by_user"]
        assert "u3" in result["by_user"]
        assert "u2" not in result["by_user"]

    async def test_summary_with_user_filter(self, service: UsageService, store: InMemoryUsageStore):
        _populate_store(store)
        result = await service.get_usage_summary(store, user_id="u2")
        assert result["total_tokens"] == 50
        assert result["total_requests"] == 1
        assert "u2" in result["by_user"]

    async def test_summary_with_empty_store(self, service: UsageService, store: InMemoryUsageStore):
        result = await service.get_usage_summary(store)
        assert result["total_tokens"] == 0
        assert result["total_cost"] == 0.0
        assert result["total_requests"] == 0
        assert result["by_model"] == {}
        assert result["by_user"] == {}
        assert result["by_department"] == {}


# ---------------------------------------------------------------------------
# get_usage_timeseries
# ---------------------------------------------------------------------------


class TestGetUsageTimeseries:
    async def test_timeseries_day_buckets(self, service: UsageService, store: InMemoryUsageStore):
        _populate_store(store)
        result = await service.get_usage_timeseries(store, interval="day")
        # All records are within the same day (today), so we expect one bucket.
        assert len(result) >= 1
        bucket = result[0]
        assert "timestamp" in bucket
        assert bucket["tokens"] == 850
        assert abs(bucket["cost"] - 0.67) < 1e-6
        assert bucket["requests"] == 4

    async def test_timeseries_hour_buckets(self, service: UsageService, store: InMemoryUsageStore):
        _populate_store(store)
        result = await service.get_usage_timeseries(store, interval="hour")
        # All records are within the same hour (now), so we expect one bucket.
        assert len(result) >= 1
        assert result[0]["tokens"] == 850

    async def test_timeseries_with_department_filter(
        self, service: UsageService, store: InMemoryUsageStore
    ):
        _populate_store(store)
        result = await service.get_usage_timeseries(store, department_id="d2", interval="day")
        assert len(result) >= 1
        assert result[0]["tokens"] == 50

    async def test_timeseries_empty_store(self, service: UsageService, store: InMemoryUsageStore):
        result = await service.get_usage_timeseries(store, interval="day")
        assert result == []


# ---------------------------------------------------------------------------
# get_usage_by_model
# ---------------------------------------------------------------------------


class TestGetUsageByModel:
    async def test_by_model_breakdown(self, service: UsageService, store: InMemoryUsageStore):
        _populate_store(store)
        result = await service.get_usage_by_model(store)
        # Sorted by model name: claude, gpt-4o
        assert len(result) == 2
        models = {row["model"] for row in result}
        assert models == {"gpt-4o", "claude"}
        gpt_row = next(r for r in result if r["model"] == "gpt-4o")
        assert gpt_row["tokens"] == 650
        assert gpt_row["requests"] == 3
        claude_row = next(r for r in result if r["model"] == "claude")
        assert claude_row["tokens"] == 200
        assert claude_row["requests"] == 1

    async def test_by_model_with_department_filter(
        self, service: UsageService, store: InMemoryUsageStore
    ):
        _populate_store(store)
        result = await service.get_usage_by_model(store, department_id="d2")
        # d2 only has gpt-4o (50 tokens, 1 request)
        assert len(result) == 1
        assert result[0]["model"] == "gpt-4o"
        assert result[0]["tokens"] == 50

    async def test_by_model_empty_store(self, service: UsageService, store: InMemoryUsageStore):
        result = await service.get_usage_by_model(store)
        assert result == []


# ---------------------------------------------------------------------------
# get_top_users
# ---------------------------------------------------------------------------


class TestGetTopUsers:
    async def test_top_users_sorted_by_tokens(
        self, service: UsageService, store: InMemoryUsageStore
    ):
        _populate_store(store)
        result = await service.get_top_users(store, limit=10)
        # u3 (500), u1 (300), u2 (50)
        assert len(result) == 3
        assert result[0]["user_id"] == "u3"
        assert result[0]["tokens"] == 500
        assert result[1]["user_id"] == "u1"
        assert result[1]["tokens"] == 300
        assert result[2]["user_id"] == "u2"
        assert result[2]["tokens"] == 50

    async def test_top_users_respects_limit(self, service: UsageService, store: InMemoryUsageStore):
        _populate_store(store)
        result = await service.get_top_users(store, limit=2)
        assert len(result) == 2
        assert result[0]["user_id"] == "u3"
        assert result[1]["user_id"] == "u1"

    async def test_top_users_with_department_filter(
        self, service: UsageService, store: InMemoryUsageStore
    ):
        _populate_store(store)
        result = await service.get_top_users(store, department_id="d1", limit=10)
        # d1 has u1 and u3
        assert len(result) == 2
        assert result[0]["user_id"] == "u3"
        assert result[1]["user_id"] == "u1"

    async def test_top_users_empty_store(self, service: UsageService, store: InMemoryUsageStore):
        result = await service.get_top_users(store, limit=10)
        assert result == []


# ---------------------------------------------------------------------------
# export_usage
# ---------------------------------------------------------------------------


class TestExportUsage:
    async def test_export_csv_has_header_and_rows(
        self, service: UsageService, store: InMemoryUsageStore
    ):
        _populate_store(store)
        body = await service.export_usage(store, format="csv")
        reader = csv.DictReader(io.StringIO(body))
        rows = list(reader)
        assert len(rows) == 4
        # Verify headers
        assert "timestamp" in rows[0]
        assert "agent_name" in rows[0]
        assert "model" in rows[0]
        assert "user_id" in rows[0]
        assert "department_id" in rows[0]
        # Verify a known record
        gpt_rows = [r for r in rows if r["model"] == "gpt-4o"]
        assert len(gpt_rows) == 3

    async def test_export_json_returns_valid_json(
        self, service: UsageService, store: InMemoryUsageStore
    ):
        _populate_store(store)
        body = await service.export_usage(store, format="json")
        data = json.loads(body)
        assert isinstance(data, list)
        assert len(data) == 4
        assert "timestamp" in data[0]
        assert "user_id" in data[0]
        assert "department_id" in data[0]

    async def test_export_csv_empty_store_returns_header_only(
        self, service: UsageService, store: InMemoryUsageStore
    ):
        body = await service.export_usage(store, format="csv")
        reader = csv.DictReader(io.StringIO(body))
        rows = list(reader)
        assert rows == []
        # Header should still be present.
        assert "timestamp" in body

    async def test_export_with_department_filter(
        self, service: UsageService, store: InMemoryUsageStore
    ):
        _populate_store(store)
        body = await service.export_usage(store, department_id="d2", format="csv")
        reader = csv.DictReader(io.StringIO(body))
        rows = list(reader)
        assert len(rows) == 1
        assert rows[0]["department_id"] == "d2"


# ---------------------------------------------------------------------------
# Singleton helpers
# ---------------------------------------------------------------------------


class TestSingletonHelpers:
    def test_get_usage_service_returns_singleton(self):
        first = get_usage_service()
        second = get_usage_service()
        assert first is second

    def test_set_usage_service_overrides(self):
        custom = UsageService()
        set_usage_service(custom)
        assert get_usage_service() is custom
        # Clearing falls back to a new lazy instance.
        set_usage_service(None)
        new_one = get_usage_service()
        assert new_one is not custom
