"""
AmplitudeAI client — the public entry point for the Amplitude AI SDK.

Uses composition over inheritance: wraps an ``Amplitude`` instance and
delegates event delivery through the customer's existing pipeline.
"""

from __future__ import annotations

import atexit
import logging
import threading
import uuid
import weakref
from collections.abc import AsyncIterator, Iterator
from concurrent.futures import Future
from contextlib import asynccontextmanager, contextmanager
from typing import Any

from amplitude import Amplitude

from .config import AIConfig, ContentMode
from .context import SessionContext, _current_session
from .core.enrichments import SessionEnrichments
from .core.privacy import PrivacyConfig
from .exceptions import ConfigurationError
from .serverless import is_serverless
from .utils.costs import calculate_cost
from .core.tracking import (
    PROP_SESSION_REPLAY_ID,
    track_ai_message,
    track_embedding,
    track_score,
    track_session_end,
    track_session_enrichment,
    track_span,
    track_tool_call,
    track_user_message,
)

_logger = logging.getLogger(__name__)

# Maximum number of session turn counters to track before evicting oldest.
# Prevents unbounded memory growth for long-running servers where sessions
# may never receive an explicit track_session_end() call.
_MAX_SESSION_TURN_COUNTERS = 10_000


class _TrackingProxy:
    """Thin proxy that intercepts ``track()`` for counting and debug/dry-run.

    All other attribute access is forwarded to the real
    :class:`~amplitude.Amplitude` instance.  This avoids monkey-patching
    ``amplitude.track`` — the real instance stays untouched, so mocks in
    tests retain their ``call_args`` / ``assert_called_*`` capabilities.
    """

    __slots__ = ("_real", "_owner")

    def __init__(self, real: Amplitude, owner: AmplitudeAI) -> None:
        object.__setattr__(self, "_real", real)
        object.__setattr__(self, "_owner", owner)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._real, name)

    def track(self, event: Any) -> None:
        owner = self._owner
        owner._track_count_since_flush += 1

        if owner._config.debug or owner._config.dry_run:
            import sys

            from .utils.debug import format_debug_line, format_dry_run_line

            if owner._config.debug:
                print(format_debug_line(event), file=sys.stderr)
            if owner._config.dry_run:
                print(format_dry_run_line(event), file=sys.stderr)
                return  # dry-run: count but don't send

        self._real.track(event)

# Weak references to all live AmplitudeAI instances for the exit warning.
_active_instances: weakref.WeakSet[AmplitudeAI] = weakref.WeakSet()
_exit_hook_registered = False


def _register_exit_hook() -> None:
    global _exit_hook_registered  # noqa: PLW0603
    if _exit_hook_registered:
        return
    _exit_hook_registered = True

    def _warn_unflushed() -> None:
        if not is_serverless():
            return
        for instance in _active_instances:
            if instance._track_count_since_flush > 0:
                _logger.warning(
                    "AmplitudeAI: %d event(s) were tracked but never flushed. "
                    "In serverless environments, call `ai.flush()` before your "
                    "handler returns, or use `with agent.session() as s:` which "
                    "auto-flushes by default.",
                    instance._track_count_since_flush,
                )
                break  # one warning is enough

    atexit.register(_warn_unflushed)


class AmplitudeAI:
    """Amplitude AI SDK client for LLM observability.

    Wraps an :class:`amplitude.Amplitude` instance (composition) to track
    Agent events without duplicating the event pipeline.

    **Recommended** — pass an existing ``Amplitude`` instance::

        from amplitude import Amplitude
        from amplitude_ai import AmplitudeAI

        amplitude = Amplitude("api-key")
        ai = AmplitudeAI(amplitude=amplitude)

    **Standalone** — creates an ``Amplitude`` instance internally::

        ai = AmplitudeAI(api_key="api-key")

    Parameters
    ----------
    amplitude:
        An existing :class:`amplitude.Amplitude` instance.  The SDK shares
        this client's event pipeline — no duplicate queues.
    api_key:
        Amplitude API key.  Used only when ``amplitude`` is not provided.
    config:
        SDK configuration (:class:`AIConfig`).  Controls privacy tier,
        PII redaction, etc.

    Raises
    ------
    ConfigurationError
        If neither ``amplitude`` nor ``api_key`` is provided.
    """

    def __init__(
        self,
        amplitude: Amplitude | None = None,
        api_key: str | None = None,
        config: AIConfig | None = None,
    ) -> None:
        if amplitude is not None:
            self._amplitude = amplitude
            self._owns_client = False
        elif api_key is not None:
            self._amplitude = Amplitude(api_key)
            self._owns_client = True
        else:
            raise ConfigurationError(
                "Provide either an existing Amplitude instance via "
                "'amplitude=' or an API key via 'api_key='."
            )

        self._config = config or AIConfig()
        self._privacy_config = self._config.to_privacy_config()

        # Per-session turn counter (spec: "auto-incremented if omitted")
        self._session_turn_counters: dict[str, int] = {}
        self._turn_lock = threading.Lock()

        # Wire per-event delivery callback into the Amplitude SDK.
        # If the caller already configured a callback on the Amplitude
        # instance (e.g., for monitoring / alerting), compose rather than
        # silently replacing it.
        if self._config.on_event_callback is not None:
            existing_cb = getattr(self._amplitude.configuration, "callback", None)
            if existing_cb is not None:
                new_cb = self._config.on_event_callback

                def _composed_callback(event: Any, code: int, message: str = "") -> None:
                    try:
                        existing_cb(event, code, message)
                    except Exception:
                        pass  # Don't let one callback failure block the other
                    try:
                        new_cb(event, code, message)
                    except Exception:
                        pass

                self._amplitude.configuration.callback = _composed_callback
            else:
                self._amplitude.configuration.callback = self._config.on_event_callback

        # Proxy intercepts track() for counting + debug/dry-run without
        # monkey-patching the real Amplitude instance.
        self._track_count_since_flush = 0
        self._tracking_proxy = _TrackingProxy(self._amplitude, self)

        _active_instances.add(self)
        _register_exit_hook()

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def amplitude(self) -> Amplitude:
        """Access the underlying :class:`amplitude.Amplitude` instance."""
        return self._amplitude

    @property
    def config(self) -> AIConfig:
        """Access the SDK configuration."""
        return self._config

    def status(self) -> dict[str, Any]:
        """Return SDK configuration summary for introspection.

        Useful for health-check endpoints and debugging.  Contains only
        client-side configuration -- no server-side metrics.
        """
        import importlib

        _availability_flags = {
            "openai": "OPENAI_AVAILABLE",
            "anthropic": "ANTHROPIC_AVAILABLE",
            "gemini": "GEMINI_AVAILABLE",
            "azure_openai": "AZURE_OPENAI_AVAILABLE",
            "bedrock": "BEDROCK_AVAILABLE",
            "mistral": "MISTRAL_AVAILABLE",
        }
        available: list[str] = []
        for mod_name, flag_name in _availability_flags.items():
            try:
                mod = importlib.import_module(
                    f".providers.{mod_name}", package="amplitude_ai"
                )
                if getattr(mod, flag_name, False):
                    available.append(mod_name.replace("_", "-"))
            except ImportError:
                pass

        from .patching import patched_providers

        return {
            "content_mode": self._config.content_mode.value,
            "debug": self._config.debug,
            "dry_run": self._config.dry_run,
            "redact_pii": self._config.redact_pii,
            "providers_available": available,
            "patched_providers": patched_providers(),
        }

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _next_turn_id(self, session_id: str) -> int:
        """Return the next auto-incremented turn ID for *session_id*.

        Thread-safe.  Each ``session_id`` gets its own monotonically
        increasing counter starting at 1.

        When the number of tracked sessions exceeds
        ``_MAX_SESSION_TURN_COUNTERS``, the oldest session entry is
        evicted to prevent unbounded memory growth in long-running
        servers where sessions may never call ``track_session_end()``.
        """
        with self._turn_lock:
            # Evict oldest entry when at capacity and this is a new session
            if (
                session_id not in self._session_turn_counters
                and len(self._session_turn_counters) >= _MAX_SESSION_TURN_COUNTERS
            ):
                oldest_key = next(iter(self._session_turn_counters))
                del self._session_turn_counters[oldest_key]

            self._session_turn_counters[session_id] = (
                self._session_turn_counters.get(session_id, 0) + 1
            )
            return self._session_turn_counters[session_id]

    # ------------------------------------------------------------------
    # Message Tracking
    # ------------------------------------------------------------------

    def track_user_message(
        self,
        user_id: str,
        content: str,
        session_id: str,
        trace_id: str | None = None,
        turn_id: int | None = None,
        message_id: str | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        # Implicit feedback signals
        is_regeneration: bool = False,
        is_edit: bool = False,
        edited_message_id: str | None = None,
        # Attachments
        attachments: list[dict[str, Any]] | None = None,
        # Message source
        message_source: str | None = "user",
        event_properties: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
    ) -> str:
        """Track a user message.

        When *turn_id* is ``None`` it is auto-incremented per *session_id*.
        When *message_id* is ``None`` a UUID is auto-generated.

        Args:
            trace_id: Groups related events within a session into a single
                trace (e.g., one user request -> tool calls -> AI response).
            message_id: Custom message identifier.  Auto-generated if omitted.
            parent_agent_id: ID of the agent that delegated to this one
                (multi-agent orchestration).
            customer_org_id: Org ID of the end customer using the agent
                (multi-tenant platforms).
            env: Environment tag (e.g., ``"production"``, ``"staging"``).
            is_regeneration: Whether the user requested a regeneration.
            is_edit: Whether the user edited a previous message.
            edited_message_id: The ``message_id`` of the original message
                that was edited.
            attachments: List of attachment metadata dicts (type, name,
                size_bytes, mime_type, etc.).  File content is never sent.

        Returns a ``message_id`` (UUID) that can be used as a ``target_id``
        in :meth:`score`.
        """
        effective_turn_id = (
            turn_id if turn_id is not None else self._next_turn_id(session_id)
        )
        return track_user_message(
            amplitude=self._tracking_proxy,
            user_id=user_id,
            message_content=content,
            session_id=session_id,
            trace_id=trace_id,
            turn_id=effective_turn_id,
            message_id=message_id,
            agent_id=agent_id,
            parent_agent_id=parent_agent_id,
            customer_org_id=customer_org_id,
            agent_version=agent_version,
            description=description,
            context=context,
            env=env,
            is_regeneration=is_regeneration,
            is_edit=is_edit,
            edited_message_id=edited_message_id,
            attachments=attachments,
            message_source=message_source,
            event_properties=event_properties,
            groups=groups,
            privacy_config=self._privacy_config,
        )

    def track_ai_message(
        self,
        user_id: str,
        content: str,
        session_id: str,
        model: str,
        provider: str,
        latency_ms: float,
        input_tokens: int | None = None,
        output_tokens: int | None = None,
        reasoning_tokens: int | None = None,
        cache_read_tokens: int | None = None,
        cache_creation_tokens: int | None = None,
        total_tokens: int | None = None,
        total_cost_usd: float | None = None,
        finish_reason: str | None = None,
        tool_calls: list[Any] | None = None,
        reasoning_content: str | None = None,
        # Tool definitions (request-side tool schemas sent to the LLM)
        tool_definitions: list[dict[str, Any]] | None = None,
        # Model configuration
        system_prompt: str | None = None,
        temperature: float | None = None,
        max_output_tokens: int | None = None,
        top_p: float | None = None,
        is_streaming: bool | None = None,
        prompt_id: str | None = None,
        # Implicit feedback signals
        was_copied: bool = False,
        # Semantic cache hit (distinct from token-level prompt caching)
        was_cached: bool = False,
        # Model tier override (auto-inferred from model name if not provided)
        model_tier: str | None = None,
        # Attachments (AI-generated images, charts, files, etc.)
        attachments: list[dict[str, Any]] | None = None,
        trace_id: str | None = None,
        turn_id: int | None = None,
        message_id: str | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        is_error: bool = False,
        error_message: str | None = None,
        ttfb_ms: float | None = None,
        event_properties: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
    ) -> str:
        """Track an AI response.

        When *turn_id* is ``None`` it is auto-incremented per *session_id*.
        When *message_id* is ``None`` a UUID is auto-generated.

        When *total_cost_usd* is ``None`` and both *input_tokens* and
        *output_tokens* are provided, cost is auto-calculated using
        cache-aware pricing (including *reasoning_tokens*,
        *cache_read_tokens*, and *cache_creation_tokens* when given).
        Pass an explicit *total_cost_usd* to override.

        Args:
            trace_id: Groups related events within a session into a single
                trace (e.g., one user request -> tool calls -> AI response).
            total_tokens: Total token count for the request.  Used by the
                enrichment pipeline for cost/usage accounting.
            message_id: Custom message identifier.  Auto-generated if omitted.
            parent_agent_id: ID of the agent that delegated to this one
                (multi-agent orchestration).
            customer_org_id: Org ID of the end customer using the agent
                (multi-tenant platforms).
            env: Environment tag (e.g., ``"production"``, ``"staging"``).
            is_error: Whether the AI response represents an error.
            error_message: Error description when ``is_error`` is ``True``.
            ttfb_ms: Time-to-first-byte in milliseconds.
            was_copied: Whether the user copied this AI response.
            was_cached: Whether this response was served from a semantic
                cache (e.g., Redis, LLM cache layer) rather than a live
                LLM call.  Distinct from ``cache_read_tokens`` which
                tracks token-level prompt caching.
            tool_definitions: Normalized list of tool schemas sent to the
                LLM in this request.  Each entry should have ``name``,
                ``description``, and ``parameters`` keys.  When provided,
                the SDK emits a hash (always) and the full definitions
                (subject to ``PrivacyConfig``).
            attachments: List of attachment dicts for AI-generated rich
                media.  Include a ``url`` key to enable rendering in the
                LLM session viewer.

        Returns a ``message_id`` (UUID) that can be used as a ``target_id``
        in :meth:`score` or as ``parent_message_id`` in :meth:`track_tool_call`.
        """
        effective_turn_id = (
            turn_id if turn_id is not None else self._next_turn_id(session_id)
        )

        # Auto-calculate cost when not explicitly provided (spec §3.2)
        effective_cost = total_cost_usd
        if effective_cost is None and input_tokens is not None and output_tokens is not None:
            try:
                effective_cost = calculate_cost(
                    model_name=model,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    reasoning_tokens=reasoning_tokens or 0,
                    cache_read_input_tokens=cache_read_tokens or 0,
                    cache_creation_input_tokens=cache_creation_tokens or 0,
                    default_provider=provider if provider else None,
                )
            except Exception:
                pass  # cost calculation is best-effort

        return track_ai_message(
            amplitude=self._tracking_proxy,
            user_id=user_id,
            model_name=model,
            provider=provider,
            response_content=content,
            latency_ms=latency_ms,
            session_id=session_id,
            trace_id=trace_id,
            turn_id=effective_turn_id,
            message_id=message_id,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            reasoning_tokens=reasoning_tokens,
            cache_read_input_tokens=cache_read_tokens,
            cache_creation_input_tokens=cache_creation_tokens,
            total_tokens=total_tokens,
            total_cost_usd=effective_cost,
            finish_reason=finish_reason,
            tool_calls=tool_calls,
            reasoning_content=reasoning_content,
            tool_definitions=tool_definitions,
            system_prompt=system_prompt,
            temperature=temperature,
            max_output_tokens=max_output_tokens,
            top_p=top_p,
            is_streaming=is_streaming,
            prompt_id=prompt_id,
            was_copied=was_copied,
            was_cached=was_cached,
            model_tier=model_tier,
            attachments=attachments,
            agent_id=agent_id,
            parent_agent_id=parent_agent_id,
            customer_org_id=customer_org_id,
            agent_version=agent_version,
            description=description,
            context=context,
            env=env,
            is_error=is_error,
            error_message=error_message,
            provider_ttfb_ms=ttfb_ms,
            event_properties=event_properties,
            groups=groups,
            privacy_config=self._privacy_config,
        )

    # ------------------------------------------------------------------
    # Operation Tracking
    # ------------------------------------------------------------------

    def track_tool_call(
        self,
        user_id: str,
        tool_name: str,
        latency_ms: float,
        success: bool,
        session_id: str | None = None,
        trace_id: str | None = None,
        turn_id: int | None = None,
        invocation_id: str | None = None,
        input: dict[str, Any] | None = None,
        output: Any | None = None,
        parent_message_id: str | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        error_message: str | None = None,
        event_properties: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
    ) -> str:
        """Track a tool call.

        Args:
            trace_id: Groups related events within a session into a single
                trace (e.g., one user request -> tool calls -> AI response).
            turn_id: Associates the tool call with a specific conversation
                turn.  Defaults to ``1`` in the underlying event.
            invocation_id: Custom invocation identifier.  Auto-generated
                (UUID) if omitted.
            parent_agent_id: ID of the agent that delegated to this one
                (multi-agent orchestration).
            customer_org_id: Org ID of the end customer using the agent
                (multi-tenant platforms).
            env: Environment tag (e.g., ``"production"``, ``"staging"``).
            error_message: Error description when ``success`` is ``False``.

        Returns an ``invocation_id`` (UUID) identifying this tool call.
        """
        kwargs: dict[str, Any] = {}
        if turn_id is not None:
            kwargs["turn_id"] = turn_id
        return track_tool_call(
            amplitude=self._tracking_proxy,
            user_id=user_id,
            tool_name=tool_name,
            success=success,
            latency_ms=latency_ms,
            session_id=session_id,
            trace_id=trace_id,
            invocation_id=invocation_id,
            tool_input=input,
            tool_output=output,
            parent_message_id=parent_message_id,
            agent_id=agent_id,
            parent_agent_id=parent_agent_id,
            customer_org_id=customer_org_id,
            agent_version=agent_version,
            description=description,
            context=context,
            env=env,
            error_message=error_message,
            event_properties=event_properties,
            groups=groups,
            privacy_config=self._privacy_config,
            **kwargs,
        )

    def track_embedding(
        self,
        user_id: str,
        model: str,
        provider: str,
        latency_ms: float,
        input_tokens: int | None = None,
        dimensions: int | None = None,
        total_cost_usd: float | None = None,
        session_id: str | None = None,
        trace_id: str | None = None,
        turn_id: int | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        event_properties: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
    ) -> str:
        """Track an embedding operation.

        Returns a ``span_id`` (UUID).
        """
        return track_embedding(
            amplitude=self._tracking_proxy,
            user_id=user_id,
            model=model,
            provider=provider,
            latency_ms=latency_ms,
            input_tokens=input_tokens,
            dimensions=dimensions,
            total_cost_usd=total_cost_usd,
            session_id=session_id,
            trace_id=trace_id,
            turn_id=turn_id,
            agent_id=agent_id,
            parent_agent_id=parent_agent_id,
            customer_org_id=customer_org_id,
            agent_version=agent_version,
            description=description,
            context=context,
            env=env,
            event_properties=event_properties,
            groups=groups,
            privacy_config=self._privacy_config,
        )

    def track_span(
        self,
        user_id: str,
        span_name: str,
        trace_id: str,
        latency_ms: float,
        input_state: dict[str, Any] | None = None,
        output_state: dict[str, Any] | None = None,
        parent_span_id: str | None = None,
        is_error: bool = False,
        error_message: str | None = None,
        session_id: str | None = None,
        turn_id: int | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        event_properties: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
    ) -> str:
        """Track a generic operation span (e.g., vector search, rerank).

        Returns a ``span_id`` (UUID).
        """
        return track_span(
            amplitude=self._tracking_proxy,
            user_id=user_id,
            span_name=span_name,
            trace_id=trace_id,
            latency_ms=latency_ms,
            input_state=input_state,
            output_state=output_state,
            parent_span_id=parent_span_id,
            is_error=is_error,
            error_message=error_message,
            session_id=session_id,
            turn_id=turn_id,
            agent_id=agent_id,
            parent_agent_id=parent_agent_id,
            customer_org_id=customer_org_id,
            agent_version=agent_version,
            description=description,
            context=context,
            env=env,
            event_properties=event_properties,
            groups=groups,
            privacy_config=self._privacy_config,
        )

    # ------------------------------------------------------------------
    # Session Management
    # ------------------------------------------------------------------

    def track_session_end(
        self,
        user_id: str,
        session_id: str,
        enrichments: SessionEnrichments | None = None,
        trace_id: str | None = None,
        turn_id: int | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        abandonment_turn: int | None = None,
        idle_timeout_minutes: int | None = None,
        event_properties: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
    ) -> None:
        """Track an explicit session close, optionally with enrichments.

        Args:
            abandonment_turn: If set, indicates the user abandoned the
                session at this turn number.
            idle_timeout_minutes: Hint to the enrichment pipeline for how
                long to wait before considering this session idle.
        """
        track_session_end(
            amplitude=self._tracking_proxy,
            user_id=user_id,
            session_id=session_id,
            enrichments=enrichments,
            trace_id=trace_id,
            turn_id=turn_id,
            agent_id=agent_id,
            parent_agent_id=parent_agent_id,
            customer_org_id=customer_org_id,
            agent_version=agent_version,
            description=description,
            context=context,
            env=env,
            abandonment_turn=abandonment_turn,
            idle_timeout_minutes=idle_timeout_minutes,
            event_properties=event_properties,
            groups=groups,
            privacy_config=self._privacy_config,
        )
        # Clean up turn counter for this session to prevent unbounded growth
        with self._turn_lock:
            self._session_turn_counters.pop(session_id, None)

    def track_session_enrichment(
        self,
        user_id: str,
        session_id: str,
        enrichments: SessionEnrichments,
        trace_id: str | None = None,
        turn_id: int | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        event_properties: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
    ) -> None:
        """Track customer-provided session classifications.

        Primarily used when ``content_mode="customer_enriched"``.
        """
        track_session_enrichment(
            amplitude=self._tracking_proxy,
            user_id=user_id,
            session_id=session_id,
            enrichments=enrichments,
            trace_id=trace_id,
            turn_id=turn_id,
            agent_id=agent_id,
            parent_agent_id=parent_agent_id,
            customer_org_id=customer_org_id,
            agent_version=agent_version,
            description=description,
            context=context,
            env=env,
            event_properties=event_properties,
            groups=groups,
            privacy_config=self._privacy_config,
        )

    # ------------------------------------------------------------------
    # Scoring
    # ------------------------------------------------------------------

    def score(
        self,
        user_id: str,
        name: str,
        value: float,
        target_id: str,
        target_type: str = "message",
        source: str = "user",
        comment: str | None = None,
        session_id: str | None = None,
        trace_id: str | None = None,
        turn_id: int | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        event_properties: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
    ) -> None:
        """Attach a quality signal to a message or session.

        Covers user feedback, automated evals, and human annotations.

        Args:
            name: Score name (e.g., ``"user-feedback"``, ``"accuracy"``).
            value: Numeric value (binary 1/0, continuous 0-1, or 1-5 rating).
            target_id: The ``message_id`` or ``session_id`` being scored.
            target_type: ``"message"`` or ``"session"``.
            source: ``"user"``, ``"ai"``, or ``"reviewer"``.
            comment: Optional text explanation (respects ``content_mode``).
        """
        track_score(
            amplitude=self._tracking_proxy,
            user_id=user_id,
            name=name,
            value=value,
            target_id=target_id,
            target_type=target_type,
            source=source,
            comment=comment,
            session_id=session_id,
            trace_id=trace_id,
            turn_id=turn_id,
            agent_id=agent_id,
            parent_agent_id=parent_agent_id,
            customer_org_id=customer_org_id,
            agent_version=agent_version,
            description=description,
            context=context,
            env=env,
            event_properties=event_properties,
            groups=groups,
            privacy_config=self._privacy_config,
        )

    # ------------------------------------------------------------------
    # Bound Agent Factory
    # ------------------------------------------------------------------

    def agent(
        self,
        agent_id: str,
        *,
        user_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        session_id: str | None = None,
        trace_id: str | None = None,
        groups: dict[str, Any] | None = None,
        device_id: str | None = None,
        browser_session_id: str | None = None,
    ) -> BoundAgent:
        """Create a :class:`BoundAgent` with pre-filled context.

        All ``track_*`` methods on the returned handle inherit the bound
        values, eliminating per-call boilerplate::

            agent = ai.agent("support-bot", user_id="u1", env="production",
                             agent_version="v4.2",
                             context={"experiment_variant": "prompt-v2"})
            agent.track_user_message(content="Hello", session_id="s1")
            agent.track_ai_message(content="Hi!", session_id="s1",
                                   model="gpt-4o", provider="openai",
                                   latency_ms=200)

        For **Session Replay linking**, pass the browser's ``device_id``
        and ``browser_session_id`` from the Amplitude Browser SDK.  These
        propagate to all events and sessions created from this agent.

        For multi-agent orchestration, use :meth:`BoundAgent.child` to
        spawn child agents that automatically set ``parent_agent_id``.
        """
        return BoundAgent(
            self,
            agent_id=agent_id,
            user_id=user_id,
            parent_agent_id=parent_agent_id,
            customer_org_id=customer_org_id,
            agent_version=agent_version,
            description=description,
            context=context,
            env=env,
            session_id=session_id,
            trace_id=trace_id,
            groups=groups,
            device_id=device_id,
            browser_session_id=browser_session_id,
        )

    # ------------------------------------------------------------------
    # Tenant Factory
    # ------------------------------------------------------------------

    def tenant(
        self,
        customer_org_id: str,
        *,
        groups: dict[str, Any] | None = None,
        env: str | None = None,
    ) -> TenantHandle:
        """Create a tenant-scoped handle for multi-tenant platforms.

        All agents created through this handle inherit the tenant's
        ``customer_org_id`` and ``groups``::

            tenant = ai.tenant("acme-corp", groups={"company": "acme-corp"})
            agent = tenant.agent("support-bot", user_id="u1")
            # agent.track_* calls carry customer_org_id="acme-corp"
        """
        return TenantHandle(
            self,
            customer_org_id=customer_org_id,
            groups=groups,
            env=env,
        )

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def flush(self) -> list[Future[Any] | None]:
        """Flush pending events.

        Delegates to :meth:`amplitude.Amplitude.flush` and returns one
        :class:`~concurrent.futures.Future` per destination plugin (or
        ``None`` when a plugin has nothing to flush).  In serverless
        environments, iterate and call ``.result()`` on each non-``None``
        future to block until delivery completes.
        """
        self._track_count_since_flush = 0
        return list(self._amplitude.flush())

    def shutdown(self) -> None:
        """Shut down the SDK.

        Only shuts down the underlying ``Amplitude`` client if it was
        created internally (via ``api_key``).  If the customer passed in
        their own instance, they manage the lifecycle.
        """
        _active_instances.discard(self)
        self._track_count_since_flush = 0
        if self._owns_client:
            self._amplitude.shutdown()


# ======================================================================
# Tenant Handle
# ======================================================================


class TenantHandle:
    """Tenant-scoped factory created by :meth:`AmplitudeAI.tenant`.

    Pre-fills ``customer_org_id``, ``groups``, and ``env`` on every
    :class:`BoundAgent` created through :meth:`agent`.
    """

    __slots__ = ("_ai", "_customer_org_id", "_groups", "_env")

    def __init__(
        self,
        ai: AmplitudeAI,
        *,
        customer_org_id: str,
        groups: dict[str, Any] | None = None,
        env: str | None = None,
    ) -> None:
        self._ai = ai
        self._customer_org_id = customer_org_id
        self._groups = groups
        self._env = env

    def agent(self, agent_id: str, **kwargs: Any) -> "BoundAgent":
        """Create a :class:`BoundAgent` inheriting tenant defaults.

        Explicit keyword arguments take precedence over tenant defaults.
        """
        kwargs.setdefault("customer_org_id", self._customer_org_id)
        if self._groups is not None:
            kwargs.setdefault("groups", self._groups)
        if self._env is not None:
            kwargs.setdefault("env", self._env)
        return self._ai.agent(agent_id, **kwargs)


# ======================================================================
# Bound Agent
# ======================================================================

# Fields that BoundAgent propagates to track_* calls.
_CONTEXT_FIELDS = (
    "user_id",
    "agent_id",
    "parent_agent_id",
    "customer_org_id",
    "agent_version",
    "description",
    "context",
    "env",
    "session_id",
    "trace_id",
    "groups",
)

# Per-method accepted context fields (subset of _CONTEXT_FIELDS).
# Methods that accept all fields don't need an entry here.
_FIELDS_EMBEDDING = _CONTEXT_FIELDS
_FIELDS_SPAN = _CONTEXT_FIELDS
_FIELDS_SESSION = _CONTEXT_FIELDS
_FIELDS_SESSION_ENRICHMENT = _CONTEXT_FIELDS
_FIELDS_SCORE = _CONTEXT_FIELDS


class BoundAgent:
    """A pre-configured agent handle that eliminates per-call boilerplate.

    Created via :meth:`AmplitudeAI.agent` or :meth:`BoundAgent.child`.
    Every ``track_*`` method mirrors :class:`AmplitudeAI` but fills in
    context fields (``user_id``, ``agent_id``, ``env``, etc.) from the
    bound defaults.  Explicit kwargs always override the bound values.

    Example::

        agent = ai.agent("orchestrator", env="production")
        with agent.session(user_id="u1") as s:
            s.track_user_message(content="Hi")

        child = agent.child("researcher")
        child.track_tool_call(tool_name="search", latency_ms=50,
                              success=True, session_id="s1")
    """

    __slots__ = ("_ai", "_defaults")

    def __init__(
        self,
        ai: AmplitudeAI,
        *,
        agent_id: str,
        user_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        description: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        session_id: str | None = None,
        trace_id: str | None = None,
        groups: dict[str, Any] | None = None,
        device_id: str | None = None,
        browser_session_id: str | None = None,
    ) -> None:
        self._ai = ai
        self._defaults: dict[str, Any] = {
            "user_id": user_id,
            "agent_id": agent_id,
            "parent_agent_id": parent_agent_id,
            "customer_org_id": customer_org_id,
            "agent_version": agent_version,
            "description": description,
            "context": context,
            "env": env,
            "session_id": session_id,
            "trace_id": trace_id,
            "groups": groups,
            "device_id": device_id,
            "browser_session_id": browser_session_id,
        }

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def agent_id(self) -> str:
        """The bound agent identifier."""
        return str(self._defaults["agent_id"])

    @property
    def ai(self) -> AmplitudeAI:
        """The parent :class:`AmplitudeAI` instance."""
        return self._ai

    # ------------------------------------------------------------------
    # Child agents
    # ------------------------------------------------------------------

    def child(
        self,
        agent_id: str,
        **overrides: Any,
    ) -> BoundAgent:
        """Create a child agent.

        Inherits ``user_id``, ``env``, ``customer_org_id``, ``session_id``,
        ``trace_id``, and ``groups`` from this agent.  Sets
        ``parent_agent_id`` to this agent's ``agent_id``.

        Override any inherited field via keyword arguments::

            child = agent.child("researcher", env="staging")
        """
        inherited = {
            "user_id": self._defaults.get("user_id"),
            "env": self._defaults.get("env"),
            "customer_org_id": self._defaults.get("customer_org_id"),
            "agent_version": self._defaults.get("agent_version"),
            "session_id": self._defaults.get("session_id"),
            "trace_id": self._defaults.get("trace_id"),
            "groups": self._defaults.get("groups"),
            "device_id": self._defaults.get("device_id"),
            "browser_session_id": self._defaults.get("browser_session_id"),
        }
        # Merge context dicts: parent context is the base, child overrides
        # specific keys.  A child can replace the entire context by passing
        # context={...} in overrides.
        parent_ctx = self._defaults.get("context")
        child_ctx = overrides.pop("context", None)
        if parent_ctx or child_ctx:
            merged_ctx = dict(parent_ctx or {})
            if child_ctx:
                merged_ctx.update(child_ctx)
            inherited["context"] = merged_ctx

        inherited.update({k: v for k, v in overrides.items() if v is not None})
        # Extract parent_agent_id from overrides (if provided) to avoid
        # duplicate keyword argument with the explicit parent_agent_id below.
        explicit_parent = inherited.pop(
            "parent_agent_id", self._defaults.get("agent_id")
        )
        return BoundAgent(
            self._ai,
            agent_id=agent_id,
            parent_agent_id=explicit_parent,
            **inherited,
        )

    # ------------------------------------------------------------------
    # Internal merge
    # ------------------------------------------------------------------

    def _merge(
        self,
        kwargs: dict[str, Any],
        fields: tuple[str, ...] = _CONTEXT_FIELDS,
    ) -> dict[str, Any]:
        """Merge explicit kwargs with bound defaults.

        Explicit non-``None`` values always win.  Bound defaults fill in
        anything the caller didn't provide.  *fields* limits which
        context keys are injected (some ``AmplitudeAI`` methods only
        accept a subset of the 8 context fields).

        When both ``device_id`` and ``browser_session_id`` are set,
        ``[Amplitude] Session Replay ID`` is auto-injected into
        ``event_properties`` (unless already present).
        """
        merged = dict(kwargs)
        for field in fields:
            if merged.get(field) is None and self._defaults.get(field) is not None:
                merged[field] = self._defaults[field]
        device_id = self._defaults.get("device_id")
        browser_session_id = self._defaults.get("browser_session_id")
        if device_id and browser_session_id:
            existing_ep = merged.get("event_properties")
            ep = dict(existing_ep) if isinstance(existing_ep, dict) else {}
            if PROP_SESSION_REPLAY_ID not in ep:
                ep[PROP_SESSION_REPLAY_ID] = f"{device_id}/{browser_session_id}"
                merged["event_properties"] = ep
        return merged

    # ------------------------------------------------------------------
    # Delegating track_* methods
    # ------------------------------------------------------------------

    def track_user_message(self, content: str, **kwargs: Any) -> str:
        """Track a user message.  See :meth:`AmplitudeAI.track_user_message`."""
        return self._ai.track_user_message(content=content, **self._merge(kwargs))

    def track_ai_message(
        self,
        content: str,
        model: str,
        provider: str,
        latency_ms: float,
        **kwargs: Any,
    ) -> str:
        """Track an AI response.  See :meth:`AmplitudeAI.track_ai_message`."""
        return self._ai.track_ai_message(
            content=content,
            model=model,
            provider=provider,
            latency_ms=latency_ms,
            **self._merge(kwargs),
        )

    def track_tool_call(
        self,
        tool_name: str,
        latency_ms: float,
        success: bool,
        **kwargs: Any,
    ) -> str:
        """Track a tool call.  See :meth:`AmplitudeAI.track_tool_call`."""
        return self._ai.track_tool_call(
            tool_name=tool_name,
            latency_ms=latency_ms,
            success=success,
            **self._merge(kwargs),
        )

    def track_embedding(
        self,
        model: str,
        provider: str,
        latency_ms: float,
        **kwargs: Any,
    ) -> str:
        """Track an embedding.  See :meth:`AmplitudeAI.track_embedding`."""
        return self._ai.track_embedding(
            model=model, provider=provider, latency_ms=latency_ms,
            **self._merge(kwargs, _FIELDS_EMBEDDING),
        )

    def track_span(
        self,
        span_name: str,
        latency_ms: float,
        **kwargs: Any,
    ) -> str:
        """Track a span.  See :meth:`AmplitudeAI.track_span`."""
        return self._ai.track_span(
            span_name=span_name, latency_ms=latency_ms,
            **self._merge(kwargs, _FIELDS_SPAN),
        )

    def track_session_end(self, **kwargs: Any) -> None:
        """End a session.  See :meth:`AmplitudeAI.track_session_end`."""
        self._ai.track_session_end(**self._merge(kwargs, _FIELDS_SESSION))

    def track_session_enrichment(
        self, enrichments: SessionEnrichments, **kwargs: Any
    ) -> None:
        """Send session enrichments.  See :meth:`AmplitudeAI.track_session_enrichment`."""
        self._ai.track_session_enrichment(
            enrichments=enrichments, **self._merge(kwargs, _FIELDS_SESSION)
        )

    def score(self, name: str, value: float, target_id: str, **kwargs: Any) -> None:
        """Attach a score.  See :meth:`AmplitudeAI.score`."""
        self._ai.score(
            name=name, value=value, target_id=target_id,
            **self._merge(kwargs, _FIELDS_SCORE),
        )

    # ------------------------------------------------------------------
    # Session context manager
    # ------------------------------------------------------------------

    def session(
        self,
        session_id: str | None = None,
        idle_timeout_minutes: int | None = None,
        user_id: str | None = None,
        device_id: str | None = None,
        browser_session_id: str | None = None,
        auto_flush: bool | None = None,
    ) -> Session:
        """Create a :class:`Session` context manager.

        Automatically calls ``track_session_end`` when the ``with``
        block exits (even on exception)::

            with agent.session(user_id="user-1") as s:
                s.new_trace()
                s.track_user_message(content="Hello")
                s.track_ai_message(content="Hi!", model="gpt-4o",
                                   provider="openai", latency_ms=200)
            # session auto-ended here + auto-flush in serverless

        For **Session Replay linking**, pass the browser's ``device_id``
        and ``browser_session_id`` from the Amplitude Browser SDK.  The
        SDK auto-computes ``[Amplitude] Session Replay ID`` and attaches
        it to every ``[Agent]`` event, enabling direct navigation from
        the LLM Analytics session explorer to the Session Replay viewer::

            with agent.session(
                user_id="user-1",
                device_id=request.headers["X-Amplitude-Device-ID"],
                browser_session_id=request.headers["X-Amplitude-Session-ID"],
            ) as s:
                ...

        Args:
            session_id: Optional session ID (auto-generated if omitted).
            idle_timeout_minutes: Hint to the enrichment pipeline for how
                long to wait before considering this session idle.
                Defaults to ``None`` (pipeline uses its own default, typically 30 min).
                Set higher for long-running sessions (e.g., 240 for dashboard monitoring).
            user_id: User ID for this session.  Overrides the agent-level
                ``user_id`` (if set) for all events within this session.
            device_id: The Amplitude Browser SDK ``device_id`` from the
                frontend.  Used together with *browser_session_id* to
                build ``[Amplitude] Session Replay ID``.
            browser_session_id: The Amplitude Browser SDK ``session_id``
                from the frontend (a millisecond timestamp).  Used
                together with *device_id* to build
                ``[Amplitude] Session Replay ID``.
            auto_flush: Flush pending events when the session exits.
                ``True`` = always flush (recommended for serverless).
                ``False`` = never flush.
                ``None`` (default) = auto-detect serverless environment.
        """
        return Session(
            self,
            session_id=session_id,
            idle_timeout_minutes=idle_timeout_minutes,
            user_id=user_id,
            device_id=device_id,
            browser_session_id=browser_session_id,
            auto_flush=auto_flush,
        )

    # ------------------------------------------------------------------
    # Lifecycle (delegates)
    # ------------------------------------------------------------------

    def flush(self) -> list[Future[Any] | None]:
        """Flush pending events.  Delegates to :meth:`AmplitudeAI.flush`."""
        return self._ai.flush()

    def shutdown(self) -> None:
        """Shut down.  Delegates to :meth:`AmplitudeAI.shutdown`."""
        self._ai.shutdown()


# ======================================================================
# Session Context Manager
# ======================================================================


class Session:
    """Context manager for session lifecycle.

    Created via :meth:`BoundAgent.session`.  Provides ``track_*``
    methods that automatically inject ``session_id`` and ``trace_id``.
    Calls ``track_session_end`` on exit (including on exception).

    When ``auto_flush`` is ``True`` (default in serverless environments),
    all pending events are flushed before the context manager exits.
    This prevents event loss when the serverless runtime freezes or
    terminates after the handler returns.

    Example::

        with agent.session("sess-1") as s:
            s.new_trace()
            msg = s.track_user_message(content="Hello")
            ai_msg = s.track_ai_message(
                content="Hi!", model="gpt-4o",
                provider="openai", latency_ms=200,
            )
            s.set_enrichments(enrichments)
        # auto track_session_end with enrichments + auto flush in serverless
    """

    __slots__ = (
        "session_id", "trace_id", "idle_timeout_minutes", "user_id",
        "device_id", "browser_session_id", "auto_flush",
        "_agent", "_enrichments", "_ctx_token", "_session_replay_id",
    )

    def __init__(
        self,
        agent: BoundAgent,
        session_id: str | None = None,
        idle_timeout_minutes: int | None = None,
        user_id: str | None = None,
        device_id: str | None = None,
        browser_session_id: str | None = None,
        auto_flush: bool | None = None,
    ) -> None:
        from .serverless import is_serverless

        self.session_id: str = session_id or str(uuid.uuid4())
        self.trace_id: str | None = None
        self.idle_timeout_minutes = idle_timeout_minutes
        self.user_id = user_id
        self.device_id = device_id or agent._defaults.get("device_id")
        self.browser_session_id = browser_session_id or agent._defaults.get("browser_session_id")
        self.auto_flush: bool = auto_flush if auto_flush is not None else is_serverless()
        self._agent = agent
        self._enrichments: SessionEnrichments | None = None
        from contextvars import Token
        self._ctx_token: Token[SessionContext | None] | None = None
        self._session_replay_id: str | None = (
            f"{self.device_id}/{self.browser_session_id}"
            if self.device_id and self.browser_session_id
            else None
        )

    # ------------------------------------------------------------------
    # Context propagation helpers
    # ------------------------------------------------------------------

    def _build_session_context(self) -> SessionContext:
        """Build a :class:`SessionContext` from this session + its agent."""
        defaults = self._agent._defaults
        ai = self._agent._ai
        sid = self.session_id

        return SessionContext(
            session_id=sid,
            trace_id=self.trace_id,
            user_id=self.user_id or defaults.get("user_id"),
            agent_id=defaults.get("agent_id"),
            parent_agent_id=defaults.get("parent_agent_id"),
            env=defaults.get("env"),
            customer_org_id=defaults.get("customer_org_id"),
            agent_version=defaults.get("agent_version"),
            description=defaults.get("description"),
            context=defaults.get("context"),
            groups=defaults.get("groups"),
            idle_timeout_minutes=self.idle_timeout_minutes,
            device_id=self.device_id or defaults.get("device_id"),
            browser_session_id=self.browser_session_id or defaults.get("browser_session_id"),
            next_turn_id=lambda: ai._next_turn_id(sid),
        )

    def _set_context(self) -> None:
        """Publish this session into the :mod:`contextvars` ContextVar."""
        self._ctx_token = _current_session.set(self._build_session_context())

    def _reset_context(self) -> None:
        """Remove this session from the ContextVar (restore previous)."""
        if self._ctx_token is not None:
            _current_session.reset(self._ctx_token)
            self._ctx_token = None

    # ------------------------------------------------------------------
    # Trace management
    # ------------------------------------------------------------------

    def new_trace(self) -> str:
        """Generate and set a new ``trace_id``.  Returns the new ID.

        Also refreshes the ContextVar so that provider wrappers called
        after this point see the updated ``trace_id``.
        """
        self.trace_id = str(uuid.uuid4())
        # Refresh the ContextVar with the new trace_id.
        ctx = _current_session.get()
        if ctx is not None:
            ctx.trace_id = self.trace_id
        return self.trace_id

    # ------------------------------------------------------------------
    # Enrichments
    # ------------------------------------------------------------------

    def set_enrichments(self, enrichments: SessionEnrichments) -> None:
        """Set enrichments to attach when the session ends."""
        self._enrichments = enrichments

    # ------------------------------------------------------------------
    # Context manager
    # ------------------------------------------------------------------

    def __enter__(self) -> Session:
        self._set_context()
        return self

    def __exit__(self, *exc: Any) -> None:
        self._reset_context()
        try:
            end_kwargs: dict[str, Any] = {
                "session_id": self.session_id,
                "enrichments": self._enrichments,
                "idle_timeout_minutes": self.idle_timeout_minutes,
            }
            if self.user_id is not None:
                end_kwargs["user_id"] = self.user_id
            self._agent.track_session_end(**self._inject(end_kwargs))
        except Exception:
            import logging
            logging.getLogger(__name__).debug(
                "Failed to auto-end session %s", self.session_id, exc_info=True,
            )
        if self.auto_flush:
            self._flush_sync()

    async def __aenter__(self) -> Session:
        self._set_context()
        return self

    async def __aexit__(self, *exc: Any) -> None:
        self._reset_context()
        try:
            end_kwargs: dict[str, Any] = {
                "session_id": self.session_id,
                "enrichments": self._enrichments,
                "idle_timeout_minutes": self.idle_timeout_minutes,
            }
            if self.user_id is not None:
                end_kwargs["user_id"] = self.user_id
            self._agent.track_session_end(**self._inject(end_kwargs))
        except Exception:
            import logging
            logging.getLogger(__name__).debug(
                "Failed to auto-end session %s", self.session_id, exc_info=True,
            )
        if self.auto_flush:
            await self._flush_async()

    def _flush_sync(self) -> None:
        """Flush pending events, blocking until delivery completes."""
        try:
            futures = self._agent.flush()
            for f in futures:
                if f is not None:
                    f.result()
        except Exception:
            import logging
            logging.getLogger(__name__).debug(
                "Failed to flush after session %s", self.session_id, exc_info=True,
            )

    async def _flush_async(self) -> None:
        """Flush pending events without blocking the event loop."""
        import asyncio

        try:
            futures = self._agent.flush()
            for f in futures:
                if f is not None:
                    await asyncio.wrap_future(f)
        except Exception:
            import logging
            logging.getLogger(__name__).debug(
                "Failed to flush after session %s", self.session_id, exc_info=True,
            )

    # ------------------------------------------------------------------
    # Child-agent delegation
    # ------------------------------------------------------------------

    @contextmanager
    def run_as(self, child_agent: BoundAgent) -> Iterator[Session]:
        """Run code as a child agent within this session.

        Provider wrappers automatically pick up the child agent's identity
        (``agent_id``, ``parent_agent_id``) while sharing this session's
        ``session_id``, ``trace_id``, and turn counter.  No
        ``track_session_end`` is emitted on exit.

        Example::

            with parent_session.run_as(recipe_agent) as cs:
                # provider calls see agent_id='recipe-agent'
                response = client.chat.completions.create(...)
                cs.track_tool_call(tool_name="search", latency_ms=50, success=True)
        """
        child_session = Session(
            child_agent,
            session_id=self.session_id,
            user_id=self.user_id,
            device_id=self.device_id,
            browser_session_id=self.browser_session_id,
        )
        child_session.trace_id = self.trace_id
        child_session._set_context()
        try:
            yield child_session
        finally:
            child_session._reset_context()

    @asynccontextmanager
    async def arun_as(self, child_agent: BoundAgent) -> AsyncIterator[Session]:
        """Async version of :meth:`run_as`.

        Example::

            async with parent_session.arun_as(recipe_agent) as cs:
                response = await client.chat.completions.create(...)
        """
        child_session = Session(
            child_agent,
            session_id=self.session_id,
            user_id=self.user_id,
            device_id=self.device_id,
            browser_session_id=self.browser_session_id,
        )
        child_session.trace_id = self.trace_id
        child_session._set_context()
        try:
            yield child_session
        finally:
            child_session._reset_context()

    # ------------------------------------------------------------------
    # Internal merge
    # ------------------------------------------------------------------

    def _inject(self, kwargs: dict[str, Any]) -> dict[str, Any]:
        """Inject session_id, trace_id, user_id, and session replay ID into kwargs."""
        merged = dict(kwargs)
        if "session_id" not in merged or merged["session_id"] is None:
            merged["session_id"] = self.session_id
        if self.trace_id is not None:
            if "trace_id" not in merged or merged["trace_id"] is None:
                merged["trace_id"] = self.trace_id
        if self.user_id is not None:
            if "user_id" not in merged or merged["user_id"] is None:
                merged["user_id"] = self.user_id
        if self._session_replay_id is not None:
            existing_ep = merged.get("event_properties")
            ep = dict(existing_ep) if isinstance(existing_ep, dict) else {}
            if PROP_SESSION_REPLAY_ID not in ep:
                ep[PROP_SESSION_REPLAY_ID] = self._session_replay_id
                merged["event_properties"] = ep
        return merged

    # ------------------------------------------------------------------
    # Delegating track_* methods
    # ------------------------------------------------------------------

    def track_user_message(self, content: str, **kwargs: Any) -> str:
        """Track a user message within this session."""
        return self._agent.track_user_message(content=content, **self._inject(kwargs))

    def track_ai_message(
        self,
        content: str,
        model: str,
        provider: str,
        latency_ms: float,
        **kwargs: Any,
    ) -> str:
        """Track an AI response within this session."""
        return self._agent.track_ai_message(
            content=content,
            model=model,
            provider=provider,
            latency_ms=latency_ms,
            **self._inject(kwargs),
        )

    def track_tool_call(
        self,
        tool_name: str,
        latency_ms: float,
        success: bool,
        **kwargs: Any,
    ) -> str:
        """Track a tool call within this session."""
        return self._agent.track_tool_call(
            tool_name=tool_name,
            latency_ms=latency_ms,
            success=success,
            **self._inject(kwargs),
        )

    def track_embedding(
        self,
        model: str,
        provider: str,
        latency_ms: float,
        **kwargs: Any,
    ) -> str:
        """Track an embedding within this session."""
        return self._agent.track_embedding(
            model=model, provider=provider, latency_ms=latency_ms, **self._inject(kwargs)
        )

    def track_span(
        self,
        span_name: str,
        latency_ms: float,
        **kwargs: Any,
    ) -> str:
        """Track a span within this session."""
        return self._agent.track_span(
            span_name=span_name, latency_ms=latency_ms, **self._inject(kwargs)
        )

    def score(self, name: str, value: float, target_id: str, **kwargs: Any) -> None:
        """Attach a score within this session."""
        self._agent.score(
            name=name, value=value, target_id=target_id, **self._inject(kwargs)
        )
