"""
OpenAI provider integration for Amplitude AI SDK.

Wraps the OpenAI Python SDK to automatically track LLM usage in Amplitude.
"""

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

from amplitude import Amplitude
from .base import BaseAIProvider
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


class OpenAI(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: Amplitude,
        api_key: Optional[str] = None,
        privacy_mode: bool = False,
        **openai_kwargs
    ):
        """
        Initialize OpenAI client with Amplitude tracking.
        
        Args:
            amplitude: Amplitude client instance
            api_key: OpenAI API key
            privacy_mode: Enable privacy mode (content hashing)
            **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)
        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.
        
        This allows access to APIs we haven't wrapped yet (embeddings, images, etc.)
        They will work but won't be tracked in Amplitude.
        """
        # Avoid recursion on internal attributes
        if name.startswith('_') or name in ('amplitude', 'privacy_config', 'provider_name'):
            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
        # Forward to parent OpenAI client
        return getattr(super(OpenAIClient, self), name)
    
    def _get_provider_name(self) -> str:
        return "openai"
    
    def _extract_model_name(self, **kwargs) -> str:
        return kwargs.get('model', 'unknown')
    
    def _extract_messages(self, **kwargs) -> List[Dict[str, Any]]:
        return kwargs.get('messages', [])

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

        Chat Completions: ``messages[0].role in ('system', 'developer')``
        Responses API:    ``kwargs['instructions']``
        """
        # Responses API passes system instructions as a top-level kwarg
        instructions = kwargs.get("instructions")
        if instructions and isinstance(instructions, str):
            return instructions
        # Fall back to Chat Completions messages-based extraction
        return super()._extract_system_prompt(**kwargs)
    
    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()
            
            # Extract text and apply universal sanitization
            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:
                    # Keep structure simple for tracking
                    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 response."""
        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 response."""
        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.

        OpenAI o1/o3 responses expose reasoning via
        ``response.choices[0].message.reasoning_content``.
        """
        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 response."""
        if hasattr(response, 'choices') and response.choices:
            choice = response.choices[0]
            if hasattr(choice, 'message') and hasattr(choice.message, 'tool_calls'):
                tool_calls = choice.message.tool_calls
                if tool_calls:
                    return [
                        {
                            'type': tool_call.type,
                            'id': tool_call.id,
                            'function': {
                                'name': tool_call.function.name,
                                'arguments': tool_call.function.arguments
                            }
                        }
                        for tool_call in tool_calls
                    ]
        return None


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)

        # Generate missing IDs
        if not amplitude_trace_id:
            amplitude_trace_id = str(uuid.uuid4())
        if not amplitude_conversation_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,
                    )
                
                # 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
        
        # Process stream with defensive error handling for wrapper logic
        try:
            for chunk in stream:
                # Process chunk with defensive wrapper logic
                try:
                    # Capture the versioned model name from the response
                    # (e.g., "gpt-4o-2024-11-20" instead of "gpt-4o").
                    if not getattr(streaming_tracker, '_response_model', None):
                        _rm = getattr(chunk, 'model', None)
                        if _rm and isinstance(_rm, str):
                            streaming_tracker._response_model = _rm

                    # Extract content from chunk
                    if 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 content from o-series models
                        reasoning = getattr(delta, 'reasoning_content', None)
                        if reasoning:
                            streaming_tracker.add_reasoning_chunk(reasoning)
                        
                        # Check for tool calls
                        if hasattr(delta, 'tool_calls') and delta.tool_calls:
                            for tool_call in delta.tool_calls:
                                if tool_call.function:
                                    streaming_tracker.add_tool_call({
                                        'type': 'function',
                                        'id': tool_call.id or '',
                                        'function': {
                                            'name': tool_call.function.name or '',
                                            'arguments': tool_call.function.arguments or ''
                                        }
                                    })
                    
                    # Update token usage if available
                    if hasattr(chunk, 'usage') and chunk.usage:
                        usage = chunk.usage
                        token_update = {}
                        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
                        
                        # Handle reasoning tokens if available
                        if hasattr(usage, 'completion_tokens_details'):
                            details = usage.completion_tokens_details
                            if hasattr(details, 'reasoning_tokens'):
                                token_update['reasoning_tokens'] = details.reasoning_tokens
                        
                        # Handle cache tokens if available
                        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)
                    
                    # Check for finish reason
                    if chunk.choices and chunk.choices[0].finish_reason:
                        streaming_tracker.set_finish_reason(chunk.choices[0].finish_reason)
                        
                except Exception as tracking_error:
                    # Log wrapper errors but don't break the stream
                    get_logger(self.client.amplitude).warning(f"Tracking failed for chunk: {tracking_error}")
                
                # Always yield the original chunk regardless of tracking success/failure
                yield chunk
            
            # Prefer the response model name (often versioned, e.g.,
            # "gpt-4o-2024-11-20") over the request model name ("gpt-4o").
            model_name = (
                getattr(streaming_tracker, '_response_model', None)
                or kwargs.get('model', 'unknown')
            )
            if streaming_tracker.token_usage:
                try:
                    total_cost_usd = calculate_cost(
                        model_name=model_name,
                        input_tokens=streaming_tracker.token_usage.get('input_tokens', 0),
                        output_tokens=streaming_tracker.token_usage.get('output_tokens', 0),
                        reasoning_tokens=streaming_tracker.token_usage.get('reasoning_tokens', 0),
                        cache_read_input_tokens=streaming_tracker.token_usage.get('cache_read_input_tokens', 0),
                        cache_creation_input_tokens=streaming_tracker.token_usage.get('cache_creation_input_tokens', 0),
                    )
                    streaming_tracker.token_usage['total_cost_usd'] = total_cost_usd
                except Exception as cost_error:
                    get_logger(self.client.amplitude).warning(f"Cost calculation failed: {cost_error}")
            
            # Finalize tracking
            try:
                streaming_tracker.finalize(model_name)
            except Exception as finalize_error:
                get_logger(self.client.amplitude).warning(f"Streaming tracker finalization failed: {finalize_error}")
                
        except Exception as e:
            # This catches errors in the stream iteration itself (network issues, etc.)
            # These could be from OpenAI or from our processing, but we should still track them
            try:
                streaming_tracker.set_error(str(e))
                streaming_tracker.finalize(kwargs.get('model', 'unknown'))
            except Exception as tracking_error:
                get_logger(self.client.amplitude).warning(f"Error tracking failed: {tracking_error}")
            raise
    
    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"))

        # Generate missing IDs
        if not amplitude_trace_id:
            amplitude_trace_id = str(uuid.uuid4())
        if not amplitude_conversation_id:
            amplitude_conversation_id = str(uuid.uuid4())
        
        # Augment event properties with responses-specific fields
        augmented_properties = amplitude_event_properties or {}
        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,
                    )
                
                # 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
        
        # Process stream with defensive error handling
        try:
            for chunk in stream:
                # Process chunk with defensive wrapper logic
                try:
                    # Capture versioned model name from response
                    if not getattr(streaming_tracker, '_response_model', None):
                        _rm = getattr(chunk, 'model', None)
                        if _rm and isinstance(_rm, str):
                            streaming_tracker._response_model = _rm

                    # The responses API has similar streaming format to chat completions
                    # Extract content from chunk
                    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 content from o-series models
                        reasoning = getattr(delta, 'reasoning_content', None)
                        if reasoning:
                            streaming_tracker.add_reasoning_chunk(reasoning)
                        
                        # Check for tool calls
                        if hasattr(delta, 'tool_calls') and delta.tool_calls:
                            for tool_call in delta.tool_calls:
                                if tool_call.function:
                                    streaming_tracker.add_tool_call({
                                        'type': 'function',
                                        'id': tool_call.id or '',
                                        'function': {
                                            'name': tool_call.function.name or '',
                                            'arguments': tool_call.function.arguments or ''
                                        }
                                    })
                    
                    # Update token usage if available
                    if hasattr(chunk, 'usage') and chunk.usage:
                        usage = chunk.usage
                        token_update = {}
                        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
                        
                        # Handle reasoning tokens if available
                        if hasattr(usage, 'completion_tokens_details'):
                            details = usage.completion_tokens_details
                            if hasattr(details, 'reasoning_tokens'):
                                token_update['reasoning_tokens'] = details.reasoning_tokens
                        
                        # Handle cache tokens if available
                        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)
                    
                    # Check for finish reason
                    if hasattr(chunk, 'choices') and chunk.choices and chunk.choices[0].finish_reason:
                        streaming_tracker.set_finish_reason(chunk.choices[0].finish_reason)
                        
                except Exception as tracking_error:
                    # Log wrapper errors but don't break the stream
                    get_logger(self.client.amplitude).warning(f"Tracking failed for chunk: {tracking_error}")
                
                # Always yield the original chunk
                yield chunk
            
            # Prefer the response model name (often versioned) over request
            model_name = (
                getattr(streaming_tracker, '_response_model', None)
                or kwargs.get('model', 'unknown')
            )
            if streaming_tracker.token_usage:
                try:
                    total_cost_usd = calculate_cost(
                        model_name=model_name,
                        input_tokens=streaming_tracker.token_usage.get('input_tokens', 0),
                        output_tokens=streaming_tracker.token_usage.get('output_tokens', 0),
                        reasoning_tokens=streaming_tracker.token_usage.get('reasoning_tokens', 0),
                        cache_read_input_tokens=streaming_tracker.token_usage.get('cache_read_input_tokens', 0),
                        cache_creation_input_tokens=streaming_tracker.token_usage.get('cache_creation_input_tokens', 0),
                    )
                    streaming_tracker.token_usage['total_cost_usd'] = total_cost_usd
                except Exception as cost_error:
                    get_logger(self.client.amplitude).warning(f"Cost calculation failed: {cost_error}")
            
            # Finalize tracking
            try:
                streaming_tracker.finalize(model_name)
            except Exception as finalize_error:
                get_logger(self.client.amplitude).warning(f"Streaming tracker finalization failed: {finalize_error}")
                
        except Exception as e:
            # This catches errors in the stream iteration itself
            try:
                streaming_tracker.set_error(str(e))
                streaming_tracker.finalize(kwargs.get('model', 'unknown'))
            except Exception as tracking_error:
                get_logger(self.client.amplitude).warning(f"Error tracking failed: {tracking_error}")
            raise
    
    def __getattr__(self, name):
        """Forward other attributes to original responses object."""
        return getattr(self._original, name)