"""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 import select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from app.gateway.admin_audit_retention import sweep_admin_audit_retention
from deerflow.agents.middlewares.llm_metrics_middleware import (
    LlmMetricsMiddleware,
    _extract_usage,
)
from deerflow.persistence.admin_stats.model import AdminLeaderboardDailyStatRow
from deerflow.persistence.base import Base
from deerflow.persistence.llm_metrics import SqlLlmMetricsStore
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsQuery
from deerflow.persistence.llm_metrics.model import LlmCallMetricRow
from deerflow.persistence.run import RunRepository
from deerflow.persistence.run.model import RunRow
from deerflow.persistence.tool_metrics import SqlToolMetricsStore
from deerflow.persistence.tool_metrics.base import ToolMetricRecord, ToolMetricsQuery
from deerflow.persistence.tool_metrics.model import ToolCallMetricRow
from deerflow.persistence.types import BEIJING_TZ


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:
        await conn.run_sync(
            lambda sync_conn: Base.metadata.create_all(
                sync_conn,
                tables=[LlmCallMetricRow.__table__],
            )
        )
    sf = async_sessionmaker(engine, expire_on_commit=False)
    store = SqlLlmMetricsStore(sf)
    try:
        yield store
    finally:
        await engine.dispose()


@pytest_asyncio.fixture
async def tool_sql_store(tmp_path):
    engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'tool-metrics.db'}")
    async with engine.begin() as conn:
        await conn.run_sync(
            lambda sync_conn: Base.metadata.create_all(
                sync_conn,
                tables=[ToolCallMetricRow.__table__],
            )
        )
    store = SqlToolMetricsStore(async_sessionmaker(engine, expire_on_commit=False))
    try:
        yield store
    finally:
        await engine.dispose()


@pytest_asyncio.fixture
async def run_sql_store(tmp_path):
    engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'runs.db'}")
    async with engine.begin() as conn:
        await conn.run_sync(
            lambda sync_conn: Base.metadata.create_all(
                sync_conn,
                tables=[RunRow.__table__],
            )
        )
    store = RunRepository(async_sessionmaker(engine, expire_on_commit=False))
    try:
        yield store
    finally:
        await engine.dispose()


@pytest_asyncio.fixture
async def snapshot_session_factory(tmp_path):
    engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'snapshots.db'}")
    async with engine.begin() as conn:
        await conn.run_sync(
            lambda sync_conn: Base.metadata.create_all(
                sync_conn,
                tables=[AdminLeaderboardDailyStatRow.__table__],
            )
        )
    session_factory = async_sessionmaker(engine, expire_on_commit=False)
    try:
        yield session_factory
    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


@pytest.mark.asyncio
async def test_llm_metric_purge_before_keeps_rows_at_or_after_cutoff(sql_store):
    """The retention boundary is exclusive: a record exactly at cutoff stays."""

    cutoff = datetime.now(UTC) - timedelta(days=90)
    for record_id, created_at in (
        ("old", cutoff - timedelta(days=1)),
        ("boundary", cutoff),
        ("fresh", cutoff + timedelta(days=1)),
    ):
        await sql_store.record(
            LlmMetricRecord(id=record_id, created_at=created_at, duration_ms=1)
        )

    assert await sql_store.purge_before(cutoff) == 1
    # A subsequent cleanup pass is idempotent.
    assert await sql_store.purge_before(cutoff) == 0

    rows = await sql_store.list(LlmMetricsQuery(limit=10))
    assert {row["id"] for row in rows} == {"boundary", "fresh"}


@pytest.mark.asyncio
async def test_tool_metric_purge_before_keeps_rows_at_or_after_cutoff(tool_sql_store):
    """Tool-call audit retention has the same strict-before-cutoff behavior."""

    cutoff = datetime.now(UTC) - timedelta(days=90)
    for record_id, created_at in (
        ("old", cutoff - timedelta(days=1)),
        ("boundary", cutoff),
        ("fresh", cutoff + timedelta(days=1)),
    ):
        await tool_sql_store.record(
            ToolMetricRecord(
                id=record_id,
                created_at=created_at,
                tool_name="web_search",
                duration_ms=1,
            )
        )

    assert await tool_sql_store.purge_before(cutoff) == 1
    assert await tool_sql_store.purge_before(cutoff) == 0

    rows = await tool_sql_store.list(ToolMetricsQuery(limit=10))
    assert {row["id"] for row in rows} == {"boundary", "fresh"}


@pytest.mark.asyncio
async def test_terminal_run_summary_purge_keeps_boundary_and_live_runs(run_sql_store):
    """Audit-summary retention never deletes a current, boundary, or late-completed run."""

    cutoff = datetime.now(UTC) - timedelta(days=90)
    async with run_sql_store._sf() as session:
        session.add_all(
            [
                RunRow(
                    run_id="old-terminal",
                    thread_id="thread-old-terminal",
                    status="success",
                    created_at=cutoff - timedelta(seconds=1),
                    updated_at=cutoff - timedelta(seconds=1),
                ),
                RunRow(
                    run_id="boundary-terminal",
                    thread_id="thread-boundary-terminal",
                    status="error",
                    created_at=cutoff,
                    updated_at=cutoff,
                ),
                RunRow(
                    run_id="old-live",
                    thread_id="thread-old-live",
                    status="running",
                    created_at=cutoff - timedelta(days=1),
                    updated_at=cutoff - timedelta(days=1),
                ),
                RunRow(
                    run_id="old-created-recently-completed",
                    thread_id="thread-old-created-recently-completed",
                    status="success",
                    created_at=cutoff - timedelta(days=1),
                    updated_at=cutoff + timedelta(seconds=1),
                ),
            ]
        )
        await session.commit()

    assert await run_sql_store.purge_terminal_before(cutoff) == 1
    assert await run_sql_store.get("old-terminal", user_id=None) is None
    assert await run_sql_store.get("boundary-terminal", user_id=None) is not None
    assert await run_sql_store.get("old-live", user_id=None) is not None
    assert await run_sql_store.get("old-created-recently-completed", user_id=None) is not None


@pytest.mark.asyncio
async def test_llm_metric_purge_caps_one_large_batch_at_500(sql_store):
    """A single metrics delete never expands into an unbounded SQL statement."""

    cutoff = datetime.now(UTC) - timedelta(days=90)
    async with sql_store._sf() as session:
        session.add_all(
            [
                LlmCallMetricRow(
                    id=f"expired-{index}",
                    created_at=cutoff - timedelta(days=1),
                    duration_ms=1,
                )
                for index in range(501)
            ]
        )
        await session.commit()

    # Even a much larger requested size must respect the portable 500-row cap.
    assert await sql_store.purge_before(cutoff, batch_size=10_000) == 500
    assert await sql_store.count(LlmMetricsQuery()) == 1
    assert await sql_store.purge_before(cutoff, batch_size=10_000) == 1


@pytest.mark.asyncio
async def test_terminal_run_summary_purge_caps_one_large_batch_at_500(run_sql_store):
    """Run-summary retention uses the same bounded primary-key deletion rule."""

    cutoff = datetime.now(UTC) - timedelta(days=90)
    async with run_sql_store._sf() as session:
        session.add_all(
            [
                RunRow(
                    run_id=f"expired-run-{index}",
                    thread_id=f"thread-{index}",
                    status="success",
                    created_at=cutoff - timedelta(days=1),
                    updated_at=cutoff - timedelta(days=1),
                )
                for index in range(501)
            ]
        )
        await session.commit()

    assert await run_sql_store.purge_terminal_before(cutoff, batch_size=10_000) == 500
    async with run_sql_store._sf() as session:
        remaining = list((await session.scalars(select(RunRow.run_id))).all())
    assert len(remaining) == 1
    assert await run_sql_store.purge_terminal_before(cutoff, batch_size=10_000) == 1


@pytest.mark.asyncio
async def test_admin_retention_sweep_deletes_only_one_configured_batch_per_store(
    sql_store,
    run_sql_store,
):
    """A small sweep budget leaves the remaining historic backlog for later ticks."""

    now = datetime(2026, 10, 26, tzinfo=UTC)
    cutoff = now - timedelta(days=90)
    async with sql_store._sf() as session:
        session.add_all(
            [
                LlmCallMetricRow(
                    id=f"metric-{index}",
                    created_at=cutoff - timedelta(days=1),
                    duration_ms=1,
                )
                for index in range(11)
            ]
        )
        await session.commit()
    async with run_sql_store._sf() as session:
        session.add_all(
            [
                RunRow(
                    run_id=f"run-{index}",
                    thread_id=f"thread-{index}",
                    status="success",
                    created_at=cutoff - timedelta(days=1),
                    updated_at=cutoff - timedelta(days=1),
                )
                for index in range(11)
            ]
        )
        await session.commit()

    result = await sweep_admin_audit_retention(
        session_factory=None,
        llm_metrics_store=sql_store,
        run_store=run_sql_store,
        now=now,
        batch_size=10,
        max_batches_per_store=1,
        inter_batch_pause_seconds=0,
        max_sweep_seconds=60,
    )

    assert result["llm_metrics"] == 10
    assert result["run_summaries"] == 10
    assert await sql_store.count(LlmMetricsQuery()) == 1
    async with run_sql_store._sf() as session:
        remaining_runs = list((await session.scalars(select(RunRow.run_id))).all())
    assert len(remaining_runs) == 1


@pytest.mark.asyncio
async def test_admin_retention_snapshot_batch_does_not_delete_an_entire_old_date(
    snapshot_session_factory,
):
    """Snapshot cleanup selects composite keys, not every settings hash for a date."""

    now = datetime(2026, 10, 26, tzinfo=UTC)
    cutoff = now - timedelta(days=90)
    cutoff_date = cutoff.astimezone(BEIJING_TZ).date().isoformat()
    old_date = (cutoff - timedelta(days=1)).astimezone(BEIJING_TZ).date().isoformat()
    async with snapshot_session_factory() as session:
        session.add_all(
            [
                AdminLeaderboardDailyStatRow(
                    stat_date=old_date,
                    settings_hash=f"old-hash-{index:02d}",
                )
                for index in range(12)
            ]
            + [
                AdminLeaderboardDailyStatRow(
                    stat_date=cutoff_date,
                    settings_hash="boundary-hash",
                )
            ]
        )
        await session.commit()

    result = await sweep_admin_audit_retention(
        session_factory=snapshot_session_factory,
        now=now,
        batch_size=10,
        max_batches_per_store=1,
        inter_batch_pause_seconds=0,
        max_sweep_seconds=60,
    )

    assert result["leaderboard_snapshots"] == 10
    async with snapshot_session_factory() as session:
        rows = list((await session.scalars(select(AdminLeaderboardDailyStatRow))).all())
    assert sum(row.stat_date == old_date for row in rows) == 2
    assert [(row.stat_date, row.settings_hash) for row in rows if row.stat_date == cutoff_date] == [
        (cutoff_date, "boundary-hash")
    ]


class _AlwaysFullRetentionStore:
    """A store that models an arbitrarily large historical backlog."""

    def __init__(self, maximum_calls: int) -> None:
        self.calls: list[tuple[datetime, int]] = []
        self.maximum_calls = maximum_calls

    async def purge_before(self, cutoff: datetime, *, batch_size: int) -> int:
        self.calls.append((cutoff, batch_size))
        # Fail fast if a future implementation loses the sweep budget instead
        # of letting this regression test loop forever.
        if len(self.calls) > self.maximum_calls:
            raise AssertionError("retention sweep exceeded its per-store work budget")
        return batch_size


class _BusyRetentionStore:
    """Models a table temporarily locked by another database writer."""

    async def purge_before(self, cutoff: datetime, *, batch_size: int) -> int:
        raise RuntimeError("database is locked")


@pytest.mark.asyncio
async def test_admin_retention_sweep_has_fixed_per_store_work_budget():
    """One scheduler tick stops after a bounded amount of historic work."""

    batch_size = 37
    max_batches = 3
    store = _AlwaysFullRetentionStore(maximum_calls=max_batches)
    now = datetime(2026, 10, 26, tzinfo=UTC)

    result = await sweep_admin_audit_retention(
        session_factory=None,
        llm_metrics_store=store,
        now=now,
        batch_size=batch_size,
        max_batches_per_store=max_batches,
        inter_batch_pause_seconds=0,
        max_sweep_seconds=60,
    )

    expected_cutoff = now - timedelta(days=90)
    assert store.calls == [(expected_cutoff, batch_size)] * max_batches
    assert result == {
        "llm_metrics": batch_size * max_batches,
        "tool_metrics": 0,
        "run_summaries": 0,
        "leaderboard_snapshots": 0,
    }


@pytest.mark.asyncio
async def test_admin_retention_busy_store_does_not_block_other_sources():
    """A lock on one audit table defers only that source to a later tick."""

    available_store = _AlwaysFullRetentionStore(maximum_calls=1)
    result = await sweep_admin_audit_retention(
        session_factory=None,
        llm_metrics_store=_BusyRetentionStore(),
        tool_metrics_store=available_store,
        now=datetime(2026, 10, 26, tzinfo=UTC),
        batch_size=10,
        max_batches_per_store=1,
        inter_batch_pause_seconds=0,
        max_sweep_seconds=60,
    )

    assert available_store.calls
    assert result == {
        "llm_metrics": 0,
        "tool_metrics": 10,
        "run_summaries": 0,
        "leaderboard_snapshots": 0,
    }


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