"""
Framework wrappers for automatic tracing of LLM providers.

Provides drop-in replacements and wrappers for:
- OpenAI (including Azure OpenAI)
- Anthropic Claude
- Google Gemini
- Generic LLM wrapper pattern

Features:
- Zero-code-change integration
- Automatic token tracking
- Cost calculation
- Streaming support
- Error handling and retries
- Context propagation
- Automatic event queueing to platform
"""

import functools
import inspect
import asyncio
from typing import Any, Optional, Dict, Union, Iterator, AsyncIterator, Callable
from datetime import datetime
import logging

from .context_manager import (
    get_current_trace_context,
    set_current_span_context,
    get_current_span_context,
    get_parent_context,
    RunContext,
    is_tracing_enabled,
)
from .decorators_v2 import _extract_inputs, _serialize_output
from .cost_tracking import extract_and_calculate_cost
from .buffer import EventType

logger = logging.getLogger(__name__)


def _extract_system_prompt(messages: list) -> Optional[str]:
    """
    Extract system prompt from messages array.

    Supports various message formats (OpenAI, Anthropic, dict, object).
    Returns the first system message content found.
    """
    if not messages:
        return None

    for msg in messages:
        # Handle dict-style messages
        if isinstance(msg, dict):
            role = msg.get("role", "")
            if role == "system":
                content = msg.get("content", "")
                if isinstance(content, str):
                    return content
                elif isinstance(content, list):
                    # Handle content blocks (e.g., [{"type": "text", "text": "..."}])
                    text_parts = []
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "text":
                            text_parts.append(block.get("text", ""))
                        elif isinstance(block, str):
                            text_parts.append(block)
                    return " ".join(text_parts) if text_parts else None
        # Handle object-style messages (e.g., ChatCompletionMessage)
        elif hasattr(msg, "role"):
            role = getattr(msg, "role", None)
            if role == "system":
                content = getattr(msg, "content", None)
                if isinstance(content, str):
                    return content

    return None


def _update_trace_system_prompt(system_prompt: str) -> None:
    """
    Update the current trace's metadata with the extracted system prompt.

    Only updates if the trace doesn't already have a system prompt set.
    """
    try:
        trace_ctx = get_current_trace_context()
        if trace_ctx and hasattr(trace_ctx, 'metadata'):
            # Only set if not already provided by user
            if "kytte.system_prompt" not in trace_ctx.metadata:
                trace_ctx.metadata["kytte.system_prompt"] = system_prompt
                logger.debug(f"[wrapper] Auto-extracted system prompt ({len(system_prompt)} chars) to trace metadata")
    except Exception as e:
        logger.debug(f"[wrapper] Could not update trace with system prompt: {e}")


async def _queue_llm_span_event(run_ctx: RunContext, trace_id: str) -> None:
    """
    Queue an LLM span event to the global Aigie client buffer.

    This ensures token usage and cost data from LLM wrappers flows to the platform.
    """
    try:
        from .client import get_aigie

        aigie = get_aigie()
        if not aigie or not aigie._buffer:
            logger.debug("[wrapper] No global Aigie client - skipping LLM span event queue")
            return

        end_time = datetime.utcnow()

        # Build span payload with LLM-specific fields
        payload = {
            "id": run_ctx.id,
            "name": run_ctx.name,
            "trace_id": trace_id,
            "parent_span_id": run_ctx.parent_id if run_ctx.parent_id != trace_id else None,
            "type": "llm",  # Always LLM type for wrapper spans
            "start_time": run_ctx.start_time.isoformat() if run_ctx.start_time else end_time.isoformat(),
            "end_time": end_time.isoformat(),
            "metadata": run_ctx.metadata,
            "tags": run_ctx.tags,
            "input": run_ctx.metadata.get("input"),
            "output": run_ctx.metadata.get("output"),
            "status": run_ctx.metadata.get("status", "success"),
            "level": "ERROR" if run_ctx.metadata.get("status") == "error" else "DEFAULT",
        }

        # Add token usage as top-level fields
        if "usage" in run_ctx.metadata:
            usage = run_ctx.metadata["usage"]
            payload["prompt_tokens"] = usage.get("prompt_tokens", 0)
            payload["completion_tokens"] = usage.get("completion_tokens", 0)
            payload["total_tokens"] = usage.get("total_tokens", 0)

        # Add cost as top-level field
        if "cost" in run_ctx.metadata:
            cost = run_ctx.metadata["cost"]
            payload["total_cost"] = cost.get("total_cost", 0)

        # Add model info
        if "model" in run_ctx.metadata:
            payload["model"] = run_ctx.metadata["model"]
        if "provider" in run_ctx.metadata:
            payload["model_provider"] = run_ctx.metadata["provider"]

        # Calculate duration
        if run_ctx.start_time:
            duration_ns = int((end_time - run_ctx.start_time).total_seconds() * 1e9)
            payload["duration"] = duration_ns

        # Queue to buffer
        await aigie._buffer.add(EventType.SPAN_CREATE, payload)
        logger.debug(f"[wrapper] Queued LLM span event: {run_ctx.name} ({run_ctx.id}) -> trace:{trace_id}")

    except Exception as e:
        logger.warning(f"[wrapper] Failed to queue LLM span event: {e}")


def _queue_llm_span_event_sync(run_ctx: RunContext, trace_id: str) -> None:
    """Synchronously queue an LLM span event (schedules async operation)."""
    try:
        loop = asyncio.get_running_loop()
        loop.create_task(_queue_llm_span_event(run_ctx, trace_id))
    except RuntimeError:
        # No running loop - try to run directly
        try:
            asyncio.run(_queue_llm_span_event(run_ctx, trace_id))
        except Exception as e:
            logger.debug(f"[wrapper] Could not queue LLM span event synchronously: {e}")


class OpenAIWrapper:
    """
    Wrapper for OpenAI client with automatic tracing.

    Usage:
        from aigie.wrappers import wrap_openai
        import openai

        client = wrap_openai(openai.OpenAI())

        # All calls are now automatically traced
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": "Hello"}]
        )
    """

    def __init__(self, client: Any, aigie_client: Optional[Any] = None):
        """
        Initialize wrapper.

        Args:
            client: OpenAI client instance
            aigie_client: Optional Aigie client for API calls
        """
        self._client = client
        self._aigie = aigie_client
        self._original_client = client

    def __getattr__(self, name: str) -> Any:
        """Proxy attribute access to wrapped client."""
        attr = getattr(self._client, name)

        # Wrap chat.completions.create
        if name == "chat":
            return self._wrap_chat(attr)

        # Wrap completions.create
        if name == "completions":
            return self._wrap_completions(attr)

        # Wrap embeddings.create
        if name == "embeddings":
            return self._wrap_embeddings(attr)

        return attr

    def _wrap_chat(self, chat_obj: Any) -> Any:
        """Wrap chat completions."""
        class ChatWrapper:
            def __init__(wrapper_self, obj):
                wrapper_self._obj = obj

            def __getattr__(wrapper_self, name: str):
                attr = getattr(wrapper_self._obj, name)
                if name == "completions":
                    return wrapper_self._wrap_completions(attr)
                return attr

            def _wrap_completions(wrapper_self, completions_obj):
                class CompletionsWrapper:
                    def __init__(comp_self, obj):
                        comp_self._obj = obj

                    def __getattr__(comp_self, name: str):
                        attr = getattr(comp_self._obj, name)
                        if name == "create":
                            return self._trace_chat_completion(attr)
                        return attr

                return CompletionsWrapper(completions_obj)

        return ChatWrapper(chat_obj)

    def _trace_chat_completion(self, create_func: Callable) -> Callable:
        """Trace chat completion calls."""
        if inspect.iscoroutinefunction(create_func):
            @functools.wraps(create_func)
            async def async_wrapper(*args, **kwargs):
                return await self._handle_chat_completion_async(create_func, *args, **kwargs)
            return async_wrapper
        else:
            @functools.wraps(create_func)
            def sync_wrapper(*args, **kwargs):
                return self._handle_chat_completion_sync(create_func, *args, **kwargs)
            return sync_wrapper

    async def _handle_chat_completion_async(self, func: Callable, *args, **kwargs) -> Any:
        """Handle async chat completion with tracing and real-time interception."""
        if not is_tracing_enabled():
            return await func(*args, **kwargs)

        from uuid import uuid4
        from .client import get_aigie

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

        # Auto-extract system prompt and update trace metadata
        system_prompt = _extract_system_prompt(messages)
        if system_prompt:
            _update_trace_system_prompt(system_prompt)

        # Create span context
        parent_ctx = get_parent_context()
        trace_ctx = get_current_trace_context()
        run_id = str(uuid4())
        run_ctx = RunContext(
            id=run_id,
            name=f"openai.chat.completions.create",
            type="span",
            span_type="llm",
            parent_id=parent_ctx.id if parent_ctx else None,
            metadata={
                "provider": "openai",
                "model": model,
                "input": {"messages": messages},
            },
            tags=["openai", "llm", model],
            start_time=datetime.utcnow(),
        )

        prev_span_ctx = get_current_span_context()
        set_current_span_context(run_ctx)

        # Real-time interception context
        interception_ctx = None
        aigie = get_aigie()

        try:
            # ==================== PRE-CALL INTERCEPTION ====================
            if aigie and aigie._interceptor_chain:
                try:
                    interception_ctx = await aigie.intercept_pre_call(
                        provider="openai",
                        model=model,
                        messages=messages,
                        trace_id=trace_ctx.id if trace_ctx else None,
                        span_id=run_id,
                        **{k: v for k, v in kwargs.items() if k not in ['messages', 'model']}
                    )

                    # Check if request was blocked
                    from .interceptor.protocols import InterceptionDecision, InterceptionBlockedError
                    if interception_ctx.decision == InterceptionDecision.BLOCK:
                        raise InterceptionBlockedError(
                            reason=interception_ctx.block_reason or "Request blocked by interception",
                            hook_name="pre_call",
                        )

                    # Apply modifications if any
                    if interception_ctx.modified_messages:
                        kwargs['messages'] = interception_ctx.modified_messages
                        messages = interception_ctx.modified_messages
                    if interception_ctx.modified_kwargs:
                        kwargs.update(interception_ctx.modified_kwargs)
                        model = kwargs.get('model', model)

                except ImportError:
                    # Interception modules not available
                    pass

            # ==================== CALL OPENAI ====================
            response = await func(*args, **kwargs)

            # Handle streaming
            if stream:
                return self._wrap_stream_async(response, run_ctx, interception_ctx)

            # Extract response data
            if hasattr(response, 'choices') and len(response.choices) > 0:
                output_message = response.choices[0].message
                output_content = output_message.content if hasattr(output_message, 'content') else str(output_message)
            else:
                output_content = str(response)

            # Track tokens and calculate cost
            if hasattr(response, 'usage'):
                run_ctx.metadata['usage'] = {
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens,
                }

                # Automatic cost tracking
                cost_info = extract_and_calculate_cost(response, 'openai')
                if cost_info:
                    run_ctx.metadata['cost'] = {
                        "input_cost": float(cost_info.input_cost),
                        "output_cost": float(cost_info.output_cost),
                        "total_cost": float(cost_info.total_cost),
                        "currency": cost_info.currency,
                    }

            run_ctx.metadata['output'] = {"content": output_content}
            run_ctx.metadata['status'] = 'success'

            # ==================== POST-CALL INTERCEPTION ====================
            if aigie and aigie._interceptor_chain and interception_ctx:
                try:
                    # Update interception context with cost/token info
                    if 'cost' in run_ctx.metadata:
                        interception_ctx.actual_cost = run_ctx.metadata['cost'].get('total_cost', 0)
                    if 'usage' in run_ctx.metadata:
                        interception_ctx.actual_input_tokens = run_ctx.metadata['usage'].get('prompt_tokens', 0)
                        interception_ctx.actual_output_tokens = run_ctx.metadata['usage'].get('completion_tokens', 0)
                    interception_ctx.response_content = output_content

                    interception_ctx = await aigie.intercept_post_call(
                        ctx=interception_ctx,
                        response=response,
                        error=None,
                    )

                    # Handle retry request
                    from .interceptor.protocols import InterceptionDecision, InterceptionRetryError
                    if interception_ctx.decision == InterceptionDecision.MODIFY and interception_ctx.should_retry:
                        # Apply retry with modified parameters
                        retry_kwargs = interception_ctx.retry_kwargs or kwargs
                        raise InterceptionRetryError(
                            reason="Post-call interception requested retry",
                            retry_kwargs=retry_kwargs,
                        )

                    # Apply response modifications if any
                    if interception_ctx.modified_response:
                        # Return modified response (for content changes)
                        logger.debug("[wrapper] Applying post-call response modification")
                        # Note: We can't easily modify the OpenAI response object,
                        # but we update our tracking metadata
                        run_ctx.metadata['output'] = {"content": interception_ctx.modified_response.get('content', output_content)}
                        run_ctx.metadata['interception_modified'] = True

                except ImportError:
                    pass

            return response

        except Exception as e:
            # Check if it's an interception retry request
            try:
                from .interceptor.protocols import InterceptionRetryError
                if isinstance(e, InterceptionRetryError) and e.retry_kwargs:
                    logger.info(f"[wrapper] Retrying with modified parameters: {e.reason}")
                    # Recursive retry with modified kwargs
                    return await self._handle_chat_completion_async(func, *args, **e.retry_kwargs)
            except ImportError:
                pass

            run_ctx.metadata['error'] = {
                'type': type(e).__name__,
                'message': str(e),
            }
            run_ctx.metadata['status'] = 'error'

            # Post-call interception for errors
            if aigie and aigie._interceptor_chain and interception_ctx:
                try:
                    interception_ctx = await aigie.intercept_post_call(
                        ctx=interception_ctx,
                        response=None,
                        error=e,
                    )

                    # Check if auto-fix suggests retry
                    from .interceptor.protocols import InterceptionDecision
                    if interception_ctx.should_retry and interception_ctx.retry_kwargs:
                        logger.info(f"[wrapper] Auto-fix retry after error")
                        return await self._handle_chat_completion_async(func, *args, **interception_ctx.retry_kwargs)
                except ImportError:
                    pass
                except Exception as intercept_error:
                    logger.debug(f"[wrapper] Post-call error interception failed: {intercept_error}")

            raise

        finally:
            # Queue LLM span event to buffer for API submission
            trace_id = trace_ctx.id if trace_ctx else run_ctx.id
            await _queue_llm_span_event(run_ctx, trace_id)

            set_current_span_context(prev_span_ctx)

    def _handle_chat_completion_sync(self, func: Callable, *args, **kwargs) -> Any:
        """Handle sync chat completion with tracing and basic interception."""
        if not is_tracing_enabled():
            return func(*args, **kwargs)

        from uuid import uuid4
        from .client import get_aigie

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

        # Auto-extract system prompt and update trace metadata
        system_prompt = _extract_system_prompt(messages)
        if system_prompt:
            _update_trace_system_prompt(system_prompt)

        parent_ctx = get_parent_context()
        trace_ctx = get_current_trace_context()
        run_id = str(uuid4())
        run_ctx = RunContext(
            id=run_id,
            name=f"openai.chat.completions.create",
            type="span",
            span_type="llm",
            parent_id=parent_ctx.id if parent_ctx else None,
            metadata={
                "provider": "openai",
                "model": model,
                "input": {"messages": messages},
            },
            tags=["openai", "llm", model],
            start_time=datetime.utcnow(),
        )

        prev_span_ctx = get_current_span_context()
        set_current_span_context(run_ctx)

        # Basic sync interception using rules engine directly
        aigie = get_aigie()

        try:
            # ==================== SYNC PRE-CALL INTERCEPTION ====================
            # For sync calls, we can only use the local rules engine (no async backend)
            if aigie and aigie._rules_engine:
                try:
                    from .interceptor.protocols import InterceptionContext, InterceptionDecision, InterceptionBlockedError

                    # Create interception context
                    interception_ctx = InterceptionContext(
                        provider="openai",
                        model=model,
                        messages=messages,
                        trace_id=trace_ctx.id if trace_ctx else None,
                        span_id=run_id,
                        request_kwargs=kwargs,
                    )

                    # Evaluate local rules synchronously (rules engine is sync)
                    # We can't call the full async chain, but we can check rules
                    loop = None
                    try:
                        loop = asyncio.get_running_loop()
                    except RuntimeError:
                        pass

                    if loop:
                        # If there's a running loop, schedule the evaluation
                        # but we can't block on it in sync code
                        logger.debug("[wrapper] Sync call in async context - rules not evaluated")
                    else:
                        # No running loop - we can run the rules evaluation
                        try:
                            result = asyncio.run(aigie._rules_engine.evaluate(interception_ctx))
                            if result.decision == InterceptionDecision.BLOCK:
                                raise InterceptionBlockedError(
                                    reason=result.reason or "Request blocked by rules",
                                    hook_name="rules_engine",
                                )
                        except Exception as rule_error:
                            if isinstance(rule_error, InterceptionBlockedError):
                                raise
                            logger.debug(f"[wrapper] Sync rules evaluation failed: {rule_error}")

                except ImportError:
                    pass

            # ==================== CALL OPENAI ====================
            response = func(*args, **kwargs)

            if stream:
                return self._wrap_stream_sync(response, run_ctx)

            if hasattr(response, 'choices') and len(response.choices) > 0:
                output_message = response.choices[0].message
                output_content = output_message.content if hasattr(output_message, 'content') else str(output_message)
            else:
                output_content = str(response)

            if hasattr(response, 'usage'):
                run_ctx.metadata['usage'] = {
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens,
                }

                # Automatic cost tracking
                cost_info = extract_and_calculate_cost(response, 'openai')
                if cost_info:
                    run_ctx.metadata['cost'] = {
                        "input_cost": float(cost_info.input_cost),
                        "output_cost": float(cost_info.output_cost),
                        "total_cost": float(cost_info.total_cost),
                        "currency": cost_info.currency,
                    }

            run_ctx.metadata['output'] = {"content": output_content}
            run_ctx.metadata['status'] = 'success'

            return response

        except Exception as e:
            run_ctx.metadata['error'] = {
                'type': type(e).__name__,
                'message': str(e),
            }
            run_ctx.metadata['status'] = 'error'
            raise

        finally:
            # Queue LLM span event to buffer for API submission
            trace_id = trace_ctx.id if trace_ctx else run_ctx.id
            _queue_llm_span_event_sync(run_ctx, trace_id)

            set_current_span_context(prev_span_ctx)

    async def _wrap_stream_async(self, stream: AsyncIterator, run_ctx: RunContext, interception_ctx: Any = None) -> AsyncIterator:
        """Wrap async stream to collect chunks and calculate streaming metrics."""
        import time
        chunks = []
        first_chunk_time = None
        last_chunk_time = None
        total_tokens = 0
        trace_ctx = get_current_trace_context()

        try:
            async for chunk in stream:
                if first_chunk_time is None:
                    first_chunk_time = time.time()
                last_chunk_time = time.time()

                chunks.append(chunk)
                yield chunk

                # Count tokens in chunk if available
                if hasattr(chunk, 'usage') and chunk.usage:
                    if hasattr(chunk.usage, 'total_tokens'):
                        total_tokens = chunk.usage.total_tokens
                    elif isinstance(chunk.usage, dict):
                        total_tokens = chunk.usage.get('total_tokens', 0)

            # Aggregate content
            full_content = ""
            for chunk in chunks:
                if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
                    delta = chunk.choices[0].delta
                    if hasattr(delta, 'content') and delta.content:
                        full_content += delta.content

            # Calculate streaming metrics
            streaming_metrics = {}
            if first_chunk_time and last_chunk_time:
                time_to_first_token_ms = (first_chunk_time - run_ctx.start_time.timestamp()) * 1000
                streaming_duration_ms = (last_chunk_time - first_chunk_time) * 1000
                streaming_metrics['time_to_first_token_ms'] = time_to_first_token_ms
                streaming_metrics['streaming_duration_ms'] = streaming_duration_ms
                streaming_metrics['chunk_count'] = len(chunks)

                # Calculate tokens per second if we have tokens and duration
                if total_tokens > 0 and streaming_duration_ms > 0:
                    streaming_metrics['tokens_per_second'] = (total_tokens / streaming_duration_ms) * 1000
                elif streaming_duration_ms > 0:
                    # Estimate tokens from content length (rough estimate: 1 token ≈ 4 chars)
                    estimated_tokens = len(full_content) / 4
                    streaming_metrics['estimated_tokens_per_second'] = (estimated_tokens / streaming_duration_ms) * 1000

            # Update token usage from final chunk if available
            if total_tokens > 0:
                run_ctx.metadata['usage'] = {
                    "total_tokens": total_tokens,
                }

            run_ctx.metadata['output'] = {"content": full_content}
            run_ctx.metadata['stream_chunks'] = len(chunks)
            run_ctx.metadata['streaming'] = True
            if streaming_metrics:
                run_ctx.metadata['streaming_metrics'] = streaming_metrics
            run_ctx.metadata['status'] = 'success'

            # Post-call interception for streaming
            if interception_ctx:
                try:
                    from .client import get_aigie
                    aigie = get_aigie()
                    if aigie and aigie._interceptor_chain:
                        interception_ctx.response_content = full_content
                        if total_tokens > 0:
                            interception_ctx.actual_output_tokens = total_tokens
                        await aigie.intercept_post_call(
                            ctx=interception_ctx,
                            response=None,  # Streaming doesn't have a single response object
                            error=None,
                        )
                except Exception as intercept_error:
                    logger.debug(f"[wrapper] Stream post-call interception failed: {intercept_error}")

        except Exception as e:
            run_ctx.metadata['error'] = {'type': type(e).__name__, 'message': str(e)}
            run_ctx.metadata['status'] = 'error'
            raise

        finally:
            # Queue LLM span event after streaming completes
            trace_id = trace_ctx.id if trace_ctx else run_ctx.id
            await _queue_llm_span_event(run_ctx, trace_id)

    def _wrap_stream_sync(self, stream: Iterator, run_ctx: RunContext, interception_ctx: Any = None) -> Iterator:
        """Wrap sync stream to collect chunks and calculate streaming metrics."""
        import time
        chunks = []
        first_chunk_time = None
        last_chunk_time = None
        total_tokens = 0
        trace_ctx = get_current_trace_context()

        try:
            for chunk in stream:
                if first_chunk_time is None:
                    first_chunk_time = time.time()
                last_chunk_time = time.time()

                chunks.append(chunk)
                yield chunk

                # Count tokens in chunk if available
                if hasattr(chunk, 'usage') and chunk.usage:
                    if hasattr(chunk.usage, 'total_tokens'):
                        total_tokens = chunk.usage.total_tokens
                    elif isinstance(chunk.usage, dict):
                        total_tokens = chunk.usage.get('total_tokens', 0)

            full_content = ""
            for chunk in chunks:
                if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
                    delta = chunk.choices[0].delta
                    if hasattr(delta, 'content') and delta.content:
                        full_content += delta.content

            # Calculate streaming metrics (same as async version)
            streaming_metrics = {}
            if first_chunk_time and last_chunk_time:
                time_to_first_token_ms = (first_chunk_time - run_ctx.start_time.timestamp()) * 1000
                streaming_duration_ms = (last_chunk_time - first_chunk_time) * 1000
                streaming_metrics['time_to_first_token_ms'] = time_to_first_token_ms
                streaming_metrics['streaming_duration_ms'] = streaming_duration_ms
                streaming_metrics['chunk_count'] = len(chunks)

                if total_tokens > 0 and streaming_duration_ms > 0:
                    streaming_metrics['tokens_per_second'] = (total_tokens / streaming_duration_ms) * 1000
                elif streaming_duration_ms > 0:
                    estimated_tokens = len(full_content) / 4
                    streaming_metrics['estimated_tokens_per_second'] = (estimated_tokens / streaming_duration_ms) * 1000

            # Update token usage from final chunk if available
            if total_tokens > 0:
                run_ctx.metadata['usage'] = {
                    "total_tokens": total_tokens,
                }

            run_ctx.metadata['output'] = {"content": full_content}
            run_ctx.metadata['stream_chunks'] = len(chunks)
            run_ctx.metadata['streaming'] = True
            if streaming_metrics:
                run_ctx.metadata['streaming_metrics'] = streaming_metrics
            run_ctx.metadata['status'] = 'success'

        except Exception as e:
            run_ctx.metadata['error'] = {'type': type(e).__name__, 'message': str(e)}
            run_ctx.metadata['status'] = 'error'
            raise

        finally:
            # Queue LLM span event after streaming completes
            trace_id = trace_ctx.id if trace_ctx else run_ctx.id
            _queue_llm_span_event_sync(run_ctx, trace_id)

    def _wrap_completions(self, completions_obj: Any) -> Any:
        """Wrap completions (legacy)."""
        # Similar pattern for legacy completions endpoint
        return completions_obj

    def _wrap_embeddings(self, embeddings_obj: Any) -> Any:
        """Wrap embeddings."""
        # Similar pattern for embeddings
        return embeddings_obj


def wrap_openai(client: Any, aigie_client: Optional[Any] = None) -> Any:
    """
    Wrap OpenAI client for automatic tracing.

    Args:
        client: OpenAI client instance
        aigie_client: Optional Aigie client

    Returns:
        Wrapped client with tracing

    Example:
        import openai
        from aigie.wrappers import wrap_openai

        client = wrap_openai(openai.OpenAI(api_key="..."))

        response = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": "Hello"}]
        )
    """
    return OpenAIWrapper(client, aigie_client)


def wrap_anthropic(client: Any, aigie_client: Optional[Any] = None) -> Any:
    """
    Wrap Anthropic client for automatic tracing.

    Args:
        client: Anthropic client instance
        aigie_client: Optional Aigie client

    Returns:
        Wrapped client with tracing

    Example:
        import anthropic
        from aigie.wrappers import wrap_anthropic

        client = wrap_anthropic(anthropic.Anthropic(api_key="..."))

        response = client.messages.create(
            model="claude-3-opus-20240229",
            messages=[{"role": "user", "content": "Hello"}]
        )
    """
    # Implementation similar to OpenAI
    # For now, return client as-is with logging
    logger.warning("Anthropic wrapper not yet fully implemented")
    return client


def wrap_gemini(client: Any, aigie_client: Optional[Any] = None) -> Any:
    """
    Wrap Google Gemini client for automatic tracing.

    Args:
        client: Gemini client instance
        aigie_client: Optional Aigie client

    Returns:
        Wrapped client with tracing

    Example:
        import google.generativeai as genai
        from aigie.wrappers import wrap_gemini

        genai.configure(api_key="...")
        model = wrap_gemini(genai.GenerativeModel('gemini-pro'))

        response = model.generate_content("Hello")
    """
    # Implementation for Gemini
    logger.warning("Gemini wrapper not yet fully implemented")
    return client


__all__ = [
    "wrap_openai",
    "wrap_anthropic",
    "wrap_gemini",
    "OpenAIWrapper",
]
