"""Integration tests for SqlUsageBackend using SQLite in-memory.

Skipped when sqlmodel or aiosqlite are not installed (requires --extra sql
and aiosqlite in test deps).
"""
# ruff: noqa: E402

from __future__ import annotations

import asyncio
from collections.abc import AsyncGenerator
from typing import Any

import pytest
import pytest_asyncio

sqlmodel = pytest.importorskip("sqlmodel")
aiosqlite = pytest.importorskip("aiosqlite")

from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlmodel import select

from agentbyte.context import AgentContext
from agentbyte.messages import Usage
from agentbyte.middleware.base import MiddlewareContext, OperationType
from agentbyte.middleware.sql_usage import (
    LlmUsageEvent,
    SqlUsageBackend,
    create_usage_tables,
)
from agentbyte.middleware.usage_logger import UsageLoggerMiddleware


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


@pytest_asyncio.fixture
async def engine():
    eng = create_async_engine("sqlite+aiosqlite:///:memory:", echo=False)
    await create_usage_tables(eng)
    yield eng
    await eng.dispose()


@pytest.fixture
def session_factory(engine):
    AsyncSessionLocal = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

    async def get_session() -> AsyncGenerator[AsyncSession, None]:
        async with AsyncSessionLocal() as session:
            yield session

    return get_session


@pytest.fixture
def backend(session_factory):
    return SqlUsageBackend(session_factory=session_factory)


@pytest.fixture
def middleware(backend):
    return UsageLoggerMiddleware(backend=backend, identifier="test-app")


@pytest.fixture
def agent_ctx():
    return AgentContext(
        session_id="sess-1",
        user_id="user-1",
        metadata={"run_id": "run-abc"},
    )


@pytest.fixture
def model_ctx(agent_ctx):
    return MiddlewareContext(
        operation=OperationType.MODEL_CALL,
        agent_name="test-agent",
        agent_context=agent_ctx,
        data={},
        metadata={"model": "gpt-4o"},
    )


def _make_result(
    tokens_input: int = 100,
    tokens_output: int = 50,
    duration_ms: int = 300,
    llm_calls: int = 1,
    tool_calls: int = 2,
    memory_operations: int = 1,
    cost_estimate: float | None = 0.003,
    model: str = "gpt-4o",
) -> Any:
    class _Result:
        pass

    r = _Result()
    r.usage = 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,
    )
    r.model = model
    return r


async def _rows(engine) -> list[LlmUsageEvent]:
    AsyncSessionLocal = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
    async with AsyncSessionLocal() as session:
        result = await session.execute(select(LlmUsageEvent))
        return list(result.scalars().all())


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


@pytest.mark.asyncio
async def test_writes_full_row(middleware, model_ctx, engine) -> None:
    result = _make_result()
    await middleware.process_response(model_ctx, result)
    await asyncio.sleep(0.1)

    rows = await _rows(engine)
    assert len(rows) == 1
    row = rows[0]
    assert row.run_id == "run-abc"
    assert row.session_id == "sess-1"
    assert row.user_id == "user-1"
    assert row.identifier == "test-app"
    assert row.agent_name == "test-agent"
    assert row.model == "gpt-4o"
    assert row.tokens_input == 100
    assert row.tokens_output == 50
    assert row.duration_ms == 300
    assert row.llm_calls == 1
    assert row.tool_calls == 2
    assert row.memory_operations == 1
    assert row.cost_estimate == pytest.approx(0.003)
    assert row.recorded_at is not None


@pytest.mark.asyncio
async def test_non_model_call_not_written(middleware, engine, agent_ctx) -> None:
    ctx = MiddlewareContext(
        operation=OperationType.TOOL_CALL,
        agent_name="agent",
        agent_context=agent_ctx,
        data={},
        metadata={},
    )
    await middleware.process_response(ctx, _make_result())
    await asyncio.sleep(0.1)
    assert await _rows(engine) == []


@pytest.mark.asyncio
async def test_null_run_id_row_still_written(middleware, engine) -> None:
    ctx = MiddlewareContext(
        operation=OperationType.MODEL_CALL,
        agent_name="agent",
        agent_context=AgentContext(session_id="s", user_id="u"),
        data={},
        metadata={"model": "gpt-4o"},
    )
    await middleware.process_response(ctx, _make_result())
    await asyncio.sleep(0.1)
    rows = await _rows(engine)
    assert len(rows) == 1
    assert rows[0].run_id is None


@pytest.mark.asyncio
async def test_multiple_calls_write_multiple_rows(middleware, model_ctx, engine) -> None:
    for _ in range(3):
        await middleware.process_response(model_ctx, _make_result())
    await asyncio.sleep(0.15)
    assert len(await _rows(engine)) == 3


@pytest.mark.asyncio
async def test_run_id_injected_by_agent_run() -> None:
    """agent.run() must stamp run_id when not already present."""
    ctx = AgentContext(session_id="s", user_id="u")
    assert "run_id" not in ctx.metadata

    # Import here to avoid circular issues at module level
    from agentbyte.agents.agent import Agent

    # We only want to verify the run_id injection, not run a real model call.
    # Patch _run_internal to return immediately.
    from unittest.mock import AsyncMock, patch
    from agentbyte.agents.types import AgentResponse
    from agentbyte.messages import Usage

    fake_response = AgentResponse(
        context=ctx,
        source="test",
        finish_reason="stop",
        usage=Usage(),
    )

    with patch.object(Agent, "_run_internal", new=AsyncMock(return_value=fake_response)):
        agent = Agent.__new__(Agent)
        agent.name = "test"
        agent.middleware_chain = type("MC", (), {"middlewares": []})()
        await agent.run(task="hello", context=ctx)

    assert "run_id" in ctx.metadata
    import uuid as _uuid
    _uuid.UUID(ctx.metadata["run_id"])  # raises if not a valid UUID


@pytest.mark.asyncio
async def test_run_id_not_overwritten() -> None:
    from unittest.mock import AsyncMock, patch
    from agentbyte.agents.agent import Agent
    from agentbyte.agents.types import AgentResponse
    from agentbyte.messages import Usage

    ctx = AgentContext(metadata={"run_id": "my-custom-id"})
    fake_response = AgentResponse(
        context=ctx,
        source="test",
        finish_reason="stop",
        usage=Usage(),
    )

    with patch.object(Agent, "_run_internal", new=AsyncMock(return_value=fake_response)):
        agent = Agent.__new__(Agent)
        agent.name = "test"
        agent.middleware_chain = type("MC", (), {"middlewares": []})()
        await agent.run(task="hello", context=ctx)

    assert ctx.metadata["run_id"] == "my-custom-id"
