"""LLM client wrappers — drop-in replacements that log all calls."""

from __future__ import annotations

import logging
import re
import time
from typing import Any

import openai

from llm_observatory.config import get_config
from llm_observatory.event import build_event, _maybe_redact
from llm_observatory.streaming import StreamWrapper, AnthropicStreamWrapper
from llm_observatory.transport import get_transport

logger = logging.getLogger("llm_observatory")

# Maximum length for error messages stored in events
_MAX_ERROR_MESSAGE_LENGTH = 500

# Patterns to strip from error messages
_SENSITIVE_PATTERNS = [
    re.compile(r"(?i)(api[_-]?key|token|secret|password|authorization)[=:]\s*\S+"),
    re.compile(r"(?i)Bearer\s+\S+"),
    re.compile(r"(?i)\?api_key=[^&\s]+"),
    re.compile(r"(?i)sk-[a-zA-Z0-9]{20,}"),
]


def _sanitize_error_message(msg: str) -> str:
    """Strip API keys, Bearer tokens, and other secrets from error messages."""
    sanitized = msg
    for pattern in _SENSITIVE_PATTERNS:
        sanitized = pattern.sub("[REDACTED]", sanitized)
    if len(sanitized) > _MAX_ERROR_MESSAGE_LENGTH:
        sanitized = sanitized[:_MAX_ERROR_MESSAGE_LENGTH] + "..."
    return sanitized


# ── Shared capture logic ─────────────────────────────────────────


class _CompletionsCaptureLogic:
    """Mixin providing capture methods for OpenAI-compatible completions."""

    _parent: Any  # _BaseClient or _AsyncBaseClient

    def _do_capture(
        self,
        request_kwargs: dict[str, Any],
        response: Any,
        start_time: float,
        end_time: float,
        metadata: dict[str, Any] | None = None,
        **extra: Any,
    ) -> None:
        try:
            cfg = get_config()
            event = build_event(
                request_kwargs=request_kwargs,
                response=response,
                start_time=start_time,
                end_time=end_time,
                metadata=metadata,
                redact_messages=cfg.get("redact_messages", False),
                **extra,
            )
            get_transport().enqueue(event)
        except Exception as exc:
            logger.debug("llm_observatory: failed to capture event: %s", exc)

    def _do_capture_error(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        end_time: float,
        exc: Exception,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        try:
            cfg = get_config()
            event = build_event(
                request_kwargs=request_kwargs,
                response=None,
                start_time=start_time,
                end_time=end_time,
                error={
                    "type": type(exc).__name__,
                    "message": _sanitize_error_message(str(exc)),
                },
                metadata=metadata,
                redact_messages=cfg.get("redact_messages", False),
            )
            get_transport().enqueue(event)
        except Exception as capture_exc:
            logger.debug("llm_observatory: failed to capture error event: %s", capture_exc)

    def _resolve_metadata(self, per_request: dict[str, Any] | None = None) -> dict[str, Any] | None:
        """Snapshot parent metadata and merge per-request overrides."""
        base = dict(self._parent._metadata) if self._parent._metadata else {}
        if per_request:
            base.update(per_request)
        return base or None


# ── Sync base client ─────────────────────────────────────────────


class _BaseClient:
    """Shared interception logic for OpenAI-compatible clients."""

    def __init__(self, client: Any) -> None:
        self._client = client
        self.chat = _ChatNamespace(self)
        self.responses = _ResponsesNamespace(self)
        self._metadata: dict[str, Any] = {}

    def set_metadata(self, **kwargs: Any) -> None:
        """Set metadata (user_id, session_id, tags) attached to all subsequent events."""
        self._metadata.update(kwargs)

    def flush(self) -> None:
        """Force-flush all buffered events to the ingestion API."""
        get_transport().flush()

    def __getattr__(self, name: str) -> Any:
        return getattr(self._client, name)


class OpenAI(_BaseClient):
    """Drop-in replacement for openai.OpenAI. All calls are transparently logged."""

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(openai.OpenAI(**kwargs))


class AzureOpenAI(_BaseClient):
    """Drop-in replacement for openai.AzureOpenAI. All calls are transparently logged."""

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(openai.AzureOpenAI(**kwargs))


class _ChatNamespace:
    def __init__(self, parent: _BaseClient) -> None:
        self._parent = parent
        self.completions = _CompletionsNamespace(parent)


class _CompletionsNamespace(_CompletionsCaptureLogic):
    def __init__(self, parent: _BaseClient) -> None:
        self._parent = parent

    def create(self, *args: Any, **kwargs: Any) -> Any:
        cfg = get_config()
        metadata = kwargs.pop("metadata", None)
        if not cfg.get("enabled"):
            return self._parent._client.chat.completions.create(*args, **kwargs)

        resolved_meta = self._resolve_metadata(metadata)
        start_time = time.time()

        # Automatically request usage in the final streaming chunk
        if kwargs.get("stream", False):
            existing = kwargs.get("stream_options") or {}
            kwargs = {**kwargs, "stream_options": {**existing, "include_usage": True}}

        try:
            response = self._parent._client.chat.completions.create(*args, **kwargs)
        except openai.APIError as exc:
            end_time = time.time()
            self._do_capture_error(kwargs, start_time, end_time, exc, resolved_meta)
            raise

        # Handle streaming
        is_stream = kwargs.get("stream", False)
        if is_stream:
            return StreamWrapper(
                stream=response,
                start_time=start_time,
                request_kwargs=kwargs,
                on_complete=lambda **kw: self._capture_stream_complete(kwargs, start_time, resolved_meta, **kw),
                on_error=lambda **kw: self._capture_stream_error(kwargs, start_time, resolved_meta, **kw),
            )

        # Non-streaming: capture immediately
        end_time = time.time()
        self._do_capture(kwargs, response, start_time, end_time, resolved_meta)
        return response

    def _capture_stream_complete(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        metadata: dict[str, Any] | None,
        *,
        response: Any,
        end_time: float,
        time_to_first_token_ms: int | None = None,
        tokens_per_second: float | None = None,
    ) -> None:
        self._do_capture(
            request_kwargs,
            response,
            start_time,
            end_time,
            metadata,
            time_to_first_token_ms=time_to_first_token_ms,
            tokens_per_second=tokens_per_second,
        )

    def _capture_stream_error(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        metadata: dict[str, Any] | None,
        *,
        exc: Exception,
        end_time: float,
    ) -> None:
        self._do_capture_error(request_kwargs, start_time, end_time, exc, metadata)


# ── OpenAI Responses API ─────────────────────────────────────────


class _ResponsesNamespace(_CompletionsCaptureLogic):
    """Intercept client.responses.create() for the OpenAI Responses API."""

    def __init__(self, parent: _BaseClient | _AsyncBaseClient) -> None:
        self._parent = parent

    def create(self, *args: Any, **kwargs: Any) -> Any:
        cfg = get_config()
        metadata = kwargs.pop("metadata", None)

        # Check if the underlying client supports responses
        real_responses = getattr(self._parent._client, "responses", None)
        if real_responses is None:
            raise AttributeError("The underlying OpenAI client does not support the responses API")

        if not cfg.get("enabled"):
            return real_responses.create(*args, **kwargs)

        resolved_meta = self._resolve_metadata(metadata)
        start_time = time.time()

        try:
            response = real_responses.create(*args, **kwargs)
        except openai.APIError as exc:
            end_time = time.time()
            # Build a minimal request_kwargs from the call
            req_kwargs = _responses_request_kwargs(kwargs)
            self._do_capture_error(req_kwargs, start_time, end_time, exc, resolved_meta)
            raise

        end_time = time.time()
        req_kwargs = _responses_request_kwargs(kwargs)
        pseudo_response = _normalize_responses_response(response)
        self._do_capture(req_kwargs, pseudo_response, start_time, end_time, resolved_meta)
        return response


def _responses_request_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
    """Normalize responses.create() kwargs to the standard request schema."""
    req: dict[str, Any] = {
        "model": kwargs.get("model", ""),
        "messages": [],
        "stream": False,
    }
    # The responses API uses 'input' instead of 'messages'
    input_val = kwargs.get("input")
    if isinstance(input_val, str):
        req["messages"] = [{"role": "user", "content": input_val}]
    elif isinstance(input_val, list):
        req["messages"] = input_val
    if "temperature" in kwargs:
        req["temperature"] = kwargs["temperature"]
    if "max_output_tokens" in kwargs:
        req["max_tokens"] = kwargs["max_output_tokens"]
    if "tools" in kwargs:
        req["tools"] = kwargs["tools"]
    return req


class _NormalizedResponsesResponse:
    """Minimal response-like object for event building from Responses API."""

    def __init__(
        self,
        id: str,
        model: str,
        content: str,
        usage: Any,
    ) -> None:
        self.id = id
        self.model = model
        self.choices = [_NormalizedResponsesChoice(content)]
        self.usage = usage


class _NormalizedResponsesChoice:
    def __init__(self, content: str) -> None:
        self.index = 0
        self.message = _NormalizedResponsesMessage(content)
        self.finish_reason = "stop"


class _NormalizedResponsesMessage:
    def __init__(self, content: str) -> None:
        self.role = "assistant"
        self.content = content
        self.reasoning_content = None
        self.tool_calls = None


class _NormalizedResponsesUsage:
    def __init__(self, input_tokens: int, output_tokens: int) -> None:
        self.prompt_tokens = input_tokens
        self.completion_tokens = output_tokens
        self.total_tokens = input_tokens + output_tokens
        self.prompt_tokens_details = None
        self.completion_tokens_details = None


def _normalize_responses_response(response: Any) -> _NormalizedResponsesResponse:
    """Convert a Responses API response to our standard format."""
    resp_id = getattr(response, "id", "")
    model = getattr(response, "model", "")

    # Extract text from output
    content_parts: list[str] = []
    output = getattr(response, "output", None)
    if output:
        for item in output:
            item_type = getattr(item, "type", "")
            if item_type == "message":
                for block in getattr(item, "content", []):
                    if getattr(block, "type", "") == "output_text":
                        content_parts.append(getattr(block, "text", ""))

    content = "".join(content_parts)

    # Extract usage
    usage_obj = getattr(response, "usage", None)
    usage = None
    if usage_obj is not None:
        input_tokens = getattr(usage_obj, "input_tokens", 0) or 0
        output_tokens = getattr(usage_obj, "output_tokens", 0) or 0
        usage = _NormalizedResponsesUsage(input_tokens, output_tokens)

    return _NormalizedResponsesResponse(
        id=resp_id,
        model=model,
        content=content,
        usage=usage,
    )


# ── Async base client ────────────────────────────────────────────


class _AsyncBaseClient:
    """Shared interception logic for async OpenAI-compatible clients."""

    def __init__(self, client: Any) -> None:
        self._client = client
        self.chat = _AsyncChatNamespace(self)
        self.responses = _AsyncResponsesNamespace(self)
        self._metadata: dict[str, Any] = {}

    def set_metadata(self, **kwargs: Any) -> None:
        self._metadata.update(kwargs)

    def flush(self) -> None:
        get_transport().flush()

    def __getattr__(self, name: str) -> Any:
        return getattr(self._client, name)


class AsyncOpenAI(_AsyncBaseClient):
    """Drop-in replacement for openai.AsyncOpenAI. All calls are transparently logged."""

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(openai.AsyncOpenAI(**kwargs))


class AsyncAzureOpenAI(_AsyncBaseClient):
    """Drop-in replacement for openai.AsyncAzureOpenAI. All calls are transparently logged."""

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(openai.AsyncAzureOpenAI(**kwargs))


class _AsyncChatNamespace:
    def __init__(self, parent: _AsyncBaseClient) -> None:
        self._parent = parent
        self.completions = _AsyncCompletionsNamespace(parent)


class _AsyncCompletionsNamespace(_CompletionsCaptureLogic):
    def __init__(self, parent: _AsyncBaseClient) -> None:
        self._parent = parent

    async def create(self, *args: Any, **kwargs: Any) -> Any:
        cfg = get_config()
        metadata = kwargs.pop("metadata", None)
        if not cfg.get("enabled"):
            return await self._parent._client.chat.completions.create(*args, **kwargs)

        resolved_meta = self._resolve_metadata(metadata)
        start_time = time.time()

        if kwargs.get("stream", False):
            existing = kwargs.get("stream_options") or {}
            kwargs = {**kwargs, "stream_options": {**existing, "include_usage": True}}

        try:
            response = await self._parent._client.chat.completions.create(*args, **kwargs)
        except openai.APIError as exc:
            end_time = time.time()
            self._do_capture_error(kwargs, start_time, end_time, exc, resolved_meta)
            raise

        is_stream = kwargs.get("stream", False)
        if is_stream:
            return StreamWrapper(
                stream=response,
                start_time=start_time,
                request_kwargs=kwargs,
                on_complete=lambda **kw: self._capture_stream_complete(kwargs, start_time, resolved_meta, **kw),
                on_error=lambda **kw: self._capture_stream_error(kwargs, start_time, resolved_meta, **kw),
            )

        end_time = time.time()
        self._do_capture(kwargs, response, start_time, end_time, resolved_meta)
        return response

    def _capture_stream_complete(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        metadata: dict[str, Any] | None,
        *,
        response: Any,
        end_time: float,
        time_to_first_token_ms: int | None = None,
        tokens_per_second: float | None = None,
    ) -> None:
        self._do_capture(
            request_kwargs,
            response,
            start_time,
            end_time,
            metadata,
            time_to_first_token_ms=time_to_first_token_ms,
            tokens_per_second=tokens_per_second,
        )

    def _capture_stream_error(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        metadata: dict[str, Any] | None,
        *,
        exc: Exception,
        end_time: float,
    ) -> None:
        self._do_capture_error(request_kwargs, start_time, end_time, exc, metadata)


class _AsyncResponsesNamespace(_CompletionsCaptureLogic):
    """Async intercept for client.responses.create() (Responses API)."""

    def __init__(self, parent: _AsyncBaseClient) -> None:
        self._parent = parent

    async def create(self, *args: Any, **kwargs: Any) -> Any:
        cfg = get_config()
        metadata = kwargs.pop("metadata", None)

        real_responses = getattr(self._parent._client, "responses", None)
        if real_responses is None:
            raise AttributeError("The underlying OpenAI client does not support the responses API")

        if not cfg.get("enabled"):
            return await real_responses.create(*args, **kwargs)

        resolved_meta = self._resolve_metadata(metadata)
        start_time = time.time()

        try:
            response = await real_responses.create(*args, **kwargs)
        except openai.APIError as exc:
            end_time = time.time()
            req_kwargs = _responses_request_kwargs(kwargs)
            self._do_capture_error(req_kwargs, start_time, end_time, exc, resolved_meta)
            raise

        end_time = time.time()
        req_kwargs = _responses_request_kwargs(kwargs)
        pseudo_response = _normalize_responses_response(response)
        self._do_capture(req_kwargs, pseudo_response, start_time, end_time, resolved_meta)
        return response


# ── Anthropic wrapper ─────────────────────────────────────────────


class Anthropic:
    """Drop-in replacement for anthropic.Anthropic. All calls are transparently logged."""

    def __init__(self, **kwargs: Any) -> None:
        import anthropic as _anthropic

        self._client = _anthropic.Anthropic(**kwargs)
        self.messages = _AnthropicMessagesNamespace(self)
        self._metadata: dict[str, Any] = {}

    def set_metadata(self, **kwargs: Any) -> None:
        self._metadata.update(kwargs)

    def flush(self) -> None:
        get_transport().flush()

    def __getattr__(self, name: str) -> Any:
        return getattr(self._client, name)


class _AnthropicMessagesNamespace:
    def __init__(self, parent: Anthropic | AsyncAnthropic) -> None:
        self._parent = parent

    def create(self, *args: Any, **kwargs: Any) -> Any:
        cfg = get_config()
        metadata = kwargs.pop("metadata", None)
        if not cfg.get("enabled"):
            return self._parent._client.messages.create(*args, **kwargs)

        resolved_meta = self._resolve_metadata(metadata)
        start_time = time.time()
        is_stream = kwargs.get("stream", False)

        try:
            response = self._parent._client.messages.create(*args, **kwargs)
        except Exception as exc:
            end_time = time.time()
            self._capture_error(kwargs, start_time, end_time, exc, resolved_meta)
            raise

        if is_stream:
            return AnthropicStreamWrapper(
                stream=response,
                start_time=start_time,
                request_kwargs=kwargs,
                on_complete=lambda **kw: self._capture_stream_complete(kwargs, start_time, resolved_meta, **kw),
                on_error=lambda **kw: self._capture_stream_error(kwargs, start_time, resolved_meta, **kw),
            )

        end_time = time.time()
        self._capture(kwargs, response, start_time, end_time, resolved_meta)
        return response

    def _resolve_metadata(self, per_request: dict[str, Any] | None = None) -> dict[str, Any] | None:
        base = dict(self._parent._metadata) if self._parent._metadata else {}
        if per_request:
            base.update(per_request)
        return base or None

    def _capture(
        self,
        request_kwargs: dict[str, Any],
        response: Any,
        start_time: float,
        end_time: float,
        metadata: dict[str, Any] | None = None,
        **extra: Any,
    ) -> None:
        try:
            cfg = get_config()
            event = _build_anthropic_event(
                request_kwargs=request_kwargs,
                response=response,
                start_time=start_time,
                end_time=end_time,
                metadata=metadata,
                redact_messages=cfg.get("redact_messages", False),
                **extra,
            )
            get_transport().enqueue(event)
        except Exception as exc:
            logger.debug("llm_observatory: failed to capture anthropic event: %s", exc)

    def _capture_stream_complete(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        metadata: dict[str, Any] | None,
        *,
        response: Any,
        end_time: float,
        time_to_first_token_ms: int | None = None,
        tokens_per_second: float | None = None,
    ) -> None:
        self._capture(
            request_kwargs,
            response,
            start_time,
            end_time,
            metadata,
            time_to_first_token_ms=time_to_first_token_ms,
            tokens_per_second=tokens_per_second,
        )

    def _capture_stream_error(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        metadata: dict[str, Any] | None,
        *,
        exc: Exception,
        end_time: float,
    ) -> None:
        self._capture_error(request_kwargs, start_time, end_time, exc, metadata)

    def _capture_error(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        end_time: float,
        exc: Exception,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        try:
            cfg = get_config()
            event = _build_anthropic_event(
                request_kwargs=request_kwargs,
                response=None,
                start_time=start_time,
                end_time=end_time,
                error={
                    "type": type(exc).__name__,
                    "message": _sanitize_error_message(str(exc)),
                },
                metadata=metadata,
                redact_messages=cfg.get("redact_messages", False),
            )
            get_transport().enqueue(event)
        except Exception as capture_exc:
            logger.debug("llm_observatory: failed to capture anthropic error event: %s", capture_exc)


class AsyncAnthropic:
    """Drop-in replacement for anthropic.AsyncAnthropic. All calls are transparently logged."""

    def __init__(self, **kwargs: Any) -> None:
        import anthropic as _anthropic

        self._client = _anthropic.AsyncAnthropic(**kwargs)
        self.messages = _AsyncAnthropicMessagesNamespace(self)
        self._metadata: dict[str, Any] = {}

    def set_metadata(self, **kwargs: Any) -> None:
        self._metadata.update(kwargs)

    def flush(self) -> None:
        get_transport().flush()

    def __getattr__(self, name: str) -> Any:
        return getattr(self._client, name)


class _AsyncAnthropicMessagesNamespace:
    def __init__(self, parent: AsyncAnthropic) -> None:
        self._parent = parent

    async def create(self, *args: Any, **kwargs: Any) -> Any:
        cfg = get_config()
        metadata = kwargs.pop("metadata", None)
        if not cfg.get("enabled"):
            return await self._parent._client.messages.create(*args, **kwargs)

        resolved_meta = self._resolve_metadata(metadata)
        start_time = time.time()
        is_stream = kwargs.get("stream", False)

        try:
            response = await self._parent._client.messages.create(*args, **kwargs)
        except Exception as exc:
            end_time = time.time()
            self._capture_error(kwargs, start_time, end_time, exc, resolved_meta)
            raise

        if is_stream:
            return AnthropicStreamWrapper(
                stream=response,
                start_time=start_time,
                request_kwargs=kwargs,
                on_complete=lambda **kw: self._capture_stream_complete(kwargs, start_time, resolved_meta, **kw),
                on_error=lambda **kw: self._capture_stream_error(kwargs, start_time, resolved_meta, **kw),
            )

        end_time = time.time()
        self._capture(kwargs, response, start_time, end_time, resolved_meta)
        return response

    def _resolve_metadata(self, per_request: dict[str, Any] | None = None) -> dict[str, Any] | None:
        base = dict(self._parent._metadata) if self._parent._metadata else {}
        if per_request:
            base.update(per_request)
        return base or None

    def _capture(
        self,
        request_kwargs: dict[str, Any],
        response: Any,
        start_time: float,
        end_time: float,
        metadata: dict[str, Any] | None = None,
        **extra: Any,
    ) -> None:
        try:
            cfg = get_config()
            event = _build_anthropic_event(
                request_kwargs=request_kwargs,
                response=response,
                start_time=start_time,
                end_time=end_time,
                metadata=metadata,
                redact_messages=cfg.get("redact_messages", False),
                **extra,
            )
            get_transport().enqueue(event)
        except Exception as exc:
            logger.debug("llm_observatory: failed to capture anthropic event: %s", exc)

    def _capture_stream_complete(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        metadata: dict[str, Any] | None,
        *,
        response: Any,
        end_time: float,
        time_to_first_token_ms: int | None = None,
        tokens_per_second: float | None = None,
    ) -> None:
        self._capture(
            request_kwargs,
            response,
            start_time,
            end_time,
            metadata,
            time_to_first_token_ms=time_to_first_token_ms,
            tokens_per_second=tokens_per_second,
        )

    def _capture_stream_error(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        metadata: dict[str, Any] | None,
        *,
        exc: Exception,
        end_time: float,
    ) -> None:
        self._capture_error(request_kwargs, start_time, end_time, exc, metadata)

    def _capture_error(
        self,
        request_kwargs: dict[str, Any],
        start_time: float,
        end_time: float,
        exc: Exception,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        try:
            cfg = get_config()
            event = _build_anthropic_event(
                request_kwargs=request_kwargs,
                response=None,
                start_time=start_time,
                end_time=end_time,
                error={
                    "type": type(exc).__name__,
                    "message": _sanitize_error_message(str(exc)),
                },
                metadata=metadata,
                redact_messages=cfg.get("redact_messages", False),
            )
            get_transport().enqueue(event)
        except Exception as capture_exc:
            logger.debug("llm_observatory: failed to capture anthropic error event: %s", capture_exc)


def _build_anthropic_event(
    *,
    request_kwargs: dict[str, Any],
    response: Any,
    start_time: float,
    end_time: float,
    error: dict[str, str] | None = None,
    metadata: dict[str, Any] | None = None,
    redact_messages: bool = False,
    time_to_first_token_ms: int | None = None,
    tokens_per_second: float | None = None,
) -> dict[str, Any]:
    """Build a normalized event from an Anthropic messages.create call."""
    import uuid
    from datetime import datetime, timezone

    latency_ms = int((end_time - start_time) * 1000)

    # Normalize request to OpenAI-like format
    messages = request_kwargs.get("messages", [])
    if redact_messages and messages:
        messages = _maybe_redact(messages, True)

    request: dict[str, Any] = {
        "model": request_kwargs.get("model", ""),
        "messages": messages,
        "stream": request_kwargs.get("stream", False),
    }

    # Anthropic uses system as a top-level param, not a message
    system_prompt = request_kwargs.get("system")
    if system_prompt:
        if redact_messages:
            system_prompt = "[REDACTED]"
        request["messages"] = [{"role": "system", "content": system_prompt}] + list(request["messages"])

    if "temperature" in request_kwargs:
        request["temperature"] = request_kwargs["temperature"]
    if "max_tokens" in request_kwargs:
        request["max_tokens"] = request_kwargs["max_tokens"]
    if "tools" in request_kwargs:
        request["tools"] = request_kwargs["tools"]

    # Build response + usage
    resp_section: dict[str, Any] = {}
    usage_section: dict[str, int] = {
        "prompt_tokens": 0,
        "completion_tokens": 0,
        "total_tokens": 0,
    }

    if error is not None:
        pass
    elif response is not None:
        # Anthropic response: id, model, content, stop_reason, usage
        content_text = ""
        thinking_text = ""
        tool_use_blocks: list[dict[str, Any]] = []
        content_blocks = getattr(response, "content", [])
        if content_blocks:
            text_parts = []
            thinking_parts = []
            for block in content_blocks:
                block_type = getattr(block, "type", "")
                if block_type == "text":
                    text_parts.append(getattr(block, "text", ""))
                elif block_type == "thinking":
                    thinking_parts.append(getattr(block, "thinking", ""))
                elif block_type == "tool_use":
                    tool_use_blocks.append({
                        "id": getattr(block, "id", ""),
                        "type": "function",
                        "function": {
                            "name": getattr(block, "name", ""),
                            "arguments": str(getattr(block, "input", {})),
                        },
                    })
            content_text = "".join(text_parts)
            thinking_text = "".join(thinking_parts)

        if redact_messages and content_text:
            content_text = "[REDACTED]"
        if redact_messages and thinking_text:
            thinking_text = "[REDACTED]"

        message: dict[str, Any] = {"role": "assistant", "content": content_text}
        if thinking_text:
            message["thinking"] = thinking_text
        if tool_use_blocks:
            if redact_messages:
                for tub in tool_use_blocks:
                    tub["function"]["arguments"] = "[REDACTED]"
            message["tool_calls"] = tool_use_blocks

        resp_section = {
            "id": getattr(response, "id", ""),
            "model": getattr(response, "model", ""),
            "finish_reason": getattr(response, "stop_reason", "") or "",
            "choices": [{
                "index": 0,
                "message": message,
                "finish_reason": getattr(response, "stop_reason", "") or "",
            }],
        }

        usage = getattr(response, "usage", None)
        if usage is not None:
            input_tokens = getattr(usage, "input_tokens", 0) or 0
            output_tokens = getattr(usage, "output_tokens", 0) or 0
            usage_section: dict[str, Any] = {
                "prompt_tokens": input_tokens,
                "completion_tokens": output_tokens,
                "total_tokens": input_tokens + output_tokens,
            }
            cache_read = getattr(usage, "cache_read_input_tokens", None)
            if cache_read:
                usage_section["cached_tokens"] = cache_read
            cache_creation = getattr(usage, "cache_creation_input_tokens", None)
            if cache_creation:
                usage_section["cache_creation_input_tokens"] = cache_creation

    # Build metrics
    metrics: dict[str, Any] = {"latency_ms": latency_ms}
    if time_to_first_token_ms is not None:
        metrics["time_to_first_token_ms"] = time_to_first_token_ms
    if tokens_per_second is not None:
        metrics["tokens_per_second"] = round(tokens_per_second, 1)

    # Build metadata
    meta = metadata or {}
    metadata_section: dict[str, Any] = {}
    if meta.get("user_id"):
        metadata_section["user_id"] = meta["user_id"]
    if meta.get("session_id"):
        metadata_section["session_id"] = meta["session_id"]
    if meta.get("tags"):
        metadata_section["tags"] = meta["tags"]
    if meta.get("custom"):
        metadata_section["custom"] = {k: str(v) for k, v in meta["custom"].items()}

    event: dict[str, Any] = {
        "event_id": str(uuid.uuid4()),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "request": request,
        "response": resp_section,
        "usage": usage_section,
        "metrics": metrics,
        "metadata": metadata_section,
        "span_type": "generation",
    }
    if error is not None:
        event["error"] = error

    # Pick up trace context from active span
    try:
        from llm_observatory.tracing import get_current_trace_id, get_current_span_id

        trace_id = get_current_trace_id()
        parent_id = get_current_span_id()
        if trace_id:
            event["trace_id"] = trace_id
        if parent_id:
            event["parent_id"] = parent_id
    except Exception as exc:
        logger.debug("llm_observatory: failed to pick up trace context: %s", exc)

    return event


# ── wrap() — universal client wrapper ─────────────────────────────


def _is_async_client(client: Any) -> bool:
    """Detect whether *client* is an async LLM client."""
    try:
        from openai._base_client import AsyncAPIClient as _OAAsync
        if isinstance(client, _OAAsync):
            return True
    except Exception:
        pass
    try:
        from anthropic._base_client import AsyncAPIClient as _AAsync
        if isinstance(client, _AAsync):
            return True
    except Exception:
        pass
    return "Async" in type(client).__name__


def _is_anthropic_client(client: Any) -> bool:
    """True if *client* looks like an Anthropic client (has .messages, no .chat)."""
    return hasattr(client, "messages") and not hasattr(client, "chat")


def wrap(client: Any) -> Any:
    """Wrap any supported LLM client with Context0 observability.

    Works with OpenAI, Azure OpenAI, Anthropic, and third-party wrappers
    (e.g. Langfuse's ``OpenAI``).  If the client is already a Context0
    wrapper, it is returned as-is (no double-wrapping).

    Returns a ``_BaseClient``, ``_AsyncBaseClient``, ``Anthropic``, or
    ``AsyncAnthropic`` instance depending on the input.

    Raises ``TypeError`` for unsupported client types.
    """
    # Already wrapped — return as-is
    if isinstance(client, (_BaseClient, _AsyncBaseClient, Anthropic, AsyncAnthropic)):
        return client

    is_async = _is_async_client(client)

    # Anthropic-shaped client
    if _is_anthropic_client(client):
        if is_async:
            wrapper = object.__new__(AsyncAnthropic)
            wrapper._client = client
            wrapper.messages = _AsyncAnthropicMessagesNamespace(wrapper)
            wrapper._metadata = {}
            return wrapper
        wrapper = object.__new__(Anthropic)
        wrapper._client = client
        wrapper.messages = _AnthropicMessagesNamespace(wrapper)
        wrapper._metadata = {}
        return wrapper

    # OpenAI-shaped client (has .chat)
    if hasattr(client, "chat"):
        if is_async:
            return _AsyncBaseClient(client)
        return _BaseClient(client)

    raise TypeError(
        f"wrap() does not support {type(client).__name__}. "
        "Expected an OpenAI, Azure OpenAI, or Anthropic client."
    )
