"""
Base provider class for Amplitude AI SDK.

Defines the interface and common functionality for all AI provider wrappers.
"""

from __future__ import annotations

import time
import uuid
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Union, Iterator, TYPE_CHECKING
from amplitude import Amplitude

from ..context import get_active_context
from ..core.tracking import track_user_message, track_ai_message
from ..core.privacy import PrivacyConfig
from ..utils.costs import calculate_cost
from ..utils.logger import get_logger

if TYPE_CHECKING:
    from ..client import AmplitudeAI

AmplitudeOrAI = Union[Amplitude, "AmplitudeAI"]


def resolve_amplitude(amplitude: AmplitudeOrAI) -> Amplitude:
    """Normalise an ``Amplitude`` or ``AmplitudeAI`` instance to a raw ``Amplitude`` client.

    This lets every provider constructor accept either type, so users can write
    ``OpenAI(amplitude=ai)`` where ``ai`` is an ``AmplitudeAI`` instance.

    When an ``AmplitudeAI`` is passed, the ``_tracking_proxy`` is returned so
    that provider-wrapper events go through the same counting and debug/dry-run
    logic as events sent via the high-level ``track_*`` methods.

    Also accepts duck-typed objects with a ``track`` method (e.g. mocks in tests).
    """
    if isinstance(amplitude, Amplitude):
        return amplitude
    # AmplitudeAI — prefer the tracking proxy so counting/debug hooks apply.
    proxy = getattr(amplitude, "_tracking_proxy", None)
    if proxy is not None:
        return proxy  # type: ignore[return-value]
    amp_attr = getattr(amplitude, "amplitude", None)
    if isinstance(amp_attr, Amplitude):
        return amp_attr
    if callable(getattr(amplitude, "track", None)):
        return amplitude  # type: ignore[return-value]
    raise TypeError(
        f"Expected Amplitude or AmplitudeAI instance, got {type(amplitude).__name__}. "
        "Pass either your Amplitude client or your AmplitudeAI instance."
    )


def _apply_session_context(
    *,
    user_id: str | None,
    session_id: str | None,
    trace_id: str | None,
    turn_id: int | None,
    event_properties: Dict[str, Any] | None,
    groups: Dict[str, Any] | None,
) -> tuple[
    str | None,       # user_id
    str | None,       # session_id
    str | None,       # trace_id
    int | None,       # turn_id
    Dict[str, Any] | None,  # event_properties (with agent context injected)
    Dict[str, Any] | None,  # groups
]:
    """Fill missing tracking fields from the active :class:`SessionContext`.

    Called at the top of ``track_completion()`` and ``create_streaming_tracker()``
    to auto-inherit session/agent context when a ``with agent.session()`` block
    is active.  Explicit values always win.

    Returns a tuple of the (possibly updated) fields.
    """
    ctx = get_active_context()
    if ctx is None:
        return user_id, session_id, trace_id, turn_id, event_properties, groups

    user_id = user_id or ctx.user_id
    session_id = session_id or ctx.session_id
    trace_id = trace_id or ctx.trace_id
    groups = groups or ctx.groups

    if turn_id is None:
        ctx_turn = ctx.next_turn_id()
        if ctx_turn is not None:
            turn_id = ctx_turn

    # Inject agent-level fields into event_properties so they flow through
    # to core.tracking.track_user_message / track_ai_message.
    _agent_props: Dict[str, Any] = {}
    if ctx.agent_id is not None:
        _agent_props["agent_id"] = ctx.agent_id
    if ctx.parent_agent_id is not None:
        _agent_props["parent_agent_id"] = ctx.parent_agent_id
    if ctx.env is not None:
        _agent_props["env"] = ctx.env
    if ctx.customer_org_id is not None:
        _agent_props["customer_org_id"] = ctx.customer_org_id
    if ctx.agent_version is not None:
        _agent_props["agent_version"] = ctx.agent_version
    if ctx.description is not None:
        _agent_props["description"] = ctx.description
    if ctx.context is not None:
        _agent_props["context"] = ctx.context
    if ctx.idle_timeout_minutes is not None:
        from ..core.tracking import PROP_IDLE_TIMEOUT_MINUTES
        _agent_props[PROP_IDLE_TIMEOUT_MINUTES] = ctx.idle_timeout_minutes
    if ctx.device_id and ctx.browser_session_id:
        from ..core.tracking import PROP_SESSION_REPLAY_ID
        _agent_props[PROP_SESSION_REPLAY_ID] = (
            f"{ctx.device_id}/{ctx.browser_session_id}"
        )

    if _agent_props:
        merged = dict(event_properties) if event_properties else {}
        for k, v in _agent_props.items():
            merged.setdefault(k, v)
        event_properties = merged

    return user_id, session_id, trace_id, turn_id, event_properties, groups


class BaseAIProvider(ABC):
    """
    Abstract base class for AI provider integrations.

    Provides common functionality for tracking LLM usage across different providers
    while allowing provider-specific implementations.

    .. note:: **Privacy parameter inconsistency (Issue 9)**

       The original providers (OpenAI, Anthropic, Gemini, AzureOpenAI) accept
       only ``privacy_mode: bool`` — a legacy boolean flag.  Newer providers
       (Bedrock, Mistral) accept ``privacy_config: Optional[PrivacyConfig]``,
       which supports the full 3-tier ``content_mode`` (full /
       metadata_only / customer_enriched).

       Users of the older wrappers cannot use ``content_mode="metadata_only"``
       or ``content_mode="customer_enriched"`` via provider wrappers — only
       the ``AmplitudeAI`` client (manual tracking) supports all tiers.

       A follow-up should add ``privacy_config`` (or ``config: AIConfig``) as
       an alternative parameter in this base class and all existing subclasses
       to unify the interface.
    """

    def __init__(
        self,
        amplitude: AmplitudeOrAI,
        privacy_config: Optional[PrivacyConfig] = None,
        **provider_kwargs
    ):
        """
        Initialize the base provider.

        Args:
            amplitude: An ``Amplitude`` client or ``AmplitudeAI`` instance.
                When an ``AmplitudeAI`` is passed, its underlying
                ``Amplitude`` client is extracted automatically.
            privacy_config: Privacy configuration for content sanitization.
                Subclasses that still expose only ``privacy_mode: bool``
                construct a ``PrivacyConfig`` internally.  See class-level
                note about the planned unification.
            **provider_kwargs: Provider-specific configuration
        """
        self.amplitude = resolve_amplitude(amplitude)
        self.privacy_config = privacy_config or PrivacyConfig()
        self.provider_name = self._get_provider_name()
        self.provider_kwargs = provider_kwargs

    @abstractmethod
    def _get_provider_name(self) -> str:
        """Return the name of this provider (e.g., 'openai', 'anthropic')."""
        pass

    @abstractmethod
    def _extract_model_name(self, **kwargs) -> str:
        """Extract model name from provider-specific kwargs."""
        pass

    @abstractmethod
    def _extract_messages(self, **kwargs) -> List[Dict[str, Any]]:
        """Extract messages from provider-specific format."""
        pass

    @abstractmethod
    def _format_messages_for_tracking(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Format messages for tracking (apply privacy settings)."""
        pass

    @abstractmethod
    def _extract_response_content(self, response: Any) -> str:
        """Extract response content from provider-specific response object."""
        pass

    def _extract_token_usage(self, response: Any) -> Dict[str, Optional[int]]:
        """
        Extract token usage from response - common implementation for most providers.

        Returns dict with keys: input_tokens, output_tokens, total_tokens,
        reasoning_tokens, cache_read_input_tokens, cache_creation_input_tokens
        """
        usage_dict = {
            'input_tokens': None,
            'output_tokens': None,
            'total_tokens': None,
            'reasoning_tokens': None,
            'cache_read_input_tokens': None,
            'cache_creation_input_tokens': None
        }

        # Handle different usage attribute names
        usage = None
        if hasattr(response, 'usage') and response.usage:
            usage = response.usage
        elif hasattr(response, 'usage_metadata') and response.usage_metadata:
            usage = response.usage_metadata

        if not usage:
            return usage_dict

        # Standard fields (handles OpenAI, Anthropic, and Gemini formats)
        for field, attr_names in [
            ('input_tokens', ['input_tokens', 'prompt_tokens', 'prompt_token_count']),
            ('output_tokens', ['output_tokens', 'completion_tokens', 'candidates_token_count']),
            ('total_tokens', ['total_tokens', 'total_token_count']),
            ('cache_read_input_tokens', ['cache_read_input_tokens']),
            ('cache_creation_input_tokens', ['cache_creation_input_tokens'])
        ]:
            for attr_name in attr_names:
                if hasattr(usage, attr_name):
                    value = getattr(usage, attr_name)
                    if value is not None:
                        usage_dict[field] = value
                        break

        # Handle nested token details (OpenAI)
        if hasattr(usage, 'completion_tokens_details'):
            details = usage.completion_tokens_details
            if hasattr(details, 'reasoning_tokens'):
                usage_dict['reasoning_tokens'] = details.reasoning_tokens

        if hasattr(usage, 'prompt_tokens_details'):
            details = usage.prompt_tokens_details
            if hasattr(details, 'cached_tokens'):
                usage_dict['cache_read_input_tokens'] = details.cached_tokens

        # Calculate total if missing but we have input+output
        if usage_dict['total_tokens'] is None and usage_dict['input_tokens'] is not None and usage_dict['output_tokens'] is not None:
            usage_dict['total_tokens'] = usage_dict['input_tokens'] + usage_dict['output_tokens']

        return usage_dict

    @abstractmethod
    def _extract_finish_reason(self, response: Any) -> Optional[str]:
        """Extract finish reason from response."""
        pass

    @abstractmethod
    def _extract_tool_calls(self, response: Any) -> Optional[List[Dict[str, Any]]]:
        """Extract tool calls from response if present."""
        pass

    def _extract_response_model(self, response: Any) -> Optional[str]:
        """Extract the model name from the provider response object.

        Most providers (OpenAI, Anthropic, Gemini, Mistral) include a
        ``model`` field on the response that contains the versioned model
        name (e.g., ``"gpt-4o-2024-11-20"`` instead of ``"gpt-4o"``).
        This captures provider-side model updates that are invisible when
        only recording the request model name.

        Returns the response model string, or *None* if unavailable.
        """
        model = getattr(response, "model", None)
        if model and isinstance(model, str):
            return model
        return None

    def _extract_reasoning_content(self, response: Any) -> Optional[str]:
        """Extract reasoning/thinking content from a response.

        Override in provider subclasses that support reasoning models
        (OpenAI o1/o3, Anthropic extended thinking, Gemini thinking).

        Returns the raw reasoning text, or *None* if not present.
        """
        return None

    def _extract_system_prompt(self, **kwargs: Any) -> Optional[str]:
        """Extract the system prompt from provider-specific kwargs.

        Default implementation checks ``messages[0]`` for
        ``role in ('system', 'developer')``.  Override in subclasses
        where the system prompt is passed differently (e.g. Anthropic's
        top-level ``system`` kwarg, Bedrock's ``system`` list).
        """
        messages = kwargs.get("messages", [])
        if messages and isinstance(messages, list) and len(messages) > 0:
            first = messages[0]
            if isinstance(first, dict) and first.get("role") in ("system", "developer"):
                content = first.get("content", "")
                if isinstance(content, str):
                    return content
                # Handle list-of-blocks content (e.g., OpenAI multimodal)
                if isinstance(content, list):
                    parts = []
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "text":
                            parts.append(block.get("text", ""))
                        elif isinstance(block, str):
                            parts.append(block)
                    return "".join(parts) if parts else None
        return None

    def _extract_tool_definitions(self, **kwargs: Any) -> Optional[List[Dict[str, Any]]]:
        """Extract request-side tool definitions from provider-specific kwargs.

        Default implementation reads ``kwargs["tools"]`` which works for
        OpenAI Chat Completions, Anthropic, and Mistral.  Override in
        subclasses where tool definitions are nested differently
        (e.g. Bedrock ``toolConfig``).
        """
        tools = kwargs.get("tools")
        if tools and isinstance(tools, list):
            return tools  # type: ignore[return-value]
        return None

    def _extract_model_params(self, **kwargs: Any) -> Dict[str, Any]:
        """Extract model configuration parameters from kwargs.

        Returns a dict with keys ``temperature``, ``top_p``,
        ``max_output_tokens``, ``is_streaming``; values are *None*
        when not present.  Override for providers that nest params
        (e.g. Bedrock ``inferenceConfig``, Gemini ``generation_config``).
        """
        return {
            "temperature": kwargs.get("temperature"),
            "top_p": kwargs.get("top_p"),
            "max_output_tokens": kwargs["max_tokens"] if "max_tokens" in kwargs else kwargs.get("max_output_tokens"),
            "is_streaming": kwargs.get("stream"),
        }

    def track_completion(
        self,
        user_id: str,
        response: Any,
        latency_ms: float,
        session_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        turn_id: Optional[int] = None,
        provider_ttfb_ms: Optional[float] = None,
        is_error: bool = False,
        error_message: Optional[str] = None,
        event_properties: Optional[Dict[str, Any]] = None,
        user_properties: Optional[Dict[str, Any]] = None,
        groups: Optional[Dict[str, Any]] = None,
        conversation_id: Optional[str] = None,
        **kwargs
    ) -> None:
        """Track a completion request and response."""
        try:
            # Auto-inherit session/agent context when inside agent.session()
            (
                user_id, session_id, trace_id, turn_id,
                event_properties, groups,
            ) = _apply_session_context(
                user_id=user_id,
                session_id=session_id or conversation_id,
                trace_id=trace_id,
                turn_id=turn_id,
                event_properties=event_properties,
                groups=groups,
            )

            # Prefer the response model name (often versioned, e.g.,
            # "gpt-4o-2024-11-20") over the request model name ("gpt-4o").
            model_name = (
                (self._extract_response_model(response) if not is_error else None)
                or self._extract_model_name(**kwargs)
            )
            response_content = self._extract_response_content(response) if not is_error else ""
            token_usage = self._extract_token_usage(response) if not is_error else {}
            finish_reason = self._extract_finish_reason(response) if not is_error else None
            tool_calls = self._extract_tool_calls(response) if not is_error else None
            reasoning_content = self._extract_reasoning_content(response) if not is_error else None
            system_prompt = self._extract_system_prompt(**kwargs)
            model_params = self._extract_model_params(**kwargs)
            tool_definitions = self._extract_tool_definitions(**kwargs)

            effective_session_id = session_id
            if not trace_id:
                trace_id = str(uuid.uuid4())

            input_messages = self._extract_messages(**kwargs)
            formatted_messages = self._format_messages_for_tracking(input_messages)

            current_turn = turn_id or 1

            # Extract agent-level fields that _apply_session_context may
            # have injected into event_properties so we can pass them as
            # first-class parameters to core.tracking functions.
            # Copy to avoid mutating the caller's dict.
            _ep = dict(event_properties) if event_properties else {}
            _ctx_agent_id = _ep.pop("agent_id", None)
            _ctx_parent_agent_id = _ep.pop("parent_agent_id", None)
            _ctx_env = _ep.pop("env", None)
            _ctx_customer_org_id = _ep.pop("customer_org_id", None)
            _ctx_agent_version = _ep.pop("agent_version", None)
            _ctx_description = _ep.pop("description", None)
            _ctx_context = _ep.pop("context", None)
            # Restore event_properties to None if we emptied it
            if not _ep and event_properties is not None:
                event_properties = None
            else:
                event_properties = _ep or event_properties

            if formatted_messages and hasattr(formatted_messages, '__iter__') and not isinstance(formatted_messages, str):
                # In agent loops each create() re-sends the full conversation.
                # Only track user messages appended after the last assistant/tool
                # reply so the same message isn't autotracked on every iteration.
                last_reply = -1
                for idx, m in enumerate(formatted_messages):
                    if m.get('role') in ('assistant', 'tool', 'model'):
                        last_reply = idx
                new_messages = formatted_messages[last_reply + 1:]
                for msg in new_messages:
                    if msg.get('role') == 'user':
                        track_user_message(
                            amplitude=self.amplitude,
                            user_id=user_id,
                            message_content=msg.get('content', ''),
                            session_id=effective_session_id,
                            trace_id=trace_id,
                            turn_id=current_turn,
                            agent_id=_ctx_agent_id,
                            parent_agent_id=_ctx_parent_agent_id,
                            env=_ctx_env,
                            customer_org_id=_ctx_customer_org_id,
                            agent_version=_ctx_agent_version,
                            description=_ctx_description,
                            context=_ctx_context,
                            message_source="agent" if _ctx_parent_agent_id else "user",
                            event_properties=event_properties,
                            user_properties=user_properties,
                            groups=groups,
                            privacy_config=self.privacy_config
                        )
                        current_turn += 1

            track_ai_message(
                amplitude=self.amplitude,
                user_id=user_id,
                model_name=model_name,
                provider=self.provider_name,
                response_content=response_content,
                latency_ms=latency_ms,
                session_id=effective_session_id,
                trace_id=trace_id,
                turn_id=current_turn,
                agent_id=_ctx_agent_id,
                parent_agent_id=_ctx_parent_agent_id,
                env=_ctx_env,
                customer_org_id=_ctx_customer_org_id,
                agent_version=_ctx_agent_version,
                description=_ctx_description,
                context=_ctx_context,
                provider_ttfb_ms=provider_ttfb_ms,
                is_error=is_error,
                error_message=error_message,
                finish_reason=finish_reason,
                tool_calls=tool_calls,
                reasoning_content=reasoning_content,
                tool_definitions=tool_definitions,
                system_prompt=system_prompt,
                temperature=model_params.get("temperature"),
                max_output_tokens=model_params.get("max_output_tokens"),
                top_p=model_params.get("top_p"),
                is_streaming=model_params.get("is_streaming"),
                prompt_id=kwargs.get("amplitude_prompt_id"),
                event_properties=event_properties,
                user_properties=user_properties,
                groups=groups,
                privacy_config=self.privacy_config,
                total_cost_usd=kwargs.get("total_cost_usd"),
                **{k: v for k, v in token_usage.items() if v is not None}
            )

        except Exception as e:
            if self.privacy_config.validate:
                raise

            import sys
            import traceback as tb

            _logger = get_logger(self.amplitude)
            _logger.error(f"Failed to track completion: {e}")

            if self.privacy_config.debug:
                print(f"[AmplitudeAI] track_completion error: {e}", file=sys.stderr)
                tb.print_exc(file=sys.stderr)
            else:
                _logger.debug(f"Traceback: {tb.format_exc()}")

    def create_streaming_tracker(
        self,
        user_id: str,
        session_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        turn_id: Optional[int] = None,
        event_properties: Optional[Dict[str, Any]] = None,
        user_properties: Optional[Dict[str, Any]] = None,
        groups: Optional[Dict[str, Any]] = None,
        conversation_id: Optional[str] = None,
        **kwargs
    ) -> 'SimpleStreamingTracker':
        """Create a streaming tracker for handling streaming responses."""
        # Auto-inherit session/agent context when inside agent.session()
        (
            user_id, session_id, trace_id, turn_id,
            event_properties, groups,
        ) = _apply_session_context(
            user_id=user_id,
            session_id=session_id or conversation_id,
            trace_id=trace_id,
            turn_id=turn_id,
            event_properties=event_properties,
            groups=groups,
        )

        # Pre-extract model config so it's captured at stream creation time
        system_prompt = self._extract_system_prompt(**kwargs)
        model_params = self._extract_model_params(**kwargs)
        tool_definitions = self._extract_tool_definitions(**kwargs)

        return SimpleStreamingTracker(
            amplitude_client=self.amplitude,
            privacy_config=self.privacy_config,
            provider_name=self.provider_name,
            user_id=user_id,
            session_id=session_id,
            trace_id=trace_id or str(uuid.uuid4()),
            turn_id=turn_id or 1,
            event_properties=event_properties,
            user_properties=user_properties,
            groups=groups,
            _system_prompt=system_prompt,
            _model_params=model_params,
            _tool_definitions=tool_definitions,
            **kwargs
        )


class SimpleStreamingTracker:
    """Streaming tracker that accumulates chunks and emits a single tracking event."""

    def __init__(
        self,
        amplitude_client: Amplitude,
        privacy_config: PrivacyConfig,
        provider_name: str,
        user_id: str,
        session_id: Optional[str],
        trace_id: str,
        turn_id: int,
        event_properties: Optional[Dict[str, Any]] = None,
        user_properties: Optional[Dict[str, Any]] = None,
        groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ):
        self.amplitude = amplitude_client
        self.privacy_config = privacy_config
        self.provider_name = provider_name
        self.user_id = user_id
        self.session_id = session_id
        self.trace_id = trace_id
        self.turn_id = turn_id
        self.event_properties = event_properties
        self.user_properties = user_properties
        self.groups = groups
        self.kwargs = kwargs

        self.content_chunks: List[str] = []
        self.reasoning_chunks: List[str] = []
        self.token_usage: Dict[str, int] = {}
        self.tool_calls: List[Dict[str, Any]] = []
        self.start_time = time.time()
        self.first_byte_time: Optional[float] = None
        self.finish_reason: Optional[str] = None
        self.is_error = False
        self.error_message: Optional[str] = None
        self.stream_abandoned = False
        self._finalized = False

        # Model configuration captured at stream creation time
        self.system_prompt: Optional[str] = kwargs.get("_system_prompt")
        self.model_params: Dict[str, Any] = kwargs.get("_model_params") or {}
        self.tool_definitions: Optional[List[Dict[str, Any]]] = kwargs.get("_tool_definitions")
        self.prompt_id: Optional[str] = kwargs.get("amplitude_prompt_id")

    def add_content_chunk(self, content: str) -> None:
        """Add a content chunk from streaming response."""
        if not self.first_byte_time:
            self.first_byte_time = time.time()
        self.content_chunks.append(content)

    def add_reasoning_chunk(self, content: str) -> None:
        """Add a reasoning/thinking content chunk from streaming response."""
        self.reasoning_chunks.append(content)

    def update_token_usage(self, usage: Dict[str, int]) -> None:
        """Update token usage with new data."""
        self.token_usage.update(usage)

    def add_tool_call(self, tool_call: Dict[str, Any]) -> None:
        """Add a tool call to the accumulated list."""
        self.tool_calls.append(tool_call)

    def set_finish_reason(self, reason: str) -> None:
        """Set the finish reason for the response."""
        self.finish_reason = reason

    def set_error(self, error_message: str) -> None:
        """Mark the stream as having an error."""
        self.is_error = True
        self.error_message = error_message

    def finalize_on_abandon(self, model_name: str) -> None:
        """Finalize an abandoned stream (e.g. consumer broke early or ``GeneratorExit``).

        Marks the stream as abandoned and delegates to :meth:`finalize`.
        No-op if the tracker was already finalized by a normal or error path.

        Uses the response-resolved model name (captured from chunks) when
        available, falling back to the caller-supplied *model_name*.
        """
        if self._finalized:
            return
        self.stream_abandoned = True
        resolved = getattr(self, "_response_model", None) or model_name
        self.finalize(resolved)

    def finalize(self, model_name: str) -> None:
        """Finalize the streaming response and emit tracking events.

        Idempotent: second and subsequent calls are silent no-ops.
        ``_finalized`` is only set to ``True`` *after* the event is
        successfully emitted so that ``finalize_on_abandon()`` can still
        retry if event construction or emission fails.
        """
        if self._finalized:
            return

        end_time = time.time()
        latency_ms = (end_time - self.start_time) * 1000

        provider_ttfb_ms = None
        if self.first_byte_time:
            provider_ttfb_ms = (self.first_byte_time - self.start_time) * 1000

        full_content = ''.join(self.content_chunks)
        reasoning_content = ''.join(self.reasoning_chunks) if self.reasoning_chunks else None

        total_cost_usd = None
        if self.token_usage.get('input_tokens') is not None and self.token_usage.get('output_tokens') is not None:
            cost_params = {
                k: v for k, v in self.token_usage.items()
                if k in ['input_tokens', 'output_tokens', 'reasoning_tokens',
                        'cache_read_input_tokens', 'cache_creation_input_tokens']
            }
            total_cost_usd = calculate_cost(
                model_name=model_name,
                default_provider=self.provider_name,
                **cost_params,
            )

        _ep = dict(self.event_properties) if self.event_properties else {}
        if self.stream_abandoned:
            _ep["[Agent] Stream Abandoned"] = True
        _ctx_agent_id = _ep.pop("agent_id", None)
        _ctx_parent_agent_id = _ep.pop("parent_agent_id", None)
        _ctx_env = _ep.pop("env", None)
        _ctx_customer_org_id = _ep.pop("customer_org_id", None)
        _ctx_agent_version = _ep.pop("agent_version", None)
        _ctx_description = _ep.pop("description", None)
        _ctx_context = _ep.pop("context", None)
        final_ep = _ep if _ep else None

        track_ai_message(
            amplitude=self.amplitude,
            user_id=self.user_id,
            model_name=model_name,
            provider=self.provider_name,
            response_content=full_content,
            latency_ms=latency_ms,
            session_id=self.session_id,
            trace_id=self.trace_id,
            turn_id=self.turn_id,
            agent_id=_ctx_agent_id,
            parent_agent_id=_ctx_parent_agent_id,
            env=_ctx_env,
            customer_org_id=_ctx_customer_org_id,
            agent_version=_ctx_agent_version,
            description=_ctx_description,
            context=_ctx_context,
            provider_ttfb_ms=provider_ttfb_ms,
            is_error=self.is_error,
            error_message=self.error_message,
            finish_reason=self.finish_reason,
            tool_calls=self.tool_calls if self.tool_calls else None,
            reasoning_content=reasoning_content,
            tool_definitions=self.tool_definitions,
            system_prompt=self.system_prompt,
            temperature=self.model_params.get("temperature"),
            max_output_tokens=self.model_params.get("max_output_tokens"),
            top_p=self.model_params.get("top_p"),
            is_streaming=True,
            prompt_id=self.prompt_id,
            total_cost_usd=total_cost_usd,
            event_properties=final_ep,
            user_properties=self.user_properties,
            groups=self.groups,
            privacy_config=self.privacy_config,
            **{k: v for k, v in self.token_usage.items() if k != 'total_cost_usd'}
        )

        # Only mark finalized after the event was successfully emitted.
        self._finalized = True
