"""
ASGI middleware for automatic session tracking in FastAPI / Starlette apps.

Usage::

    from amplitude_ai.middleware import AmplitudeAIMiddleware

    app.add_middleware(
        AmplitudeAIMiddleware,
        amplitude_ai=ai,
        user_id_resolver=lambda request: request.state.user.id,
    )

Each HTTP request gets an auto-created session in the ``_current_session``
ContextVar.  Provider wrappers and ``@tool`` calls within the request
handler automatically inherit this session context.

The middleware flushes pending events after the response completes.
"""

from __future__ import annotations

import logging
import time
import uuid
from typing import Any, Callable

logger = logging.getLogger(__name__)

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response

from .context import SessionContext, _current_session


class AmplitudeAIMiddleware(BaseHTTPMiddleware):  # type: ignore[misc]
    """Starlette/FastAPI middleware that creates a session per request.

    Parameters
    ----------
    app:
        The ASGI application.
    amplitude_ai:
        An :class:`~amplitude_ai.AmplitudeAI` instance.
    user_id_resolver:
        A callable ``(Request) -> str`` that extracts the user ID from
        the request.  Required -- there is no default because guessing
        the user ID from request headers is fragile.
    session_id_resolver:
        Optional callable ``(Request) -> str``.  Defaults to generating
        a UUID per request.
    agent_id:
        Optional default agent ID for all requests.
    env:
        Optional environment tag (e.g., ``"production"``).
    track_session_events:
        When *True* (default), the middleware calls
        ``track_session_end()`` after the response completes.
    flush_on_response:
        When *True* (default), flush pending events after each response.
        Disable for high-throughput services where the background flush
        interval is sufficient.
    """

    def __init__(
        self,
        app: Any,
        *,
        amplitude_ai: Any,
        user_id_resolver: Callable[[Request], str | None],
        session_id_resolver: Callable[[Request], str] | None = None,
        agent_id: str | None = None,
        env: str | None = None,
        track_session_events: bool = True,
        flush_on_response: bool = True,
    ) -> None:
        super().__init__(app)
        self._ai = amplitude_ai
        self._user_id_resolver = user_id_resolver
        self._session_id_resolver = session_id_resolver or (lambda _: str(uuid.uuid4()))
        self._agent_id = agent_id
        self._env = env
        self._track_session_events = track_session_events
        self._flush_on_response = flush_on_response

    async def dispatch(self, request: Request, call_next: Any) -> Response:
        user_id = self._user_id_resolver(request)
        session_id = self._session_id_resolver(request)

        # Extract trace ID from incoming headers (W3C traceparent or custom)
        trace_id = request.headers.get("x-trace-id")
        if not trace_id and "traceparent" in request.headers:
            parts = request.headers["traceparent"].split("-")
            trace_id = parts[1] if len(parts) >= 2 else None
        if not trace_id:
            trace_id = str(uuid.uuid4())

        ctx = SessionContext(
            session_id=session_id,
            trace_id=trace_id,
            user_id=user_id,
            agent_id=self._agent_id,
            env=self._env,
            next_turn_id=lambda: self._ai._next_turn_id(session_id),
        )

        token = _current_session.set(ctx)
        start = time.time()

        try:
            response = await call_next(request)
            return response
        finally:
            _current_session.reset(token)
            latency_ms = (time.time() - start) * 1000

            if self._track_session_events and user_id is not None:
                try:
                    self._ai.track_session_end(
                        user_id=user_id,
                        session_id=session_id,
                        trace_id=trace_id,
                        env=self._env,
                        agent_id=self._agent_id,
                    )
                except Exception:
                    logger.warning("Failed to track session end in middleware", exc_info=True)

            if self._flush_on_response:
                try:
                    self._ai.flush()
                except Exception:
                    logger.warning("Failed to flush events in middleware", exc_info=True)
