"""
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, 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
    ChatCompletion = object  # Fallback for type hints
    ChatCompletionChunk = 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,
        content_mode: Optional[str] = None,
        **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). 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.
        
        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)

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

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