"""
OpenAI provider integration for Visibe.

Tracks direct OpenAI API calls (chat.completions.create) by wrapping
the client. Also works with Azure OpenAI and OpenAI-compatible providers
(Groq, Together.ai, DeepSeek) since they use the same openai package.

Two modes:
1. instrument(client) — wrap once, every call auto-sends its own trace
2. track(client, name) — context manager to group multiple calls into one trace
"""
import inspect
import logging
import time
import uuid
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional

from .base import TraceSummary
from ..utils import calculate_cost

logger = logging.getLogger("visibe.openai")

# Check if openai package is available
try:
    import openai
    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False


class OpenAIIntegration:
    """
    OpenAI provider integration for Visibe observability.

    Tracks every chat.completions.create() call with token usage, cost,
    model info, and tool calls.

    Usage (instrument mode — each call = one trace):
        from visibe import Visibe
        from visibe.integrations.openai import OpenAIIntegration
        from openai import OpenAI

        obs = Visibe(api_key="sk_live_abc123")
        tracker = OpenAIIntegration(obs)

        client = OpenAI()
        tracker.instrument(client)

        # Every call is now automatically tracked
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "Hello!"}]
        )

        # When done, restore the original method
        tracker.uninstrument(client)

    Usage (track mode — group multiple calls into one trace):
        with tracker.track(client=client, name="multi-step-chat"):
            response1 = client.chat.completions.create(...)
            response2 = client.chat.completions.create(...)
        # Both calls sent as one trace
    """

    def __init__(self, client):
        """
        Initialize OpenAI integration.

        Args:
            client: Visibe client instance
        """
        if not OPENAI_AVAILABLE:
            raise ImportError(
                "OpenAI integration requires the openai package.\n"
                "Install with: pip install openai>=1.0.0\n"
            )
        self.client = client
        self._instrumented_clients = {}  # client id -> {'chat_create': fn, 'responses_create': fn}

    def instrument(self, openai_client, name=None):
        """
        Instrument an OpenAI client so every chat.completions.create()
        call automatically sends its own trace.

        Args:
            openai_client: An openai.OpenAI or openai.AsyncOpenAI client instance
            name: Optional trace name shown in the Visibe dashboard.
                  If already instrumented without a name, providing a name here will
                  override the default model-based name for future calls.
        """
        client_id = id(openai_client)
        already_instrumented = client_id in self._instrumented_clients
        
        if already_instrumented and not name:
            logger.debug("Client is already instrumented")
            return
        
        if already_instrumented and name:
            # Update name for future calls
            logger.debug(f"Updating trace name to: {name}")
            # Note: name is captured in the closure above, so new wrappers will use it
            # For now, just return silently — the existing wrappers will still use
            # the default model-based name. This is acceptable behavior.
            return

        original_create = openai_client.chat.completions.create
        original_responses_create = None
        if _has_responses_api(openai_client):
            original_responses_create = openai_client.responses.create

        self._instrumented_clients[client_id] = {
            'chat_create': original_create,
            'responses_create': original_responses_create,
        }

        visibe_client = self.client
        integration = self
        trace_name = name
        is_async = inspect.iscoroutinefunction(original_create)

        if is_async:
            async def wrapped_create(*args, **kwargs):
                return await integration._execute_and_trace_async(
                    original_create, visibe_client, None, args, kwargs,
                    name=trace_name,
                )
        else:
            def wrapped_create(*args, **kwargs):
                return integration._execute_and_trace(
                    original_create, visibe_client, None, args, kwargs,
                    name=trace_name,
                )

        openai_client.chat.completions.create = wrapped_create

        if original_responses_create is not None:
            if is_async:
                async def wrapped_responses_create(*args, **kwargs):
                    return await integration._execute_and_trace_response_async(
                        original_responses_create, visibe_client, None,
                        args, kwargs, name=trace_name,
                    )
            else:
                def wrapped_responses_create(*args, **kwargs):
                    return integration._execute_and_trace_response(
                        original_responses_create, visibe_client, None,
                        args, kwargs, name=trace_name,
                    )

            openai_client.responses.create = wrapped_responses_create
        logger.debug("Instrumented OpenAI client")

    def uninstrument(self, openai_client):
        """
        Restore the original chat.completions.create() method.

        Args:
            openai_client: The previously instrumented openai.OpenAI client
        """
        client_id = id(openai_client)
        original = self._instrumented_clients.pop(client_id, None)
        if original:
            if isinstance(original, dict):
                if original.get('chat_create') is not None:
                    openai_client.chat.completions.create = original['chat_create']
                if (
                    original.get('responses_create') is not None and
                    _has_responses_api(openai_client)
                ):
                    openai_client.responses.create = original['responses_create']
            else:
                openai_client.chat.completions.create = original
            logger.debug("Uninstrumented OpenAI client")
        else:
            logger.warning("Client was not instrumented")

    @contextmanager
    def track(self, client: Any, name: str = "openai-chat"):
        """
        Track multiple OpenAI calls grouped into a single trace.

        Args:
            client: An openai.OpenAI client instance
            name: Human-readable name for this trace

        Yields:
            None

        Example:
            with tracker.track(client=client, name="conversation"):
                r1 = client.chat.completions.create(...)
                r2 = client.chat.completions.create(...)
        """
        tracker = _OpenAIGroupTracker(self.client, client, name, self)
        tracker.start()
        try:
            yield
        except Exception as e:
            tracker.record_error(e)
            raise
        finally:
            tracker.stop()

    @staticmethod
    def _execute_and_trace(original_create, visibe_client, group_tracker,
                           args, kwargs, name=None):
        """
        Execute an OpenAI API call and track it.

        If group_tracker is provided, sends span to the group's trace.
        Otherwise, creates its own trace (create → span → complete).
        """
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        # Extract request data
        model = kwargs.get('model', 'unknown')
        messages = kwargs.get('messages', [])
        stream = kwargs.get('stream', False)

        # Extract prompt from messages
        prompt_text = _extract_prompt(messages)

        # If a LangChain/LangGraph trace is already active, skip standalone
        # OpenAI trace — the LangGraph callback already captures the LLM call.
        try:
            from visibe.integrations.langchain import _active_root_callback
            if _active_root_callback is not None:
                return original_create(*args, **kwargs)
        except ImportError:
            pass

        # If a CrewAI trace is already active, skip standalone OpenAI trace —
        # the CrewAI OTEL exporter already captures the LLM call.
        try:
            from visibe.integrations.crewai import _active_crew_trace
            if _active_crew_trace:
                return original_create(*args, **kwargs)
        except ImportError:
            pass

        # For instrument mode (no group_tracker), create trace header upfront
        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            # Default name to model if not provided
            trace_name = name or f"openai-{model}"
            trace_header = {
                'trace_id': trace_id,
                'name': trace_name,
                'framework': 'openai',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            }
            visibe_client.create_trace(trace_header)

        try:
            if stream:
                return _handle_streaming(
                    original_create, visibe_client, group_tracker, trace_id,
                    args, kwargs, model, messages, prompt_text,
                    call_start, call_timestamp,
                )

            # Non-streaming call
            response = original_create(*args, **kwargs)

            duration_sec = time.time() - call_start
            duration_ms = int(duration_sec * 1000)

            # Extract usage
            input_tokens = 0
            output_tokens = 0
            if hasattr(response, 'usage') and response.usage:
                input_tokens = response.usage.prompt_tokens or 0
                output_tokens = response.usage.completion_tokens or 0

            # Extract actual model from response
            response_model = getattr(response, 'model', model)

            # Extract output text
            output_text = ""
            if hasattr(response, 'choices') and response.choices:
                msg = response.choices[0].message
                if msg and msg.content:
                    output_text = msg.content
                elif msg and hasattr(msg, 'tool_calls') and msg.tool_calls:
                    parts = []
                    for tc in msg.tool_calls:
                        parts.append(f"{tc.function.name}({tc.function.arguments})")
                    output_text = "; ".join(parts)

            # Extract tool calls
            tools_used = _extract_tool_calls(response)

            # Calculate cost
            cost = calculate_cost(response_model, input_tokens, output_tokens)

            call_data = {
                'model': response_model,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'cost': cost,
                'duration_ms': duration_ms,
                'prompt': prompt_text,
                'output': output_text,
                'tools_used': tools_used,
                'timestamp': call_timestamp,
            }

            if group_tracker:
                group_tracker.add_call(call_data)
            else:
                _send_call_as_trace(visibe_client, trace_id, call_data, name=name)

            return response

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)

            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )

            raise

    @staticmethod
    def _execute_and_trace_response(original_create, visibe_client,
                                    group_tracker, args, kwargs, name=None):
        """
        Execute an OpenAI Responses API call and track it.

        If group_tracker is provided, sends span to the group's trace.
        Otherwise, creates its own trace (create → span → complete).
        """
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        model = kwargs.get('model', 'unknown')
        input_data = kwargs.get('input', '')
        stream = kwargs.get('stream', False)

        prompt_text = _extract_response_prompt(input_data)

        # If a LangChain/LangGraph trace is already active, skip standalone trace.
        try:
            from visibe.integrations.langchain import _active_root_callback
            if _active_root_callback is not None:
                return original_create(*args, **kwargs)
        except ImportError:
            pass

        # If a CrewAI trace is already active, skip standalone OpenAI trace.
        try:
            from visibe.integrations.crewai import _active_crew_trace
            if _active_crew_trace:
                return original_create(*args, **kwargs)
        except ImportError:
            pass

        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            # Default name to model if not provided
            trace_name = name or f"openai-{model}"
            trace_header = {
                'trace_id': trace_id,
                'name': trace_name,
                'framework': 'openai',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            }
            visibe_client.create_trace(trace_header)

        try:
            if stream:
                return _handle_response_streaming(
                    original_create, visibe_client, group_tracker, trace_id,
                    args, kwargs, model, prompt_text, call_start, call_timestamp
                )

            response = original_create(*args, **kwargs)

            duration_sec = time.time() - call_start
            duration_ms = int(duration_sec * 1000)

            input_tokens = 0
            output_tokens = 0
            if hasattr(response, 'usage') and response.usage:
                input_tokens = getattr(response.usage, 'input_tokens', 0) or 0
                output_tokens = getattr(response.usage, 'output_tokens', 0) or 0

            output_text = _extract_response_output_text(response)
            tools_used = _extract_response_tools(response)

            response_model = getattr(response, 'model', model)
            cost = calculate_cost(response_model, input_tokens, output_tokens)

            call_data = {
                'model': response_model,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'cost': cost,
                'duration_ms': duration_ms,
                'prompt': prompt_text,
                'output': output_text,
                'tools_used': tools_used,
                'timestamp': call_timestamp,
            }

            if group_tracker:
                group_tracker.add_call(call_data)
            else:
                _send_call_as_trace(visibe_client, trace_id, call_data, name=name)

            return response

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)

            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )

            raise

    @staticmethod
    async def _execute_and_trace_async(original_create, visibe_client,
                                       group_tracker, args, kwargs, name=None):
        """Async version of _execute_and_trace for AsyncOpenAI clients."""
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        model = kwargs.get('model', 'unknown')
        messages = kwargs.get('messages', [])
        stream = kwargs.get('stream', False)
        prompt_text = _extract_prompt(messages)

        # If a LangChain/LangGraph trace is already active, skip standalone trace.
        try:
            from visibe.integrations.langchain import _active_root_callback
            if _active_root_callback is not None:
                return await original_create(*args, **kwargs)
        except ImportError:
            pass

        # If a CrewAI trace is already active, skip standalone OpenAI trace.
        try:
            from visibe.integrations.crewai import _active_crew_trace
            if _active_crew_trace:
                return await original_create(*args, **kwargs)
        except ImportError:
            pass

        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            # Default name to model if not provided
            trace_name = name or f"openai-{model}"
            trace_header = {
                'trace_id': trace_id,
                'name': trace_name,
                'framework': 'openai',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            }
            visibe_client.create_trace(trace_header)

        try:
            if stream:
                stream_obj = await original_create(*args, **kwargs)
                return _AsyncStreamWrapper(
                    stream_obj, visibe_client, group_tracker, trace_id,
                    model, prompt_text, call_start, call_timestamp, name=name,
                )

            response = await original_create(*args, **kwargs)

            duration_sec = time.time() - call_start
            duration_ms = int(duration_sec * 1000)

            input_tokens = 0
            output_tokens = 0
            if hasattr(response, 'usage') and response.usage:
                input_tokens = response.usage.prompt_tokens or 0
                output_tokens = response.usage.completion_tokens or 0

            response_model = getattr(response, 'model', model)
            output_text = ""
            if hasattr(response, 'choices') and response.choices:
                msg = response.choices[0].message
                if msg and msg.content:
                    output_text = msg.content
                elif msg and hasattr(msg, 'tool_calls') and msg.tool_calls:
                    parts = []
                    for tc in msg.tool_calls:
                        parts.append(
                            f"{tc.function.name}({tc.function.arguments})"
                        )
                    output_text = "; ".join(parts)

            tools_used = _extract_tool_calls(response)
            cost = calculate_cost(response_model, input_tokens, output_tokens)

            call_data = {
                'model': response_model,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'cost': cost,
                'duration_ms': duration_ms,
                'prompt': prompt_text,
                'output': output_text,
                'tools_used': tools_used,
                'timestamp': call_timestamp,
            }

            if group_tracker:
                group_tracker.add_call(call_data)
            else:
                _send_call_as_trace(visibe_client, trace_id, call_data, name=name)

            return response

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)

            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )

            raise

    @staticmethod
    async def _execute_and_trace_response_async(original_create, visibe_client,
                                                group_tracker, args, kwargs,
                                                name=None):
        """Async version of _execute_and_trace_response for AsyncOpenAI."""
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        model = kwargs.get('model', 'unknown')
        input_data = kwargs.get('input', '')
        stream = kwargs.get('stream', False)
        prompt_text = _extract_response_prompt(input_data)

        # If a LangChain/LangGraph trace is already active, skip standalone trace.
        try:
            from visibe.integrations.langchain import _active_root_callback
            if _active_root_callback is not None:
                return await original_create(*args, **kwargs)
        except ImportError:
            pass

        # If a CrewAI trace is already active, skip standalone OpenAI trace.
        try:
            from visibe.integrations.crewai import _active_crew_trace
            if _active_crew_trace:
                return await original_create(*args, **kwargs)
        except ImportError:
            pass

        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            # Default name to model if not provided
            trace_name = name or f"openai-{model}"
            trace_header = {
                'trace_id': trace_id,
                'name': trace_name,
                'framework': 'openai',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            }
            visibe_client.create_trace(trace_header)

        try:
            if stream:
                stream_obj = await original_create(*args, **kwargs)
                return _AsyncResponseStreamWrapper(
                    stream_obj, visibe_client, group_tracker, trace_id,
                    model, prompt_text, call_start, call_timestamp, name=name,
                )

            response = await original_create(*args, **kwargs)

            duration_sec = time.time() - call_start
            duration_ms = int(duration_sec * 1000)

            input_tokens = 0
            output_tokens = 0
            if hasattr(response, 'usage') and response.usage:
                input_tokens = getattr(response.usage, 'input_tokens', 0) or 0
                output_tokens = getattr(response.usage, 'output_tokens', 0) or 0

            output_text = _extract_response_output_text(response)
            tools_used = _extract_response_tools(response)
            response_model = getattr(response, 'model', model)
            cost = calculate_cost(response_model, input_tokens, output_tokens)

            call_data = {
                'model': response_model,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'cost': cost,
                'duration_ms': duration_ms,
                'prompt': prompt_text,
                'output': output_text,
                'tools_used': tools_used,
                'timestamp': call_timestamp,
            }

            if group_tracker:
                group_tracker.add_call(call_data)
            else:
                _send_call_as_trace(visibe_client, trace_id, call_data, name=name)

            return response

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)

            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )

            raise


class _OpenAIGroupTracker:
    """Internal tracker for grouping multiple OpenAI calls into one trace."""

    def __init__(self, visibe_client, openai_client, name: str,
                 integration: 'OpenAIIntegration'):
        self.visibe_client = visibe_client
        self.openai_client = openai_client
        self.name = name
        self._integration = integration
        self.trace_id = str(uuid.uuid4())

        self.start_time = None
        self.end_time = None
        self.calls: List[Dict] = []
        self.errors: List[Dict] = []
        self._step_counter = 0

        self._original_create = None
        self._saved_create = None
        self._original_responses_create = None
        self._saved_responses_create = None

    def _next_step_id(self) -> str:
        self._step_counter += 1
        return f"step_{self._step_counter}"

    def start(self):
        """Patch the client and start tracking."""
        self.start_time = datetime.now(timezone.utc)

        # Create trace header on backend
        self.visibe_client.create_trace({
            'trace_id': self.trace_id,
            'name': self.name,
            'framework': 'openai',
            'started_at': self.start_time.isoformat(),
        })

        # Save whatever is currently on the client (may be instrument wrapper)
        self._saved_create = self.openai_client.chat.completions.create
        if _has_responses_api(self.openai_client):
            self._saved_responses_create = self.openai_client.responses.create

        # If client is already instrumented, use the real original to avoid
        # double-tracking (instrument + track both sending traces)
        client_id = id(self.openai_client)
        if client_id in self._integration._instrumented_clients:
            entry = self._integration._instrumented_clients[client_id]
            if isinstance(entry, dict):
                self._original_create = entry.get('chat_create')
                self._original_responses_create = entry.get('responses_create')
            else:
                self._original_create = entry
        else:
            self._original_create = self.openai_client.chat.completions.create
            if _has_responses_api(self.openai_client):
                self._original_responses_create = self.openai_client.responses.create

        tracker = self
        original = self._original_create
        integration_cls = OpenAIIntegration
        is_async = inspect.iscoroutinefunction(original)

        if is_async:
            async def wrapped_create(*args, **kwargs):
                return await integration_cls._execute_and_trace_async(
                    original, tracker.visibe_client, tracker, args, kwargs
                )
        else:
            def wrapped_create(*args, **kwargs):
                return integration_cls._execute_and_trace(
                    original, tracker.visibe_client, tracker, args, kwargs
                )

        self.openai_client.chat.completions.create = wrapped_create

        if self._original_responses_create is not None:
            original_responses = self._original_responses_create

            if is_async:
                async def wrapped_responses_create(*args, **kwargs):
                    return await integration_cls._execute_and_trace_response_async(
                        original_responses, tracker.visibe_client, tracker,
                        args, kwargs
                    )
            else:
                def wrapped_responses_create(*args, **kwargs):
                    return integration_cls._execute_and_trace_response(
                        original_responses, tracker.visibe_client, tracker,
                        args, kwargs
                    )

            self.openai_client.responses.create = wrapped_responses_create

    def stop(self):
        """Restore original method and complete the trace."""
        self.end_time = datetime.now(timezone.utc)

        # Restore what was on the client before track() started
        if self._saved_create:
            self.openai_client.chat.completions.create = self._saved_create
        if self._saved_responses_create is not None:
            self.openai_client.responses.create = self._saved_responses_create

        self._complete_trace()

    def add_call(self, call_data: Dict):
        """Add a completed call — sends span to backend immediately."""
        self.calls.append(call_data)

        span = {
            'span_id': self._next_step_id(),
            'type': 'llm_call',
            'timestamp': call_data['timestamp'].isoformat(),
            'agent_name': 'openai',
            'model': call_data['model'],
            'provider': 'openai',
            'status': 'success',
            'description': f"LLM Call using {call_data['model']}",
            'input_tokens': call_data['input_tokens'],
            'output_tokens': call_data['output_tokens'],
            'cost': call_data['cost'],
            'duration_ms': call_data['duration_ms'],
            'input_text': (call_data.get('prompt') or '')[:1000],
            'output_text': (call_data.get('output') or '')[:1000],
        }
        self.visibe_client.queue_span(self.trace_id, span)

    def record_error(self, exc: Exception):
        """Record an error — sends error span to backend."""
        self.errors.append({
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'error_type': type(exc).__name__,
            'message': str(exc)[:500],
        })

        span = {
            'span_id': self._next_step_id(),
            'type': 'error',
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'description': f"Error: {type(exc).__name__}",
            'error_type': type(exc).__name__,
            'error_message': str(exc)[:500],
        }
        self.visibe_client.queue_span(self.trace_id, span)

    def _complete_trace(self):
        """Complete the trace with summary aggregates."""
        duration_ms = int(
            (self.end_time - self.start_time).total_seconds() * 1000
        )

        total_input = sum(c['input_tokens'] for c in self.calls)
        total_output = sum(c['output_tokens'] for c in self.calls)
        total_cost = round(sum(c['cost'] for c in self.calls), 6)

        prompt = self.calls[0]['prompt'] if self.calls else self.name
        model = self.calls[0]['model'] if self.calls else None
        status = 'failed' if self.errors else 'completed'

        summary = {
            'status': status,
            'ended_at': self.end_time.isoformat(),
            'duration_ms': duration_ms,
            'prompt': prompt,
            'model': model,
            'total_cost': total_cost,
            'total_tokens': total_input + total_output,
            'total_input_tokens': total_input,
            'total_output_tokens': total_output,
        }

        success = self.visibe_client.complete_trace(self.trace_id, summary)
        if success:
            logger.info(f"OpenAI trace completed: {self.trace_id}")
        else:
            logger.error(f"Failed to complete OpenAI trace: {self.trace_id}")

        # Print TraceSummary to stdout for developer feedback
        total_tool_calls = sum(
            len(c.get('tools_used', [])) for c in self.calls
        )
        trace_summary = TraceSummary(
            name=self.name,
            status=status,
            llm_calls=len(self.calls),
            tool_calls=total_tool_calls,
            total_tokens=total_input + total_output,
            total_cost=total_cost,
            duration_s=round(duration_ms / 1000, 1),
            agents=['openai'],
            sent=success,
        )
        print(f"[Visibe] {trace_summary}")


# ------------------------------------------------------------------
# Helper functions
# ------------------------------------------------------------------

def _has_responses_api(openai_client) -> bool:
    """Check if this OpenAI client supports the Responses API."""
    return (
        hasattr(openai_client, 'responses') and
        hasattr(openai_client.responses, 'create')
    )


def _extract_prompt(messages: list) -> str:
    """Extract the user's prompt from the messages list."""
    if not messages:
        return ""

    # Find the last user message
    for msg in reversed(messages):
        if isinstance(msg, dict) and msg.get('role') == 'user':
            content = msg.get('content', '')
            if isinstance(content, str):
                return content[:500]
            # Handle content as list (multimodal)
            if isinstance(content, list):
                text_parts = []
                for part in content:
                    if isinstance(part, dict) and part.get('type') == 'text':
                        text_parts.append(part.get('text', ''))
                return ' '.join(text_parts)[:500]

    return ""


def _extract_tool_calls(response) -> list:
    """Extract tool calls from an OpenAI response."""
    tools = []
    try:
        if not hasattr(response, 'choices') or not response.choices:
            return tools

        msg = response.choices[0].message
        if not msg or not hasattr(msg, 'tool_calls') or not msg.tool_calls:
            return tools

        for tc in msg.tool_calls:
            tool_entry = {
                'tool_name': tc.function.name,
                'agent_name': 'openai',
                'status': 'success',
                'input': tc.function.arguments[:500] if tc.function.arguments else '',
            }
            tools.append(tool_entry)

    except Exception as e:
        logger.debug(f"Failed to extract tool calls: {e}")

    return tools


def _extract_response_prompt(input_data) -> str:
    """Extract prompt text from Responses API input."""
    if isinstance(input_data, str):
        return input_data[:1000]

    if isinstance(input_data, list):
        for msg in reversed(input_data):
            if isinstance(msg, dict) and msg.get('role') == 'user':
                content = msg.get('content', '')
                if isinstance(content, str):
                    return content[:1000]
                return str(content)[:1000]

    return str(input_data)[:1000]


def _extract_response_output_text(response) -> str:
    """Extract text output from Responses API response.output items."""
    parts = []
    if hasattr(response, 'output') and response.output:
        for item in response.output:
            if getattr(item, 'type', None) == 'message':
                for content in getattr(item, 'content', []):
                    if getattr(content, 'type', None) == 'output_text':
                        parts.append(getattr(content, 'text', ''))
    return ''.join(parts)


def _extract_response_tools(response) -> list:
    """Extract tool calls from Responses API output items."""
    tools = []

    if hasattr(response, 'output') and response.output:
        for item in response.output:
            item_type = getattr(item, 'type', None)
            if item_type == 'function_call':
                tools.append({
                    'tool_name': getattr(item, 'name', 'unknown'),
                    'input': getattr(item, 'arguments', ''),
                    'call_id': getattr(item, 'call_id', ''),
                    'status': 'success',
                })
            elif item_type in (
                'web_search_call',
                'code_interpreter_call',
                'file_search_call',
            ):
                tools.append({
                    'tool_name': item_type.replace('_call', ''),
                    'status': getattr(item, 'status', 'completed'),
                })

    return tools


def _send_call_as_trace(visibe_client, trace_id: str, call_data: dict,
                       name: str = None):
    """Send a single-call trace: span + complete (instrument mode)."""
    now = call_data['timestamp']
    end_time = now + timedelta(milliseconds=call_data['duration_ms'])

    # Send LLM span
    span = {
        'span_id': 'step_1',
        'type': 'llm_call',
        'timestamp': now.isoformat(),
        'agent_name': 'openai',
        'model': call_data['model'],
        'provider': 'openai',
        'status': 'success',
        'description': f"LLM Call using {call_data['model']}",
        'input_tokens': call_data['input_tokens'],
        'output_tokens': call_data['output_tokens'],
        'cost': call_data['cost'],
        'duration_ms': call_data['duration_ms'],
        'input_text': (call_data.get('prompt') or '')[:1000],
        'output_text': (call_data.get('output') or '')[:1000],
    }
    visibe_client.queue_span(trace_id, span)

    # Send tool spans if any
    for i, tool in enumerate(call_data.get('tools_used', []), start=2):
        tool_span = {
            'span_id': f'step_{i}',
            'type': 'tool_call',
            'timestamp': now.isoformat(),
            'tool_name': tool['tool_name'],
            'agent_name': tool.get('agent_name', 'openai'),
            'status': tool.get('status', 'success'),
            'input_text': tool.get('input', ''),
            'output_text': tool.get('output', ''),
        }
        visibe_client.queue_span(trace_id, tool_span)

    # Complete trace
    summary = {
        'status': 'completed',
        'ended_at': end_time.isoformat(),
        'duration_ms': call_data['duration_ms'],
        'prompt': call_data.get('prompt') or 'No prompt captured',
        'model': call_data['model'],
        'total_cost': call_data['cost'],
        'total_tokens': call_data['input_tokens'] + call_data['output_tokens'],
        'total_input_tokens': call_data['input_tokens'],
        'total_output_tokens': call_data['output_tokens'],
    }
    success = visibe_client.complete_trace(trace_id, summary)

    # Print TraceSummary to stdout for developer feedback
    trace_summary = TraceSummary(
        name=name or call_data['model'],
        status='completed',
        llm_calls=1,
        tool_calls=len(call_data.get('tools_used', [])),
        total_tokens=call_data['input_tokens'] + call_data['output_tokens'],
        total_cost=call_data['cost'],
        duration_s=round(call_data['duration_ms'] / 1000, 1),
        agents=['openai'],
        sent=success,
    )
    print(f"[Visibe] {trace_summary}")


def _send_error_as_trace(visibe_client, trace_id: str, model: str,
                         prompt: str, exc: Exception,
                         timestamp: datetime, duration_ms: int,
                         name: str = None):
    """Send a failed trace: error span + complete as failed (instrument mode)."""
    end_time = timestamp + timedelta(milliseconds=duration_ms)

    # Send error span
    span = {
        'span_id': 'step_1',
        'type': 'error',
        'timestamp': timestamp.isoformat(),
        'description': f"Error: {type(exc).__name__}",
        'error_type': type(exc).__name__,
        'error_message': str(exc)[:500],
        'duration_ms': duration_ms,
    }
    visibe_client.queue_span(trace_id, span)

    # Complete trace as failed
    summary = {
        'status': 'failed',
        'ended_at': end_time.isoformat(),
        'duration_ms': duration_ms,
        'prompt': prompt or 'No prompt captured',
        'total_cost': 0.0,
        'total_tokens': 0,
        'total_input_tokens': 0,
        'total_output_tokens': 0,
    }
    visibe_client.complete_trace(trace_id, summary)


def _handle_streaming(original_create, visibe_client, group_tracker, trace_id,
                      args, kwargs, model, messages, prompt_text,
                      call_start, call_timestamp):
    """
    Handle streaming OpenAI calls.

    Wraps the stream iterator to capture the final usage chunk
    (when stream_options={"include_usage": True}).
    """

    # Ensure stream_options includes usage
    stream_options = kwargs.get('stream_options', {})
    if not stream_options.get('include_usage'):
        if 'stream_options' not in kwargs:
            kwargs['stream_options'] = {'include_usage': True}
        else:
            kwargs['stream_options']['include_usage'] = True

    stream = original_create(*args, **kwargs)

    return _StreamWrapper(
        stream, visibe_client, group_tracker, trace_id, model, prompt_text,
        call_start, call_timestamp,
    )


def _handle_response_streaming(original_create, visibe_client, group_tracker,
                               trace_id, args, kwargs, model, prompt_text,
                               call_start, call_timestamp):
    """Handle streaming Responses API calls."""
    stream = original_create(*args, **kwargs)

    return _ResponseStreamWrapper(
        stream, visibe_client, group_tracker, trace_id, model, prompt_text,
        call_start, call_timestamp,
    )


class _StreamWrapper:
    """Wraps an OpenAI stream to capture usage from the final chunk."""

    def __init__(self, stream, visibe_client, group_tracker, trace_id, model,
                 prompt_text, call_start, call_timestamp, name=None):
        self._stream = stream
        self._visibe_client = visibe_client
        self._group_tracker = group_tracker
        self._trace_id = trace_id
        self._model = model
        self._prompt_text = prompt_text
        self._call_start = call_start
        self._call_timestamp = call_timestamp
        self._name = name
        self._collected_content = []
        self._input_tokens = 0
        self._output_tokens = 0
        self._response_model = model
        self._tools_used = []

    def __iter__(self):
        return self

    def __next__(self):
        try:
            chunk = next(self._stream)
            self._process_chunk(chunk)
            return chunk
        except StopIteration:
            self._finalize()
            raise

    def __enter__(self):
        if hasattr(self._stream, '__enter__'):
            self._stream.__enter__()
        return self

    def __exit__(self, *args):
        self._finalize()
        if hasattr(self._stream, '__exit__'):
            return self._stream.__exit__(*args)
        return False

    def _process_chunk(self, chunk):
        """Process each streaming chunk."""
        # Capture content
        if hasattr(chunk, 'choices') and chunk.choices:
            delta = chunk.choices[0].delta
            if delta and hasattr(delta, 'content') and delta.content:
                self._collected_content.append(delta.content)

        # Capture usage from final chunk
        if hasattr(chunk, 'usage') and chunk.usage:
            self._input_tokens = chunk.usage.prompt_tokens or 0
            self._output_tokens = chunk.usage.completion_tokens or 0

        # Capture model
        if hasattr(chunk, 'model') and chunk.model:
            self._response_model = chunk.model

    def _finalize(self):
        """Send trace data after stream completes."""
        if hasattr(self, '_finalized'):
            return
        self._finalized = True

        duration_sec = time.time() - self._call_start
        duration_ms = int(duration_sec * 1000)
        output_text = ''.join(self._collected_content)

        cost = calculate_cost(
            self._response_model, self._input_tokens, self._output_tokens
        )

        call_data = {
            'model': self._response_model,
            'input_tokens': self._input_tokens,
            'output_tokens': self._output_tokens,
            'cost': cost,
            'duration_ms': duration_ms,
            'prompt': self._prompt_text,
            'output': output_text,
            'tools_used': self._tools_used,
            'timestamp': self._call_timestamp,
        }

        if self._group_tracker:
            self._group_tracker.add_call(call_data)
        else:
            _send_call_as_trace(self._visibe_client, self._trace_id, call_data,
                               name=self._name)


class _ResponseStreamWrapper:
    """Wraps an OpenAI Responses API stream."""

    def __init__(self, stream, visibe_client, group_tracker, trace_id, model,
                 prompt_text, call_start, call_timestamp, name=None):
        self._stream = stream
        self._visibe_client = visibe_client
        self._group_tracker = group_tracker
        self._trace_id = trace_id
        self._model = model
        self._prompt_text = prompt_text
        self._call_start = call_start
        self._call_timestamp = call_timestamp
        self._name = name
        self._collected_content = []
        self._input_tokens = 0
        self._output_tokens = 0
        self._response_model = model
        self._response = None

    def __iter__(self):
        return self

    def __next__(self):
        try:
            event = next(self._stream)
            self._process_event(event)
            return event
        except StopIteration:
            self._finalize()
            raise

    def __enter__(self):
        if hasattr(self._stream, '__enter__'):
            self._stream.__enter__()
        return self

    def __exit__(self, *args):
        self._finalize()
        if hasattr(self._stream, '__exit__'):
            return self._stream.__exit__(*args)
        return False

    def _process_event(self, event):
        event_type = getattr(event, 'type', '')

        if event_type == 'response.output_text.delta':
            self._collected_content.append(getattr(event, 'delta', ''))
        elif event_type == 'response.completed':
            response = getattr(event, 'response', None)
            if response is not None:
                self._response = response
                if hasattr(response, 'usage') and response.usage:
                    self._input_tokens = getattr(
                        response.usage, 'input_tokens', 0
                    ) or 0
                    self._output_tokens = getattr(
                        response.usage, 'output_tokens', 0
                    ) or 0
                if hasattr(response, 'model') and response.model:
                    self._response_model = response.model

    def _finalize(self):
        if hasattr(self, '_finalized'):
            return
        self._finalized = True

        duration_sec = time.time() - self._call_start
        duration_ms = int(duration_sec * 1000)

        output_text = ''.join(self._collected_content)
        tools_used = []
        if self._response is not None:
            if not output_text:
                output_text = _extract_response_output_text(self._response)
            tools_used = _extract_response_tools(self._response)

        cost = calculate_cost(
            self._response_model, self._input_tokens, self._output_tokens
        )

        call_data = {
            'model': self._response_model,
            'input_tokens': self._input_tokens,
            'output_tokens': self._output_tokens,
            'cost': cost,
            'duration_ms': duration_ms,
            'prompt': self._prompt_text,
            'output': output_text,
            'tools_used': tools_used,
            'timestamp': self._call_timestamp,
        }

        if self._group_tracker:
            self._group_tracker.add_call(call_data)
        else:
            _send_call_as_trace(self._visibe_client, self._trace_id, call_data,
                               name=self._name)


# ------------------------------------------------------------------
# Async stream wrappers (for AsyncOpenAI clients)
# ------------------------------------------------------------------

class _AsyncStreamWrapper:
    """Wraps an async OpenAI stream to capture usage from the final chunk."""

    def __init__(self, stream, visibe_client, group_tracker, trace_id, model,
                 prompt_text, call_start, call_timestamp, name=None):
        self._stream = stream
        self._visibe_client = visibe_client
        self._group_tracker = group_tracker
        self._trace_id = trace_id
        self._model = model
        self._prompt_text = prompt_text
        self._call_start = call_start
        self._call_timestamp = call_timestamp
        self._name = name
        self._collected_content = []
        self._input_tokens = 0
        self._output_tokens = 0
        self._response_model = model
        self._tools_used = []

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            chunk = await self._stream.__anext__()
            self._process_chunk(chunk)
            return chunk
        except StopAsyncIteration:
            self._finalize()
            raise

    async def __aenter__(self):
        if hasattr(self._stream, '__aenter__'):
            await self._stream.__aenter__()
        return self

    async def __aexit__(self, *args):
        self._finalize()
        if hasattr(self._stream, '__aexit__'):
            return await self._stream.__aexit__(*args)
        return False

    def _process_chunk(self, chunk):
        """Process each streaming chunk."""
        if hasattr(chunk, 'choices') and chunk.choices:
            delta = chunk.choices[0].delta
            if delta and hasattr(delta, 'content') and delta.content:
                self._collected_content.append(delta.content)

        if hasattr(chunk, 'usage') and chunk.usage:
            self._input_tokens = chunk.usage.prompt_tokens or 0
            self._output_tokens = chunk.usage.completion_tokens or 0

        if hasattr(chunk, 'model') and chunk.model:
            self._response_model = chunk.model

    def _finalize(self):
        """Send trace data after stream completes."""
        if hasattr(self, '_finalized'):
            return
        self._finalized = True

        duration_sec = time.time() - self._call_start
        duration_ms = int(duration_sec * 1000)
        output_text = ''.join(self._collected_content)

        cost = calculate_cost(
            self._response_model, self._input_tokens, self._output_tokens
        )

        call_data = {
            'model': self._response_model,
            'input_tokens': self._input_tokens,
            'output_tokens': self._output_tokens,
            'cost': cost,
            'duration_ms': duration_ms,
            'prompt': self._prompt_text,
            'output': output_text,
            'tools_used': self._tools_used,
            'timestamp': self._call_timestamp,
        }

        if self._group_tracker:
            self._group_tracker.add_call(call_data)
        else:
            _send_call_as_trace(self._visibe_client, self._trace_id, call_data,
                               name=self._name)


class _AsyncResponseStreamWrapper:
    """Wraps an async OpenAI Responses API stream."""

    def __init__(self, stream, visibe_client, group_tracker, trace_id, model,
                 prompt_text, call_start, call_timestamp, name=None):
        self._stream = stream
        self._visibe_client = visibe_client
        self._group_tracker = group_tracker
        self._trace_id = trace_id
        self._model = model
        self._prompt_text = prompt_text
        self._call_start = call_start
        self._call_timestamp = call_timestamp
        self._name = name
        self._collected_content = []
        self._input_tokens = 0
        self._output_tokens = 0
        self._response_model = model
        self._response = None

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            event = await self._stream.__anext__()
            self._process_event(event)
            return event
        except StopAsyncIteration:
            self._finalize()
            raise

    async def __aenter__(self):
        if hasattr(self._stream, '__aenter__'):
            await self._stream.__aenter__()
        return self

    async def __aexit__(self, *args):
        self._finalize()
        if hasattr(self._stream, '__aexit__'):
            return await self._stream.__aexit__(*args)
        return False

    def _process_event(self, event):
        event_type = getattr(event, 'type', '')

        if event_type == 'response.output_text.delta':
            self._collected_content.append(getattr(event, 'delta', ''))
        elif event_type == 'response.completed':
            response = getattr(event, 'response', None)
            if response is not None:
                self._response = response
                if hasattr(response, 'usage') and response.usage:
                    self._input_tokens = getattr(
                        response.usage, 'input_tokens', 0
                    ) or 0
                    self._output_tokens = getattr(
                        response.usage, 'output_tokens', 0
                    ) or 0
                if hasattr(response, 'model') and response.model:
                    self._response_model = response.model

    def _finalize(self):
        if hasattr(self, '_finalized'):
            return
        self._finalized = True

        duration_sec = time.time() - self._call_start
        duration_ms = int(duration_sec * 1000)

        output_text = ''.join(self._collected_content)
        tools_used = []
        if self._response is not None:
            if not output_text:
                output_text = _extract_response_output_text(self._response)
            tools_used = _extract_response_tools(self._response)

        cost = calculate_cost(
            self._response_model, self._input_tokens, self._output_tokens
        )

        call_data = {
            'model': self._response_model,
            'input_tokens': self._input_tokens,
            'output_tokens': self._output_tokens,
            'cost': cost,
            'duration_ms': duration_ms,
            'prompt': self._prompt_text,
            'output': output_text,
            'tools_used': tools_used,
            'timestamp': self._call_timestamp,
        }

        if self._group_tracker:
            self._group_tracker.add_call(call_data)
        else:
            _send_call_as_trace(self._visibe_client, self._trace_id, call_data,
                               name=self._name)