"""Tests for UsageLoggerMiddleware and UsagePersistenceBackend interface.

No SQL dependency — uses a mock backend so these tests run without the
``sql`` extra installed.
"""

from __future__ import annotations

from typing import Any

import pytest

from agentbyte.context import AgentContext
from agentbyte.messages import Usage
from agentbyte.middleware.base import MiddlewareContext, OperationType
from agentbyte.middleware.usage_logger import (
    UsageEvent,
    UsageLoggerMiddleware,
    UsagePersistenceBackend,
)


# ---------------------------------------------------------------------------
# Test helpers
# ---------------------------------------------------------------------------


class RecordingBackend(UsagePersistenceBackend):
    """Captures every write() call for assertion."""

    def __init__(self) -> None:
        self.events: list[UsageEvent] = []

    async def write(self, event: UsageEvent) -> None:
        self.events.append(event)


class FailingBackend(UsagePersistenceBackend):
    """Always raises on write()."""

    async def write(self, event: UsageEvent) -> None:
        raise RuntimeError("simulated DB failure")


def _make_usage(
    tokens_input: int = 100,
    tokens_output: int = 50,
    duration_ms: int = 250,
    llm_calls: int = 1,
    tool_calls: int = 0,
    memory_operations: int = 0,
    cost_estimate: float | None = 0.002,
) -> Usage:
    return Usage(
        tokens_input=tokens_input,
        tokens_output=tokens_output,
        duration_ms=duration_ms,
        llm_calls=llm_calls,
        tool_calls=tool_calls,
        memory_operations=memory_operations,
        cost_estimate=cost_estimate,
    )


def _make_result(usage: Usage | None = None, model: str = "gpt-4o") -> Any:
    class _Result:
        pass

    r = _Result()
    r.usage = usage if usage is not None else _make_usage()
    r.model = model
    return r


def _make_context(
    operation: str = OperationType.MODEL_CALL,
    agent_name: str = "test-agent",
    run_id: str | None = "run-abc",
    session_id: str | None = "sess-1",
    user_id: str | None = "user-1",
    model: str = "gpt-4o",
) -> MiddlewareContext:
    metadata: dict[str, Any] = {}
    if run_id is not None:
        metadata["run_id"] = run_id

    agent_ctx = AgentContext(session_id=session_id, user_id=user_id, metadata=metadata)
    return MiddlewareContext(
        operation=operation,
        agent_name=agent_name,
        agent_context=agent_ctx,
        data={},
        metadata={"model": model},
    )


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_request_passthrough() -> None:
    mw = UsageLoggerMiddleware(backend=RecordingBackend(), identifier="app")
    ctx = _make_context()
    returned = await mw.process_request(ctx)
    assert returned is ctx


@pytest.mark.asyncio
async def test_process_error_returns_none() -> None:
    mw = UsageLoggerMiddleware(backend=RecordingBackend())
    ctx = _make_context()
    result = await mw.process_error(ctx, ValueError("boom"))
    assert result is None


@pytest.mark.asyncio
async def test_writes_event_on_model_call() -> None:
    backend = RecordingBackend()
    mw = UsageLoggerMiddleware(backend=backend, identifier="my-app")
    ctx = _make_context(run_id="run-xyz", session_id="s1", user_id="u1", model="gpt-4o-mini")
    result = _make_result(model="gpt-4o-mini")

    await mw.process_response(ctx, result)
    # allow fire-and-forget task to run
    import asyncio
    await asyncio.sleep(0.05)

    assert len(backend.events) == 1
    ev = backend.events[0]
    assert ev.run_id == "run-xyz"
    assert ev.session_id == "s1"
    assert ev.user_id == "u1"
    assert ev.identifier == "my-app"
    assert ev.agent_name == "test-agent"
    assert ev.model == "gpt-4o-mini"
    assert ev.tokens_input == 100
    assert ev.tokens_output == 50
    assert ev.duration_ms == 250
    assert ev.llm_calls == 1
    assert ev.cost_estimate == pytest.approx(0.002)


@pytest.mark.asyncio
async def test_non_model_call_not_written() -> None:
    import asyncio

    backend = RecordingBackend()
    mw = UsageLoggerMiddleware(backend=backend)
    ctx = _make_context(operation=OperationType.TOOL_CALL)
    await mw.process_response(ctx, _make_result())
    await asyncio.sleep(0.05)
    assert backend.events == []


@pytest.mark.asyncio
async def test_missing_usage_not_written() -> None:
    import asyncio

    backend = RecordingBackend()
    mw = UsageLoggerMiddleware(backend=backend)

    class _NoUsage:
        usage = None
        model = "gpt-4o"

    await mw.process_response(_make_context(), _NoUsage())
    await asyncio.sleep(0.05)
    assert backend.events == []


@pytest.mark.asyncio
async def test_backend_failure_swallowed_as_warning() -> None:
    import asyncio

    mw = UsageLoggerMiddleware(backend=FailingBackend())
    ctx = _make_context()
    result = _make_result()

    with pytest.warns(RuntimeWarning, match="backend write failed"):
        returned = await mw.process_response(ctx, result)
        await asyncio.sleep(0.05)

    assert returned is result


@pytest.mark.asyncio
async def test_null_run_id_in_metadata() -> None:
    import asyncio

    backend = RecordingBackend()
    mw = UsageLoggerMiddleware(backend=backend)
    ctx = _make_context(run_id=None)
    await mw.process_response(ctx, _make_result())
    await asyncio.sleep(0.05)
    assert len(backend.events) == 1
    assert backend.events[0].run_id is None


@pytest.mark.asyncio
async def test_no_agent_context() -> None:
    """Middleware must not crash when agent_context is None."""
    import asyncio

    backend = RecordingBackend()
    mw = UsageLoggerMiddleware(backend=backend, identifier="app")
    ctx = MiddlewareContext(
        operation=OperationType.MODEL_CALL,
        agent_name="agent",
        agent_context=None,
        data={},
        metadata={"model": "gpt-4o"},
    )
    await mw.process_response(ctx, _make_result())
    await asyncio.sleep(0.05)
    assert len(backend.events) == 1
    ev = backend.events[0]
    assert ev.run_id is None
    assert ev.session_id is None
    assert ev.user_id is None
