"""Usage logging middleware with a pluggable persistence backend.

Writes one :class:`UsageEvent` per LLM ``MODEL_CALL`` to whatever
:class:`UsagePersistenceBackend` the caller supplies.  The first concrete
backend is :class:`~agentbyte.middleware.sql_usage.SqlUsageBackend`; future
backends (CosmosDB, DynamoDB, …) implement the same interface.

Usage::

    from agentbyte.middleware.usage_logger import UsageLoggerMiddleware
    from agentbyte.middleware.sql_usage import SqlUsageBackend, create_usage_tables

    backend = SqlUsageBackend(session_factory=get_session)
    mw = UsageLoggerMiddleware(backend=backend, identifier="my-app")
    agent = Agent(..., middlewares=[mw])
"""

from __future__ import annotations

import asyncio
import warnings
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from typing import Any, Optional

from pydantic import BaseModel, Field

from .base import BaseMiddleware, MiddlewareContext, OperationType


class UsageEvent(BaseModel):
    """Data contract passed to every :class:`UsagePersistenceBackend`.

    All fields are derived from the :class:`~agentbyte.middleware.base.MiddlewareContext`
    and the :class:`~agentbyte.llm.types.ChatCompletionResult` returned by a
    ``MODEL_CALL`` operation.  Backends receive this object and persist it in
    whatever form suits their storage technology.
    """

    run_id: Optional[str] = Field(
        default=None,
        description="Groups all LLM calls that belong to one agent.run() invocation",
    )
    session_id: Optional[str] = Field(
        default=None,
        description="Session identifier from AgentContext",
    )
    user_id: Optional[str] = Field(
        default=None,
        description="User identifier from AgentContext",
    )
    identifier: Optional[str] = Field(
        default=None,
        description="Application or service label set at middleware init",
    )
    agent_name: str = Field(description="Name of the agent that made the call")
    model: str = Field(description="Model name used for the call")
    # Full Usage fields
    duration_ms: int = Field(default=0, description="Wall-clock time in milliseconds")
    llm_calls: int = Field(default=0, description="LLM call count (typically 1 per event)")
    tokens_input: int = Field(default=0, description="Prompt tokens consumed")
    tokens_output: int = Field(default=0, description="Completion tokens generated")
    tool_calls: int = Field(default=0, description="Tool invocations in Usage")
    memory_operations: int = Field(default=0, description="Memory operations in Usage")
    cost_estimate: Optional[float] = Field(
        default=None,
        description="Estimated USD cost (None when provider does not supply pricing)",
    )
    recorded_at: datetime = Field(
        default_factory=lambda: datetime.now(timezone.utc),
        description="UTC timestamp when the event was created",
    )


class UsagePersistenceBackend(ABC):
    """Interface that all usage storage backends must satisfy.

    Implement :meth:`write` to persist a :class:`UsageEvent`.  The optional
    :meth:`setup` and :meth:`close` hooks let backends manage connection
    lifecycle when the caller needs explicit control.

    Example — minimal custom backend::

        class PrintBackend(UsagePersistenceBackend):
            async def write(self, event: UsageEvent) -> None:
                print(event.model_dump_json())
    """

    @abstractmethod
    async def write(self, event: UsageEvent) -> None:
        """Persist one usage event.  Must not raise — errors should be handled internally."""

    async def setup(self) -> None:
        """Optional: initialise connections, create tables, etc."""

    async def close(self) -> None:
        """Optional: clean up connections and flush buffers."""


class UsageLoggerMiddleware(BaseMiddleware):
    """Middleware that writes one :class:`UsageEvent` per LLM call to a backend.

    Writes are fire-and-forget (``asyncio.create_task``).  Backend failures are
    swallowed as ``RuntimeWarning`` and never interrupt agent execution.

    Args:
        backend: Any :class:`UsagePersistenceBackend` implementation.
        identifier: Optional application/service label stored on every event.
    """

    def __init__(
        self,
        backend: UsagePersistenceBackend,
        identifier: Optional[str] = None,
    ) -> None:
        self._backend = backend
        self._identifier = identifier

    async def _write(self, event: UsageEvent) -> None:
        try:
            await self._backend.write(event)
        except Exception as exc:
            warnings.warn(
                f"UsageLoggerMiddleware: backend write failed: {exc}",
                RuntimeWarning,
                stacklevel=2,
            )

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        return context

    async def process_response(
        self,
        context: MiddlewareContext,
        result: Any,
    ) -> Any:
        if context.operation != OperationType.MODEL_CALL:
            return result

        usage = getattr(result, "usage", None)
        if usage is None:
            return result

        agent_ctx = context.agent_context
        run_id: Optional[str] = None
        if agent_ctx is not None:
            raw = agent_ctx.metadata.get("run_id")
            if raw is not None:
                run_id = str(raw)

        model: str = context.metadata.get("model") or getattr(result, "model", "unknown")

        event = UsageEvent(
            run_id=run_id,
            session_id=agent_ctx.session_id if agent_ctx is not None else None,
            user_id=agent_ctx.user_id if agent_ctx is not None else None,
            identifier=self._identifier,
            agent_name=context.agent_name,
            model=model,
            duration_ms=usage.duration_ms,
            llm_calls=usage.llm_calls,
            tokens_input=usage.tokens_input,
            tokens_output=usage.tokens_output,
            tool_calls=usage.tool_calls,
            memory_operations=usage.memory_operations,
            cost_estimate=usage.cost_estimate,
        )

        asyncio.create_task(self._write(event))
        return result

    async def process_error(
        self,
        context: MiddlewareContext,
        error: Exception,
    ) -> Optional[Any]:
        return None
