"""
Session context propagation for provider wrappers.

When a :class:`~amplitude_ai.client.Session` context manager is active,
provider wrappers (``OpenAI``, ``Anthropic``, etc.) automatically inherit
``session_id``, ``trace_id``, ``agent_id``, ``turn_id``, and other
context fields — no extra arguments needed on the ``create()`` call.

This uses Python's :mod:`contextvars` module, which is:

- **Thread-safe** — each thread has its own context by default.
- **Async-safe** — each :func:`asyncio.Task` gets a copy of the
  parent context, so concurrent sessions never interfere.

Pattern::

    # Provider created once (outside any session)
    client = OpenAI(amplitude=ai.amplitude, api_key="sk-...")

    agent = ai.agent("support-bot", user_id="u1")
    with agent.session() as s:
        s.track_user_message(content="hello")
        # Provider auto-inherits session context:
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": "hello"}],
        )
"""

from __future__ import annotations

from contextvars import ContextVar
from typing import Any, Callable, Optional


class SessionContext:
    """Snapshot of the active session/agent context.

    Created by :meth:`Session.__enter__` and read by provider wrappers
    in ``track_completion()`` / ``create_streaming_tracker()``.

    Fields mirror the BoundAgent defaults plus session-level state.
    ``trace_id`` is mutable so that :meth:`Session.new_trace` can
    update it without replacing the entire context object.
    """

    __slots__ = (
        "session_id",
        "trace_id",
        "user_id",
        "agent_id",
        "parent_agent_id",
        "env",
        "customer_org_id",
        "agent_version",
        "description",
        "context",
        "groups",
        "idle_timeout_minutes",
        "device_id",
        "browser_session_id",
        "_next_turn_id",
    )

    def __init__(
        self,
        *,
        session_id: str,
        trace_id: str | None = None,
        user_id: str | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        env: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
        idle_timeout_minutes: int | None = None,
        device_id: str | None = None,
        browser_session_id: str | None = None,
        next_turn_id: Callable[[], int] | None = None,
    ) -> None:
        self.session_id = session_id
        self.trace_id = trace_id
        self.user_id = user_id
        self.agent_id = agent_id
        self.parent_agent_id = parent_agent_id
        self.env = env
        self.customer_org_id = customer_org_id
        self.agent_version = agent_version
        self.description = description
        self.context = context
        self.groups = groups
        self.idle_timeout_minutes = idle_timeout_minutes
        self.device_id = device_id
        self.browser_session_id = browser_session_id
        self._next_turn_id = next_turn_id

    def next_turn_id(self) -> Optional[int]:
        """Return the next auto-incremented turn ID, or ``None``
        if no turn counter is available."""
        if self._next_turn_id is not None:
            return self._next_turn_id()
        return None


# The ContextVar holding the active session context.
# Default is ``None`` (no session active).
_current_session: ContextVar[Optional[SessionContext]] = ContextVar(
    "amplitude_ai_session", default=None
)


def get_active_context() -> Optional[SessionContext]:
    """Return the active :class:`SessionContext`, or ``None``.

    Called by provider wrappers to auto-fill session/agent fields.
    """
    return _current_session.get()
