"""
OpenAI provider integration for Amplitude AI SDK.

Wraps the OpenAI Python SDK to automatically track LLM usage in Amplitude.
"""
from __future__ import annotations

import time
import uuid
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
try:
    import openai
    from openai import OpenAI as OpenAIClient
    from openai import AsyncOpenAI as AsyncOpenAIClient
    from openai.types.chat import ChatCompletion, ChatCompletionChunk
    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False
    OpenAIClient = object  # Fallback for type hints
    AsyncOpenAIClient = object  # Fallback for type hints
    ChatCompletion = object  # Fallback for type hints
    ChatCompletionChunk = object  # Fallback for type hints

from amplitude import Amplitude
from .base import BaseAIProvider, AmplitudeOrAI, resolve_amplitude
from ..core.privacy import PrivacyConfig
from ..exceptions import ProviderError
from ..utils.logger import get_logger
from ..utils.tokens import count_message_tokens
from ..utils.costs import calculate_cost


# ---------------------------------------------------------------------------
# Shared extraction mixin — used by both OpenAI and AsyncOpenAI
# ---------------------------------------------------------------------------

class OpenAIExtractionMixin:
    """Methods shared between sync ``OpenAI`` and ``AsyncOpenAI`` wrappers.

    Extracts model names, messages, system prompts, response content,
    finish reasons, reasoning content, tool calls, and applies privacy
    formatting — all from the same OpenAI Chat / Responses API shapes.

    Requires the composing class to provide ``self.privacy_config: PrivacyConfig``
    (satisfied by ``BaseAIProvider``).
    """

    privacy_config: PrivacyConfig

    def _get_provider_name(self) -> str:
        return "openai"

    def _extract_model_name(self, **kwargs: Any) -> str:
        return kwargs.get("model", "unknown")

    def _extract_messages(self, **kwargs: Any) -> List[Dict[str, Any]]:
        return kwargs.get("messages", [])

    def _extract_system_prompt(self, **kwargs: Any) -> Optional[str]:
        """Extract system prompt from OpenAI Chat or Responses API kwargs.

        Chat Completions: ``messages[0].role in ('system', 'developer')``
        Responses API:    ``kwargs['instructions']``
        """
        instructions = kwargs.get("instructions")
        if instructions and isinstance(instructions, str):
            return instructions
        return super()._extract_system_prompt(**kwargs)  # type: ignore[misc]

    def _format_messages_for_tracking(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Format OpenAI messages for tracking (apply privacy settings)."""
        from ..core.privacy import extract_text_from_openai_message, get_text_from_llm_message, sanitize_any_content

        formatted = []
        for msg in messages:
            formatted_msg = msg.copy()
            text_content = extract_text_from_openai_message(msg)
            if text_content:
                sanitized_result = sanitize_any_content(
                    text_content,
                    privacy_mode=self.privacy_config.privacy_mode,
                    redact_pii=self.privacy_config.redact_pii,
                )
                if self.privacy_config.privacy_mode:
                    formatted_msg["content_hash"] = sanitized_result.get("content_hash")
                    formatted_msg.pop("content", None)
                else:
                    formatted_msg["content"] = get_text_from_llm_message(sanitized_result["$llm_message"])
            formatted.append(formatted_msg)
        return formatted

    def _extract_response_content(self, response: Any) -> str:
        """Extract text content from OpenAI Chat or Responses API response."""
        if hasattr(response, "output_text"):
            return response.output_text or ""
        if hasattr(response, "choices") and response.choices:
            choice = response.choices[0]
            if hasattr(choice, "message") and hasattr(choice.message, "content"):
                return choice.message.content or ""
        return ""

    def _extract_finish_reason(self, response: Any) -> Optional[str]:
        """Extract finish reason from OpenAI Chat or Responses API response."""
        if hasattr(response, "status") and not hasattr(response, "choices"):
            return "stop" if response.status == "completed" else response.status
        if hasattr(response, "choices") and response.choices:
            choice = response.choices[0]
            if hasattr(choice, "finish_reason"):
                return choice.finish_reason
        return None

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

        Chat Completions: ``response.choices[0].message.reasoning_content``
        Responses API: ``output`` items with ``type == "reasoning"`` and
        ``summary`` sub-items containing the reasoning text.
        """
        if hasattr(response, "output") and isinstance(response.output, list):
            reasoning_parts: list[str] = []
            for item in response.output:
                if getattr(item, "type", None) == "reasoning":
                    for summary in getattr(item, "summary", []):
                        text = getattr(summary, "text", None)
                        if text:
                            reasoning_parts.append(text)
            if reasoning_parts:
                return "\n".join(reasoning_parts)
        if hasattr(response, "choices") and response.choices:
            message = getattr(response.choices[0], "message", None)
            if message is not None:
                reasoning = getattr(message, "reasoning_content", None)
                if reasoning:
                    return reasoning
        return None

    def _extract_tool_calls(self, response: Any) -> Optional[List[Dict[str, Any]]]:
        """Extract tool calls from OpenAI Chat or Responses API response."""
        if hasattr(response, "output") and isinstance(response.output, list):
            tool_calls = []
            for item in response.output:
                if getattr(item, "type", None) == "function_call":
                    tool_calls.append({
                        "type": "function",
                        "id": getattr(item, "call_id", getattr(item, "id", "")),
                        "function": {
                            "name": getattr(item, "name", ""),
                            "arguments": getattr(item, "arguments", ""),
                        },
                    })
            if tool_calls:
                return tool_calls
        if hasattr(response, "choices") and response.choices:
            choice = response.choices[0]
            if hasattr(choice, "message") and hasattr(choice.message, "tool_calls"):
                tool_calls_raw = choice.message.tool_calls
                if tool_calls_raw:
                    return [
                        {
                            "type": tc.type,
                            "id": tc.id,
                            "function": {"name": tc.function.name, "arguments": tc.function.arguments},
                        }
                        for tc in tool_calls_raw
                    ]
        return None


# ---------------------------------------------------------------------------
# Shared streaming chunk processing
# ---------------------------------------------------------------------------

def _process_streaming_chunk(streaming_tracker: Any, chunk: Any) -> None:
    """Process a single OpenAI streaming chunk, updating the tracker.

    Handles delta content, reasoning, tool calls, token usage, finish reason,
    and response model capture. Used by all four streaming wrappers.
    """
    if not getattr(streaming_tracker, "_response_model", None):
        _rm = getattr(chunk, "model", None)
        if _rm and isinstance(_rm, str):
            streaming_tracker._response_model = _rm

    if hasattr(chunk, "choices") and chunk.choices and chunk.choices[0].delta:
        delta = chunk.choices[0].delta
        if hasattr(delta, "content") and delta.content:
            streaming_tracker.add_content_chunk(delta.content)

        reasoning = getattr(delta, "reasoning_content", None)
        if reasoning:
            streaming_tracker.add_reasoning_chunk(reasoning)

        if hasattr(delta, "tool_calls") and delta.tool_calls:
            for tool_call in delta.tool_calls:
                tc_index = getattr(tool_call, "index", None)
                tc_id = getattr(tool_call, "id", None)
                func = getattr(tool_call, "function", None)
                tc_name = getattr(func, "name", None) if func else None
                tc_args = getattr(func, "arguments", "") if func else ""
                if tc_id and tc_name is not None:
                    while tc_index is not None and tc_index >= len(streaming_tracker.tool_calls):
                        streaming_tracker.tool_calls.append({
                            "type": "function", "id": "", "function": {"name": "", "arguments": ""}
                        })
                    if tc_index is not None:
                        streaming_tracker.tool_calls[tc_index] = {
                            "type": "function",
                            "id": tc_id,
                            "function": {"name": tc_name, "arguments": tc_args or ""},
                        }
                    else:
                        streaming_tracker.add_tool_call({
                            "type": "function",
                            "id": tc_id,
                            "function": {"name": tc_name, "arguments": tc_args or ""},
                        })
                elif tc_args:
                    if tc_index is not None and tc_index < len(streaming_tracker.tool_calls):
                        streaming_tracker.tool_calls[tc_index]["function"]["arguments"] += tc_args
                    elif streaming_tracker.tool_calls:
                        streaming_tracker.tool_calls[-1]["function"]["arguments"] += tc_args

    if hasattr(chunk, "usage") and chunk.usage:
        usage = chunk.usage
        token_update: Dict[str, Any] = {}
        if hasattr(usage, "prompt_tokens"):
            token_update["input_tokens"] = usage.prompt_tokens
        if hasattr(usage, "completion_tokens"):
            token_update["output_tokens"] = usage.completion_tokens
        if hasattr(usage, "total_tokens"):
            token_update["total_tokens"] = usage.total_tokens
        if hasattr(usage, "completion_tokens_details"):
            details = usage.completion_tokens_details
            if hasattr(details, "reasoning_tokens"):
                token_update["reasoning_tokens"] = details.reasoning_tokens
        if hasattr(usage, "prompt_tokens_details"):
            details = usage.prompt_tokens_details
            if hasattr(details, "cached_tokens"):
                token_update["cache_read_input_tokens"] = details.cached_tokens
        streaming_tracker.update_token_usage(token_update)

    if hasattr(chunk, "choices") and chunk.choices and chunk.choices[0].finish_reason:
        streaming_tracker.set_finish_reason(chunk.choices[0].finish_reason)


def _finalize_streaming(streaming_tracker: Any, model_kwarg: str, amplitude: Any) -> None:
    """Finalize a streaming tracker after all chunks have been yielded.

    Cost calculation is handled inside ``streaming_tracker.finalize()``
    so we only need to resolve the model name and delegate.
    """
    model_name = (
        getattr(streaming_tracker, "_response_model", None) or model_kwarg
    )
    try:
        streaming_tracker.finalize(model_name)
    except Exception as finalize_error:
        get_logger(amplitude).warning(f"Streaming tracker finalization failed: {finalize_error}")


class OpenAI(OpenAIExtractionMixin, BaseAIProvider, OpenAIClient):
    """
    OpenAI client wrapper that automatically tracks usage in Amplitude.
    
    This class inherits from both BaseAIProvider (for tracking) and OpenAI's client
    (for API functionality), providing seamless drop-in replacement.
    """
    
    def __init__(
        self,
        amplitude: AmplitudeOrAI,
        api_key: Optional[str] = None,
        privacy_mode: bool = False,
        content_mode: Optional[str] = None,
        **openai_kwargs
    ):
        """
        Initialize OpenAI client with Amplitude tracking.
        
        Args:
            amplitude: ``Amplitude`` client or ``AmplitudeAI`` instance
            api_key: OpenAI API key
            privacy_mode: Enable privacy mode (content hashing). Deprecated
                in favor of *content_mode*.
            content_mode: One of ``"full"``, ``"metadata_only"``,
                ``"customer_enriched"``. When provided, takes precedence
                over *privacy_mode*.
            **openai_kwargs: Additional arguments passed to OpenAI client
        """
        if not OPENAI_AVAILABLE:
            raise ProviderError(
                "OpenAI SDK not installed. Install with: pip install openai"
            )
        
        # Initialize base provider
        privacy_config = PrivacyConfig(privacy_mode=privacy_mode, content_mode=content_mode)
        BaseAIProvider.__init__(
            self,
            amplitude=amplitude,
            privacy_config=privacy_config
        )
        
        # Initialize OpenAI client
        self._propagate_context = openai_kwargs.pop("propagate_context", False)
        OpenAIClient.__init__(self, api_key=api_key, **openai_kwargs)
        
        # Wrap the chat completions interface
        self._original_chat = self.chat
        self.chat = WrappedChat(self)
        
        # Wrap the responses interface (if available)
        if hasattr(self, 'responses'):
            self._original_responses = self.responses
            self.responses = WrappedResponses(self)
    
    def __getattr__(self, name):
        """Forward unimplemented attributes to original OpenAI client."""
        if name.startswith('_') or name in ('amplitude', 'privacy_config', 'provider_name'):
            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
        return getattr(super(OpenAIClient, self), name)


class WrappedChat:
    """Wrapper for OpenAI chat completions that adds tracking."""
    
    def __init__(self, client: OpenAI):
        self.client = client
        self._original = client._original_chat
        self.completions = WrappedCompletions(client)
    
    def __getattr__(self, name):
        """Forward other attributes to original chat object."""
        return getattr(self._original, name)


class WrappedCompletions:
    """Wrapper for chat completions that adds tracking."""
    
    def __init__(self, client: OpenAI):
        self.client = client
        self._original = client._original_chat.completions
    
    def create(
        self,
        *,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ):
        """
        Create a chat completion with automatic Amplitude tracking.
        
        Args:
            amplitude_user_id: User ID for tracking
            amplitude_conversation_id: Conversation ID for linking events
            amplitude_trace_id: Trace ID for linking related events
            amplitude_turn_id: Turn number in conversation
            amplitude_event_properties: Additional event properties
            amplitude_user_properties: User properties to set
            amplitude_groups: Group analytics properties
            **kwargs: Arguments passed to OpenAI's create method
        """
        amplitude_user_id = amplitude_user_id or getattr(self.client, '_default_amplitude_user_id', None)

        # Only auto-generate IDs when no active session context can supply them
        if not amplitude_trace_id:
            amplitude_trace_id = str(uuid.uuid4())
        if not amplitude_conversation_id:
            from ..context import get_active_context
            ctx = get_active_context()
            if ctx is None or not ctx.session_id:
                amplitude_conversation_id = str(uuid.uuid4())
        
        # Check if streaming
        is_streaming = kwargs.get('stream', False)
        
        if is_streaming:
            return self._create_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=amplitude_event_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs
            )
        else:
            return self._create_non_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=amplitude_event_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs
            )
    
    def _create_non_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> ChatCompletion:
        """Handle non-streaming completion with tracking."""
        if getattr(self.client, "_propagate_context", False):
            from ..propagation import inject_context
            kwargs["extra_headers"] = inject_context(kwargs.get("extra_headers"))

        start_time = time.time()
        
        try:
            # Make the actual OpenAI API call
            response = self._original.create(**kwargs)
            
            # Calculate latency
            end_time = time.time()
            latency_ms = (end_time - start_time) * 1000
            
            # Track the completion if user ID is provided
            if amplitude_user_id:
                # Calculate cost -- prefer versioned model name from
                # the response (e.g., "gpt-4o-2024-11-20").
                token_usage = self.client._extract_token_usage(response)
                model_name = (
                    self.client._extract_response_model(response)
                    or kwargs.get('model', 'unknown')
                )
                
                total_cost_usd = None
                if token_usage['input_tokens'] is not None and token_usage['output_tokens'] is not None:
                    total_cost_usd = calculate_cost(
                        model_name=model_name,
                        input_tokens=token_usage['input_tokens'] or 0,
                        output_tokens=token_usage['output_tokens'] or 0,
                        reasoning_tokens=token_usage['reasoning_tokens'] or 0,
                        cache_read_input_tokens=token_usage['cache_read_input_tokens'] or 0,
                        cache_creation_input_tokens=token_usage['cache_creation_input_tokens'] or 0,
                        default_provider='openai',
                    )
                
                # Track the completion
                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=response,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    total_cost_usd=total_cost_usd,
                    **kwargs
                )
            
            return response
            
        except Exception as e:
            # Calculate latency even for errors
            end_time = time.time()
            latency_ms = (end_time - start_time) * 1000
            
            # Track error if user ID is provided
            if amplitude_user_id:
                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=None,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=True,
                    error_message=str(e),
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    **kwargs
                )
            
            # Re-raise the original exception
            raise
    
    def _create_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> Iterator[ChatCompletionChunk]:
        """Handle streaming completion with tracking."""
        if getattr(self.client, "_propagate_context", False):
            from ..propagation import inject_context
            kwargs["extra_headers"] = inject_context(kwargs.get("extra_headers"))

        if not amplitude_user_id:
            # If no user ID, just return the original stream without tracking
            return self._original.create(**kwargs)
        
        # Create streaming tracker
        streaming_tracker = self.client.create_streaming_tracker(
            user_id=amplitude_user_id,
            conversation_id=amplitude_conversation_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            user_properties=amplitude_user_properties,
            groups=amplitude_groups,
            **kwargs
        )
        
        # Ensure OpenAI includes token usage in the final streaming chunk.
        # Without this, chunk.usage is always None and cost tracking is silent no-op.
        if "stream_options" not in kwargs:
            kwargs["stream_options"] = {"include_usage": True}

        # Separate API call from processing to avoid masking OpenAI errors
        try:
            # Make the streaming API call - this is the only part that should fail "legitimately"
            stream = self._original.create(**kwargs)
        except Exception as e:
            # This is definitely an OpenAI API error
            streaming_tracker.set_error(str(e))
            streaming_tracker.finalize(kwargs.get('model', 'unknown'))
            raise
        
        model_fallback = kwargs.get('model', 'unknown')
        try:
            for chunk in stream:
                try:
                    _process_streaming_chunk(streaming_tracker, chunk)
                except Exception as tracking_error:
                    get_logger(self.client.amplitude).warning(f"Tracking failed for chunk: {tracking_error}")
                yield chunk

            _finalize_streaming(streaming_tracker, model_fallback, self.client.amplitude)

        except Exception as e:
            try:
                streaming_tracker.set_error(str(e))
                streaming_tracker.finalize(model_fallback)
            except Exception as tracking_error:
                get_logger(self.client.amplitude).warning(f"Error tracking failed: {tracking_error}")
            raise
        finally:
            try:
                streaming_tracker.finalize_on_abandon(model_fallback)
            except Exception:
                pass
    
    def __getattr__(self, name):
        """Forward other attributes to original completions object."""
        return getattr(self._original, name)


class WrappedResponses:
    """Wrapper for OpenAI responses API that adds tracking."""
    
    def __init__(self, client: OpenAI):
        self.client = client
        self._original = client._original_responses
    
    def create(
        self,
        *,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ):
        """
        Create a response with automatic Amplitude tracking.
        
        This wraps the OpenAI Responses API which supports stateful conversations,
        structured outputs, tool integration, and reasoning controls.
        
        Args:
            amplitude_user_id: User ID for tracking
            amplitude_conversation_id: Conversation ID for linking events
            amplitude_trace_id: Trace ID for linking related events
            amplitude_turn_id: Turn number in conversation
            amplitude_event_properties: Additional event properties
            amplitude_user_properties: User properties to set
            amplitude_groups: Group analytics properties
            **kwargs: Arguments passed to OpenAI's responses.create method
                     Including: model, input, modalities, instructions, tools,
                     previous_response_id, reasoning, response_format, etc.
        """
        amplitude_user_id = amplitude_user_id or getattr(self.client, '_default_amplitude_user_id', None)

        if getattr(self.client, "_propagate_context", False):
            from ..propagation import inject_context
            kwargs["extra_headers"] = inject_context(kwargs.get("extra_headers"))

        # Only auto-generate IDs when no active session context can supply them
        if not amplitude_trace_id:
            amplitude_trace_id = str(uuid.uuid4())
        if not amplitude_conversation_id:
            from ..context import get_active_context
            ctx = get_active_context()
            if ctx is None or not ctx.session_id:
                amplitude_conversation_id = str(uuid.uuid4())
        
        # Augment event properties with responses-specific fields — copy
        # the caller's dict so we don't mutate it in place.
        augmented_properties = dict(amplitude_event_properties) if amplitude_event_properties else {}
        if 'previous_response_id' in kwargs:
            augmented_properties['previous_response_id'] = kwargs['previous_response_id']
        if 'reasoning' in kwargs:
            augmented_properties['reasoning_effort'] = kwargs['reasoning'].get('effort', 'default')
        if 'modalities' in kwargs:
            augmented_properties['modalities'] = kwargs['modalities']
        
        # Check if streaming
        is_streaming = kwargs.get('stream', False)
        
        if is_streaming:
            return self._create_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=augmented_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs
            )
        else:
            return self._create_non_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=augmented_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs
            )
    
    def _create_non_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ):
        """Handle non-streaming response with tracking."""
        start_time = time.time()
        
        try:
            # Make the actual OpenAI API call
            response = self._original.create(**kwargs)
            
            # Calculate latency
            end_time = time.time()
            latency_ms = (end_time - start_time) * 1000
            
            # Track the completion if user ID is provided
            if amplitude_user_id:
                # Calculate cost -- prefer versioned model from response
                token_usage = self.client._extract_token_usage(response)
                model_name = (
                    self.client._extract_response_model(response)
                    or kwargs.get('model', 'unknown')
                )
                
                total_cost_usd = None
                if token_usage['input_tokens'] is not None and token_usage['output_tokens'] is not None:
                    total_cost_usd = calculate_cost(
                        model_name=model_name,
                        input_tokens=token_usage['input_tokens'] or 0,
                        output_tokens=token_usage['output_tokens'] or 0,
                        reasoning_tokens=token_usage['reasoning_tokens'] or 0,
                        cache_read_input_tokens=token_usage['cache_read_input_tokens'] or 0,
                        cache_creation_input_tokens=token_usage['cache_creation_input_tokens'] or 0,
                        default_provider='openai',
                    )
                
                # Track the completion
                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=response,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    total_cost_usd=total_cost_usd,
                    **kwargs
                )
            
            return response
            
        except Exception as e:
            # Calculate latency even for errors
            end_time = time.time()
            latency_ms = (end_time - start_time) * 1000
            
            # Track error if user ID is provided
            if amplitude_user_id:
                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=None,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=True,
                    error_message=str(e),
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    **kwargs
                )
            
            # Re-raise the original exception
            raise
    
    def _create_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs
    ):
        """Handle streaming response with tracking."""
        if not amplitude_user_id:
            # If no user ID, just return the original stream without tracking
            return self._original.create(**kwargs)
        
        # Create streaming tracker
        streaming_tracker = self.client.create_streaming_tracker(
            user_id=amplitude_user_id,
            conversation_id=amplitude_conversation_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            user_properties=amplitude_user_properties,
            groups=amplitude_groups,
            **kwargs
        )
        
        # Separate API call from processing to avoid masking OpenAI errors
        try:
            # Make the streaming API call
            stream = self._original.create(**kwargs)
        except Exception as e:
            # This is definitely an OpenAI API error
            streaming_tracker.set_error(str(e))
            streaming_tracker.finalize(kwargs.get('model', 'unknown'))
            raise
        
        model_fallback = kwargs.get('model', 'unknown')
        try:
            for chunk in stream:
                try:
                    _process_streaming_chunk(streaming_tracker, chunk)
                except Exception as tracking_error:
                    get_logger(self.client.amplitude).warning(f"Tracking failed for chunk: {tracking_error}")
                yield chunk

            _finalize_streaming(streaming_tracker, model_fallback, self.client.amplitude)

        except Exception as e:
            try:
                streaming_tracker.set_error(str(e))
                streaming_tracker.finalize(model_fallback)
            except Exception as tracking_error:
                get_logger(self.client.amplitude).warning(f"Error tracking failed: {tracking_error}")
            raise
        finally:
            try:
                streaming_tracker.finalize_on_abandon(model_fallback)
            except Exception:
                pass
    
    def __getattr__(self, name):
        """Forward other attributes to original responses object."""
        return getattr(self._original, name)


# ---------------------------------------------------------------------------
# Async variants
# ---------------------------------------------------------------------------


class AsyncOpenAI(OpenAIExtractionMixin, BaseAIProvider, AsyncOpenAIClient):
    """Async OpenAI client wrapper that automatically tracks usage in Amplitude.

    Drop-in replacement for ``openai.AsyncOpenAI`` — all tracking calls are
    synchronous (they only queue events in memory), so the only ``async``
    surface is the actual OpenAI API call.
    """

    def __init__(
        self,
        amplitude: AmplitudeOrAI,
        api_key: Optional[str] = None,
        privacy_mode: bool = False,
        content_mode: Optional[str] = None,
        **openai_kwargs,
    ):
        if not OPENAI_AVAILABLE:
            raise ProviderError(
                "OpenAI SDK not installed. Install with: pip install openai"
            )

        privacy_config = PrivacyConfig(privacy_mode=privacy_mode, content_mode=content_mode)
        BaseAIProvider.__init__(self, amplitude=amplitude, privacy_config=privacy_config)

        self._propagate_context = openai_kwargs.pop("propagate_context", False)
        AsyncOpenAIClient.__init__(self, api_key=api_key, **openai_kwargs)

        self._original_chat = self.chat
        self.chat = WrappedAsyncChat(self)

        if hasattr(self, "responses"):
            self._original_responses = self.responses
            self.responses = WrappedAsyncResponses(self)

    def __getattr__(self, name):
        if name.startswith("_") or name in ("amplitude", "privacy_config", "provider_name"):
            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
        return getattr(super(AsyncOpenAIClient, self), name)


class WrappedAsyncChat:
    """Async wrapper for OpenAI chat completions."""

    def __init__(self, client: AsyncOpenAI):
        self.client = client
        self._original = client._original_chat
        self.completions = WrappedAsyncCompletions(client)

    def __getattr__(self, name):
        return getattr(self._original, name)


class WrappedAsyncCompletions:
    """Async wrapper for chat completions that adds tracking."""

    def __init__(self, client: AsyncOpenAI):
        self.client = client
        self._original = client._original_chat.completions

    async def create(
        self,
        *,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        amplitude_user_id = amplitude_user_id or getattr(self.client, "_default_amplitude_user_id", None)

        if not amplitude_trace_id:
            amplitude_trace_id = str(uuid.uuid4())
        if not amplitude_conversation_id:
            from ..context import get_active_context
            ctx = get_active_context()
            if ctx is None or not ctx.session_id:
                amplitude_conversation_id = str(uuid.uuid4())

        is_streaming = kwargs.get("stream", False)

        if is_streaming:
            # No await: _create_streaming is an async generator, not a coroutine.
            return self._create_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=amplitude_event_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs,
            )
        else:
            return await self._create_non_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=amplitude_event_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs,
            )

    async def _create_non_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> ChatCompletion:
        if getattr(self.client, "_propagate_context", False):
            from ..propagation import inject_context
            kwargs["extra_headers"] = inject_context(kwargs.get("extra_headers"))

        start_time = time.time()

        try:
            response = await self._original.create(**kwargs)
            latency_ms = (time.time() - start_time) * 1000

            if amplitude_user_id:
                token_usage = self.client._extract_token_usage(response)
                model_name = self.client._extract_response_model(response) or kwargs.get("model", "unknown")

                total_cost_usd = None
                if token_usage["input_tokens"] is not None and token_usage["output_tokens"] is not None:
                    total_cost_usd = calculate_cost(
                        model_name=model_name,
                        input_tokens=token_usage["input_tokens"] or 0,
                        output_tokens=token_usage["output_tokens"] or 0,
                        reasoning_tokens=token_usage["reasoning_tokens"] or 0,
                        cache_read_input_tokens=token_usage["cache_read_input_tokens"] or 0,
                        cache_creation_input_tokens=token_usage["cache_creation_input_tokens"] or 0,
                        default_provider="openai",
                    )

                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=response,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    total_cost_usd=total_cost_usd,
                    **kwargs,
                )

            return response

        except Exception as e:
            latency_ms = (time.time() - start_time) * 1000
            if amplitude_user_id:
                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=None,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=True,
                    error_message=str(e),
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    **kwargs,
                )
            raise

    async def _create_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> AsyncIterator[ChatCompletionChunk]:
        if getattr(self.client, "_propagate_context", False):
            from ..propagation import inject_context
            kwargs["extra_headers"] = inject_context(kwargs.get("extra_headers"))

        if not amplitude_user_id:
            async for chunk in await self._original.create(**kwargs):
                yield chunk
            return

        streaming_tracker = self.client.create_streaming_tracker(
            user_id=amplitude_user_id,
            conversation_id=amplitude_conversation_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            user_properties=amplitude_user_properties,
            groups=amplitude_groups,
            **kwargs,
        )

        if "stream_options" not in kwargs:
            kwargs["stream_options"] = {"include_usage": True}

        try:
            stream = await self._original.create(**kwargs)
        except Exception as e:
            streaming_tracker.set_error(str(e))
            streaming_tracker.finalize(kwargs.get("model", "unknown"))
            raise

        amplitude_ref = self.client.amplitude
        model_fallback = kwargs.get("model", "unknown")

        try:
            async for chunk in stream:
                try:
                    _process_streaming_chunk(streaming_tracker, chunk)
                except Exception as tracking_error:
                    get_logger(amplitude_ref).warning(f"Tracking failed for chunk: {tracking_error}")
                yield chunk

            _finalize_streaming(streaming_tracker, model_fallback, amplitude_ref)

        except Exception as e:
            try:
                streaming_tracker.set_error(str(e))
                streaming_tracker.finalize(model_fallback)
            except Exception as tracking_error:
                get_logger(amplitude_ref).warning(f"Error tracking failed: {tracking_error}")
            raise
        finally:
            try:
                streaming_tracker.finalize_on_abandon(model_fallback)
            except Exception:
                pass

    def __getattr__(self, name):
        return getattr(self._original, name)


class WrappedAsyncResponses:
    """Async wrapper for OpenAI responses API that adds tracking."""

    def __init__(self, client: AsyncOpenAI):
        self.client = client
        self._original = client._original_responses

    async def create(
        self,
        *,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        amplitude_user_id = amplitude_user_id or getattr(self.client, "_default_amplitude_user_id", None)

        if getattr(self.client, "_propagate_context", False):
            from ..propagation import inject_context
            kwargs["extra_headers"] = inject_context(kwargs.get("extra_headers"))

        if not amplitude_trace_id:
            amplitude_trace_id = str(uuid.uuid4())
        if not amplitude_conversation_id:
            from ..context import get_active_context
            ctx = get_active_context()
            if ctx is None or not ctx.session_id:
                amplitude_conversation_id = str(uuid.uuid4())

        augmented_properties = dict(amplitude_event_properties) if amplitude_event_properties else {}
        if "previous_response_id" in kwargs:
            augmented_properties["previous_response_id"] = kwargs["previous_response_id"]
        if "reasoning" in kwargs:
            augmented_properties["reasoning_effort"] = kwargs["reasoning"].get("effort", "default")
        if "modalities" in kwargs:
            augmented_properties["modalities"] = kwargs["modalities"]

        is_streaming = kwargs.get("stream", False)

        if is_streaming:
            # No await: _create_streaming is an async generator, not a coroutine.
            return self._create_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=augmented_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs,
            )
        else:
            return await self._create_non_streaming(
                amplitude_user_id=amplitude_user_id,
                amplitude_conversation_id=amplitude_conversation_id,
                amplitude_trace_id=amplitude_trace_id,
                amplitude_turn_id=amplitude_turn_id,
                amplitude_event_properties=augmented_properties,
                amplitude_user_properties=amplitude_user_properties,
                amplitude_groups=amplitude_groups,
                **kwargs,
            )

    async def _create_non_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        start_time = time.time()

        try:
            response = await self._original.create(**kwargs)
            latency_ms = (time.time() - start_time) * 1000

            if amplitude_user_id:
                token_usage = self.client._extract_token_usage(response)
                model_name = self.client._extract_response_model(response) or kwargs.get("model", "unknown")

                total_cost_usd = None
                if token_usage["input_tokens"] is not None and token_usage["output_tokens"] is not None:
                    total_cost_usd = calculate_cost(
                        model_name=model_name,
                        input_tokens=token_usage["input_tokens"] or 0,
                        output_tokens=token_usage["output_tokens"] or 0,
                        reasoning_tokens=token_usage["reasoning_tokens"] or 0,
                        cache_read_input_tokens=token_usage["cache_read_input_tokens"] or 0,
                        cache_creation_input_tokens=token_usage["cache_creation_input_tokens"] or 0,
                        default_provider="openai",
                    )

                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=response,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    total_cost_usd=total_cost_usd,
                    **kwargs,
                )

            return response

        except Exception as e:
            latency_ms = (time.time() - start_time) * 1000
            if amplitude_user_id:
                self.client.track_completion(
                    user_id=amplitude_user_id,
                    response=None,
                    latency_ms=latency_ms,
                    conversation_id=amplitude_conversation_id,
                    trace_id=amplitude_trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=True,
                    error_message=str(e),
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    **kwargs,
                )
            raise

    async def _create_streaming(
        self,
        amplitude_user_id: Optional[str] = None,
        amplitude_conversation_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> AsyncIterator[Any]:
        if not amplitude_user_id:
            async for chunk in await self._original.create(**kwargs):
                yield chunk
            return

        streaming_tracker = self.client.create_streaming_tracker(
            user_id=amplitude_user_id,
            conversation_id=amplitude_conversation_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            user_properties=amplitude_user_properties,
            groups=amplitude_groups,
            **kwargs,
        )

        try:
            stream = await self._original.create(**kwargs)
        except Exception as e:
            streaming_tracker.set_error(str(e))
            streaming_tracker.finalize(kwargs.get("model", "unknown"))
            raise

        amplitude_ref = self.client.amplitude
        model_fallback = kwargs.get("model", "unknown")

        try:
            async for chunk in stream:
                try:
                    _process_streaming_chunk(streaming_tracker, chunk)
                except Exception as tracking_error:
                    get_logger(amplitude_ref).warning(f"Tracking failed for chunk: {tracking_error}")
                yield chunk

            _finalize_streaming(streaming_tracker, model_fallback, amplitude_ref)

        except Exception as e:
            try:
                streaming_tracker.set_error(str(e))
                streaming_tracker.finalize(model_fallback)
            except Exception as tracking_error:
                get_logger(amplitude_ref).warning(f"Error tracking failed: {tracking_error}")
            raise
        finally:
            try:
                streaming_tracker.finalize_on_abandon(model_fallback)
            except Exception:
                pass

    def __getattr__(self, name):
        return getattr(self._original, name)