"""Tests for the LLM call metrics middleware + repository.

Covers:
- Successful calls emit a row with duration, token counts and tokens_per_sec.
- Failed calls record the exception class as ``error_type`` plus a trimmed
  ``error_message``.
- The middleware never blocks on a backend exception (best-effort writes).
- The SQL store applies filters and pagination for the admin list endpoint.
"""

from __future__ import annotations

import asyncio
import uuid
from datetime import UTC, datetime, timedelta

import pytest
import pytest_asyncio
from langchain_core.messages import AIMessage
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from deerflow.agents.middlewares.llm_metrics_middleware import (
    LlmMetricsMiddleware,
    _extract_usage,
)
from deerflow.persistence.base import Base
from deerflow.persistence.llm_metrics import SqlLlmMetricsStore
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsQuery


class _RecordingStore:
    """Minimal in-memory store used to assert what the middleware persists."""

    def __init__(self) -> None:
        self.records: list[LlmMetricRecord] = []

    async def record(self, metric: LlmMetricRecord) -> None:
        self.records.append(metric)

    async def list(self, query: LlmMetricsQuery):
        return [r.__dict__ for r in self.records]

    async def count(self, query: LlmMetricsQuery) -> int:
        return len(self.records)

    async def iter_all(self, query: LlmMetricsQuery):
        for record in self.records:
            yield record.__dict__


@pytest_asyncio.fixture
async def sql_store(tmp_path):
    engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'metrics.db'}")
    async with engine.begin() as conn:
        # Importing the package registers the LlmCallMetricRow with Base.metadata.
        import deerflow.persistence.llm_metrics.model  # noqa: F401

        await conn.run_sync(Base.metadata.create_all)
    sf = async_sessionmaker(engine, expire_on_commit=False)
    store = SqlLlmMetricsStore(sf)
    try:
        yield store
    finally:
        await engine.dispose()


def _make_request():
    return type(
        "FakeRequest",
        (),
        {
            "model": type("FakeModel", (), {"model_name": "gpt-test"}),
            "runtime": type(
                "FakeRuntime",
                (),
                {
                    "config": {"configurable": {"thread_id": "t1"}},
                    "context": {"agent_name": "lead_agent", "run_id": "r1"},
                },
            ),
            "state": {"thread_id": "t1"},
        },
    )()


def test_middleware_records_success_with_speed():
    import time as _time

    store = _RecordingStore()
    mw = LlmMetricsMiddleware(store=store)
    response = AIMessage(
        content="hi",
        usage_metadata={"input_tokens": 5, "output_tokens": 20, "total_tokens": 25},
    )

    def slow_handler(_req):
        # Force a measurable elapsed window so tokens_per_sec is populated.
        # 50ms is comfortably above Windows monotonic-clock granularity.
        _time.sleep(0.05)
        return response

    result = mw.wrap_model_call(_make_request(), slow_handler)

    assert result is response
    # awrap_model_call uses a fire-and-forget task; here the sync path runs
    # asyncio.run synchronously, so the record landed already.
    assert len(store.records) == 1
    record = store.records[0]
    assert record.status == "success"
    assert record.model_name == "gpt-test"
    assert record.thread_id == "t1"
    assert record.run_id == "r1"
    assert record.input_tokens == 5
    assert record.output_tokens == 20
    assert record.total_tokens == 25
    # tokens_per_sec may be very high in tests because the handler is
    # ~instant; we only assert that the field is populated.
    assert record.tokens_per_sec is not None
    assert record.duration_ms >= 0


def test_middleware_records_error_with_classification():
    store = _RecordingStore()
    mw = LlmMetricsMiddleware(store=store)

    class BoomError(RuntimeError):
        pass

    def explode(_req):
        raise BoomError("provider exploded")

    with pytest.raises(BoomError):
        mw.wrap_model_call(_make_request(), explode)

    assert len(store.records) == 1
    record = store.records[0]
    assert record.status == "error"
    assert record.error_type == "BoomError"
    assert record.error_message == "provider exploded"
    assert record.tokens_per_sec is None


def test_middleware_swallows_store_errors():
    """A persistence outage must not break the agent runtime."""

    class FaultyStore:
        async def record(self, metric):
            raise RuntimeError("disk gone")

        async def list(self, q):
            return []

        async def count(self, q):
            return 0

        async def iter_all(self, q):
            if False:
                yield {}

    mw = LlmMetricsMiddleware(store=FaultyStore())
    # The wrap returns normally even though _flush would raise; the call below
    # must not raise.
    response = AIMessage(
        content="ok",
        usage_metadata={"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
    )
    out = mw.wrap_model_call(_make_request(), lambda req: response)
    assert out is response


@pytest.mark.asyncio
async def test_sql_store_round_trip(sql_store):
    base = datetime.now(UTC)
    for i in range(3):
        await sql_store.record(
            LlmMetricRecord(
                id=uuid.uuid4().hex,
                created_at=base - timedelta(minutes=i),
                user_id="alice",
                model_name="gpt-test" if i % 2 == 0 else "gpt-other",
                duration_ms=100 * (i + 1),
                output_tokens=10 * (i + 1),
                tokens_per_sec=10.0 * (i + 1),
                status="success",
            )
        )
    rows = await sql_store.list(LlmMetricsQuery(limit=10))
    assert len(rows) == 3
    # Newest first.
    assert rows[0]["duration_ms"] == 100

    filtered = await sql_store.list(LlmMetricsQuery(model_name="gpt-test", limit=10))
    assert all(r["model_name"] == "gpt-test" for r in filtered)
    assert len(filtered) == 2

    assert await sql_store.count(LlmMetricsQuery(model_name="gpt-test")) == 2


@pytest.mark.asyncio
async def test_sql_store_filters_errors_only(sql_store):
    base = datetime.now(UTC)
    await sql_store.record(
        LlmMetricRecord(id=uuid.uuid4().hex, created_at=base, status="success", duration_ms=10, output_tokens=1)
    )
    await sql_store.record(
        LlmMetricRecord(
            id=uuid.uuid4().hex,
            created_at=base,
            status="error",
            duration_ms=5,
            error_type="APITimeoutError",
            error_message="timed out",
        )
    )
    only_errors = await sql_store.list(LlmMetricsQuery(status="error", limit=10))
    assert len(only_errors) == 1
    assert only_errors[0]["error_type"] == "APITimeoutError"


@pytest.mark.asyncio
async def test_sql_store_filters_by_user_ids(sql_store):
    """``user_ids`` (set by a fuzzy username search) matches several accounts."""
    base = datetime.now(UTC)
    for uid in ("alice", "alicia", "bob"):
        await sql_store.record(
            LlmMetricRecord(
                id=uuid.uuid4().hex,
                created_at=base,
                user_id=uid,
                duration_ms=1,
                status="success",
            )
        )
    # Fuzzy "ali" → {alice, alicia}; bob must be excluded.
    rows = await sql_store.list(LlmMetricsQuery(user_ids=["alice", "alicia"], limit=10))
    assert {r["user_id"] for r in rows} == {"alice", "alicia"}
    assert await sql_store.count(LlmMetricsQuery(user_ids=["alice", "alicia"])) == 2
    # user_ids takes precedence over a single user_id.
    rows = await sql_store.list(LlmMetricsQuery(user_id="bob", user_ids=["alice"], limit=10))
    assert {r["user_id"] for r in rows} == {"alice"}


@pytest.mark.asyncio
async def test_sql_store_iter_all_paginates(sql_store):
    base = datetime.now(UTC)
    for i in range(1200):
        await sql_store.record(
            LlmMetricRecord(
                id=uuid.uuid4().hex,
                created_at=base - timedelta(seconds=i),
                duration_ms=1,
                status="success",
            )
        )
    seen = 0
    async for _ in sql_store.iter_all(LlmMetricsQuery()):
        seen += 1
    assert seen == 1200


def test_extract_usage_prefers_usage_metadata():
    response = AIMessage(
        content="x",
        usage_metadata={"input_tokens": 3, "output_tokens": 7, "total_tokens": 10},
    )
    out = _extract_usage(response)
    assert out == {"input_tokens": 3, "output_tokens": 7, "total_tokens": 10}


def test_extract_usage_falls_back_to_response_metadata_token_usage():
    """DeepSeek / vLLM-style: only total in usage_metadata, breakdown elsewhere."""
    response = AIMessage(
        content="x",
        usage_metadata={"input_tokens": 0, "output_tokens": 0, "total_tokens": 11093},
        response_metadata={
            "token_usage": {
                "prompt_tokens": 9000,
                "completion_tokens": 2093,
                "total_tokens": 11093,
            },
        },
    )
    out = _extract_usage(response)
    assert out["input_tokens"] == 9000
    assert out["output_tokens"] == 2093
    assert out["total_tokens"] == 11093


def test_extract_usage_unwraps_model_response():
    """The handler hands us a ModelResponse dataclass, not a bare AIMessage."""
    from langchain.agents.middleware.types import ModelResponse

    inner = AIMessage(
        content="hi",
        usage_metadata={"input_tokens": 7, "output_tokens": 9, "total_tokens": 16},
    )
    response = ModelResponse(result=[inner])
    out = _extract_usage(response)
    assert out == {"input_tokens": 7, "output_tokens": 9, "total_tokens": 16}


def test_extract_usage_unwraps_extended_model_response():
    """ExtendedModelResponse wraps a ModelResponse in a ``model_response`` field."""
    from langchain.agents.middleware.types import ExtendedModelResponse, ModelResponse

    inner = AIMessage(
        content="hi",
        usage_metadata={"input_tokens": 1, "output_tokens": 2, "total_tokens": 3},
    )
    extended = ExtendedModelResponse(model_response=ModelResponse(result=[inner]))
    out = _extract_usage(extended)
    assert out["output_tokens"] == 2


def test_extract_usage_derives_missing_field_when_two_present():
    response = AIMessage(
        content="x",
        usage_metadata={"input_tokens": 9000, "output_tokens": 0, "total_tokens": 11093},
    )
    out = _extract_usage(response)
    # output should be derived as total - input = 2093.
    assert out["output_tokens"] == 2093


@pytest.mark.asyncio
async def test_async_middleware_records_success():
    store = _RecordingStore()
    mw = LlmMetricsMiddleware(store=store)

    async def handler(_req):
        return AIMessage(
            content="hi",
            usage_metadata={"input_tokens": 1, "output_tokens": 4, "total_tokens": 5},
        )

    response = await mw.awrap_model_call(_make_request(), handler)
    assert isinstance(response, AIMessage)

    # awrap_model_call schedules the write as an asyncio.task; give the loop a
    # chance to drain.
    await asyncio.sleep(0)
    assert len(store.records) == 1
    assert store.records[0].status == "success"
    assert store.records[0].output_tokens == 4