"""
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 threading
import uuid
from concurrent.futures import Future
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 .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,
)

# 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 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

        if self._config.debug or self._config.dry_run:
            self._install_debug_hook()

    # ------------------------------------------------------------------
    # Debug / Dry-run hook
    # ------------------------------------------------------------------

    def _install_debug_hook(self) -> None:
        """Wrap ``amplitude.track()`` to print debug output and/or suppress sends.

        .. warning:: **Shared-client clobber risk**

           This method replaces ``self._amplitude.track`` with a wrapper
           function.  If two :class:`AmplitudeAI` instances share the
           same underlying :class:`~amplitude.Amplitude` client, the
           second call to ``_install_debug_hook`` overwrites the first
           instance's hook — only the last-installed hook survives.
           Avoid sharing a single ``Amplitude`` client across multiple
           ``AmplitudeAI`` instances when ``debug=True`` or
           ``dry_run=True``.
        """
        import sys

        from .utils.debug import format_debug_line, format_dry_run_line

        original_track = self._amplitude.track
        debug = self._config.debug
        dry_run = self._config.dry_run

        def _hooked_track(event: Any) -> None:
            if debug:
                print(format_debug_line(event), file=sys.stderr)
            elif dry_run:
                print(format_dry_run_line(event), file=sys.stderr)
            if not dry_run:
                original_track(event)

        self._amplitude.track = _hooked_track

    # ------------------------------------------------------------------
    # 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

        available: list[str] = []
        for mod_name in (
            "openai", "anthropic", "gemini", "azure_openai", "bedrock", "mistral",
        ):
            try:
                importlib.import_module(
                    f".providers.{mod_name}", package="amplitude_ai"
                )
                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,
        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,
        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._amplitude,
            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,
            context=context,
            env=env,
            is_regeneration=is_regeneration,
            is_edit=is_edit,
            edited_message_id=edited_message_id,
            attachments=attachments,
            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,
        # 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,
        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.
            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,
                )
            except Exception:
                pass  # cost calculation is best-effort

        return track_ai_message(
            amplitude=self._amplitude,
            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,
            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,
            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,
        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._amplitude,
            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,
            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,
        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._amplitude,
            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,
            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,
        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._amplitude,
            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,
            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,
        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._amplitude,
            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,
            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,
        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._amplitude,
            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,
            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,
        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._amplitude,
            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,
            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,
        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,
    ) -> 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 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,
            context=context,
            env=env,
            session_id=session_id,
            trace_id=trace_id,
            groups=groups,
        )

    # ------------------------------------------------------------------
    # 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.
        """
        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.
        """
        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",
    "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,
        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,
            "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:
            ep = merged.get("event_properties") or {}
            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,
    ) -> 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

        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``.
        """
        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,
        )

    # ------------------------------------------------------------------
    # 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).

    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
    """

    __slots__ = (
        "session_id", "trace_id", "idle_timeout_minutes", "user_id",
        "device_id", "browser_session_id",
        "_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,
    ) -> None:
        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._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"),
            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,
            )

    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,
            )

    # ------------------------------------------------------------------
    # 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:
            ep = merged.get("event_properties") or {}
            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)
        )
