"""wrap() — transparent LLM client wrapping for automatic trace capture."""

from __future__ import annotations

import asyncio
import functools
import json
import logging
import time
from datetime import datetime, timezone
from typing import Any, TypeVar

from carrot_ai._context import current_trace_id

logger = logging.getLogger("carrot_ai")

T = TypeVar("T")

# ---------------------------------------------------------------------------
# Resource paths that we intercept per provider
# ---------------------------------------------------------------------------
_OPENAI_RESOURCES: list[tuple[list[str], str]] = [
    (["chat", "completions"], "chat.completions"),
    (["completions"], "completions"),
    (["embeddings"], "embeddings"),
    (["responses"], "responses"),
]
_ANTHROPIC_RESOURCES: list[tuple[list[str], str]] = [
    (["messages"], "messages"),
]

# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def wrap(client: T, *, name: str | None = None, tags: list[str] | None = None) -> T:
    """Wrap an OpenAI or Anthropic client for automatic trace capture.

    The original client is modified in-place (its ``.create()`` methods are
    wrapped) and returned.  Type checkers still see the original type.

    Args:
        client: An OpenAI or Anthropic client instance.
        name: Override the default trace name (derived from the API resource,
            e.g. ``"chat.completions"``).
        tags: Tags attached to every trace.

    Usage::

        import carrot_ai
        from openai import OpenAI

        carrot_ai.init(api_key="ck-...")
        client = carrot_ai.wrap(OpenAI(), name="my-service")

        # use client as normal — traces are captured automatically
        client.chat.completions.create(model="gpt-4o", messages=[...])
    """
    provider = _detect_provider(client)
    if provider is None:
        raise ValueError(
            f"Unsupported client type: {type(client).__name__}. "
            "carrot_ai.wrap() supports OpenAI and Anthropic clients."
        )

    is_async = "Async" in type(client).__name__
    extra: dict[str, Any] = {}
    if name:
        extra["name"] = name
    if tags:
        extra["tags"] = tags

    resources = _OPENAI_RESOURCES if provider == "openai" else _ANTHROPIC_RESOURCES
    for attr_path, resource_name in resources:
        _patch_resource(client, attr_path, provider, resource_name, is_async, extra)

    return client


# ---------------------------------------------------------------------------
# Provider detection
# ---------------------------------------------------------------------------


def _detect_provider(client: Any) -> str | None:
    module = getattr(type(client), "__module__", "") or ""
    if module.startswith("openai"):
        return "openai"
    if module.startswith("anthropic"):
        return "anthropic"
    return None


# ---------------------------------------------------------------------------
# Monkey-patching
# ---------------------------------------------------------------------------


def _patch_resource(
    client: Any,
    attr_path: list[str],
    provider: str,
    resource_name: str,
    is_async: bool,
    extra: dict[str, Any],
) -> None:
    """Navigate to the resource object and replace its ``create`` method."""
    obj = client
    for attr in attr_path:
        obj = getattr(obj, attr, None)
        if obj is None:
            return

    original_create = getattr(obj, "create", None)
    if original_create is None:
        return

    wrapped = _make_create_wrapper(original_create, provider, resource_name, is_async, extra)
    try:
        obj.create = wrapped
    except (AttributeError, TypeError):
        logger.debug("carrot_ai: could not patch %s.create", resource_name)


def _make_create_wrapper(
    original_fn: Any,
    provider: str,
    resource_name: str,
    is_async: bool,
    extra: dict[str, Any],
) -> Any:
    if is_async:

        @functools.wraps(original_fn)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            return await _trace_async_call(
                original_fn, provider, resource_name, extra, args, kwargs
            )

        return async_wrapper

    @functools.wraps(original_fn)
    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
        return _trace_sync_call(
            original_fn, provider, resource_name, extra, args, kwargs
        )

    return sync_wrapper


# ---------------------------------------------------------------------------
# Sync call tracing
# ---------------------------------------------------------------------------


def _trace_sync_call(
    fn: Any,
    provider: str,
    resource_name: str,
    extra: dict[str, Any],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> Any:
    started_at = time.time()
    is_stream = kwargs.get("stream", False)
    trace_input = _capture_input(kwargs)

    if is_stream and provider == "openai" and resource_name in ("chat.completions", "completions"):
        kwargs = dict(kwargs)
        opts = dict(kwargs.get("stream_options") or {})
        opts["include_usage"] = True
        kwargs["stream_options"] = opts

    try:
        result = fn(*args, **kwargs)
    except Exception as exc:
        _submit_error(provider, resource_name, extra, trace_input, exc, started_at)
        raise

    if is_stream:
        accumulator = _make_accumulator(provider, resource_name)

        def on_done() -> None:
            _submit_success(
                provider, resource_name, extra, trace_input,
                accumulator.result(), started_at, time.time(),
            )

        return _SyncStreamProxy(result, accumulator, on_done)

    output = _normalize_response(provider, result, resource_name)
    _submit_success(
        provider, resource_name, extra, trace_input, output, started_at, time.time()
    )
    return result


# ---------------------------------------------------------------------------
# Async call tracing
# ---------------------------------------------------------------------------


async def _trace_async_call(
    fn: Any,
    provider: str,
    resource_name: str,
    extra: dict[str, Any],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> Any:
    started_at = time.time()
    is_stream = kwargs.get("stream", False)
    trace_input = _capture_input(kwargs)

    if is_stream and provider == "openai" and resource_name in ("chat.completions", "completions"):
        kwargs = dict(kwargs)
        opts = dict(kwargs.get("stream_options") or {})
        opts["include_usage"] = True
        kwargs["stream_options"] = opts

    try:
        result = await fn(*args, **kwargs)
    except Exception as exc:
        _submit_error(provider, resource_name, extra, trace_input, exc, started_at)
        raise

    if is_stream:
        accumulator = _make_accumulator(provider, resource_name)

        def on_done() -> None:
            _submit_success(
                provider, resource_name, extra, trace_input,
                accumulator.result(), started_at, time.time(),
            )

        return _AsyncStreamProxy(result, accumulator, on_done)

    output = _normalize_response(provider, result, resource_name)
    _submit_success(
        provider, resource_name, extra, trace_input, output, started_at, time.time()
    )
    return result


# ---------------------------------------------------------------------------
# Stream proxies
# ---------------------------------------------------------------------------


class _SyncStreamProxy:
    """Wraps a sync streaming response, accumulating data for the trace."""

    def __init__(self, stream: Any, accumulator: Any, on_done: Any) -> None:
        self._stream = stream
        self._accumulator = accumulator
        self._on_done = on_done
        self._done = False

    def __iter__(self) -> Any:
        try:
            for chunk in self._stream:
                self._accumulator(chunk)
                yield chunk
        finally:
            self._finish()

    def __enter__(self) -> _SyncStreamProxy:
        if hasattr(self._stream, "__enter__"):
            self._stream.__enter__()
        return self

    def __exit__(self, *args: Any) -> Any:
        self._finish()
        if hasattr(self._stream, "__exit__"):
            return self._stream.__exit__(*args)
        return None

    def close(self) -> None:
        self._finish()
        if hasattr(self._stream, "close"):
            self._stream.close()

    def _finish(self) -> None:
        if not self._done:
            self._done = True
            try:
                self._on_done()
            except Exception:
                logger.debug("carrot_ai: error in stream trace callback", exc_info=True)

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


class _AsyncStreamProxy:
    """Wraps an async streaming response, accumulating data for the trace."""

    def __init__(self, stream: Any, accumulator: Any, on_done: Any) -> None:
        self._stream = stream
        self._accumulator = accumulator
        self._on_done = on_done
        self._done = False

    async def __aiter__(self) -> Any:
        try:
            async for chunk in self._stream:
                self._accumulator(chunk)
                yield chunk
        finally:
            self._finish()

    async def __aenter__(self) -> _AsyncStreamProxy:
        if hasattr(self._stream, "__aenter__"):
            await self._stream.__aenter__()
        return self

    async def __aexit__(self, *args: Any) -> Any:
        self._finish()
        if hasattr(self._stream, "__aexit__"):
            return await self._stream.__aexit__(*args)
        return None

    def _finish(self) -> None:
        if not self._done:
            self._done = True
            try:
                self._on_done()
            except Exception:
                logger.debug("carrot_ai: error in stream trace callback", exc_info=True)

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


# ---------------------------------------------------------------------------
# Stream accumulators
# ---------------------------------------------------------------------------


class _OpenAIAccumulator:
    def __init__(self) -> None:
        self.content_parts: list[str] = []
        self.tool_calls: dict[int, dict[str, Any]] = {}
        self.model: str | None = None
        self.usage: dict[str, int] | None = None

    def __call__(self, chunk: Any) -> None:
        if getattr(chunk, "model", None):
            self.model = chunk.model
        usage = getattr(chunk, "usage", None)
        if usage:
            self.usage = {
                "input_tokens": getattr(usage, "prompt_tokens", 0) or 0,
                "output_tokens": getattr(usage, "completion_tokens", 0) or 0,
                "total_tokens": getattr(usage, "total_tokens", 0) or 0,
            }
        choices = getattr(chunk, "choices", None)
        if not choices:
            return
        delta = getattr(choices[0], "delta", None)
        if delta is None:
            return
        content = getattr(delta, "content", None)
        if content:
            self.content_parts.append(content)
        tc_list = getattr(delta, "tool_calls", None)
        if tc_list:
            for tc in tc_list:
                idx = tc.index
                if idx not in self.tool_calls:
                    self.tool_calls[idx] = {
                        "id": getattr(tc, "id", "") or "",
                        "type": "function",
                        "function": {"name": "", "arguments": ""},
                    }
                if tc.id:
                    self.tool_calls[idx]["id"] = tc.id
                fn = getattr(tc, "function", None)
                if fn:
                    if fn.name:
                        self.tool_calls[idx]["function"]["name"] = fn.name
                    if fn.arguments:
                        self.tool_calls[idx]["function"]["arguments"] += fn.arguments

    def result(self) -> dict[str, Any]:
        message: dict[str, Any] = {"role": "assistant"}
        content = "".join(self.content_parts)
        if content:
            message["content"] = content
        if self.tool_calls:
            message["tool_calls"] = [
                self.tool_calls[i] for i in sorted(self.tool_calls)
            ]
        out: dict[str, Any] = {"message": message}
        if self.model:
            out["model"] = self.model
        if self.usage:
            out["usage"] = self.usage
        return out


class _AnthropicAccumulator:
    def __init__(self) -> None:
        self.content_parts: list[str] = []
        self.tool_calls: list[dict[str, Any]] = []
        self.model: str | None = None
        self.input_tokens: int = 0
        self.output_tokens: int = 0
        self._current_tool: dict[str, Any] | None = None

    def __call__(self, event: Any) -> None:
        event_type = getattr(event, "type", "")
        if event_type == "message_start":
            msg = getattr(event, "message", None)
            if msg:
                self.model = getattr(msg, "model", None)
                usage = getattr(msg, "usage", None)
                if usage:
                    self.input_tokens = getattr(usage, "input_tokens", 0) or 0
        elif event_type == "content_block_start":
            block = getattr(event, "content_block", None)
            if block and getattr(block, "type", "") == "tool_use":
                self._current_tool = {
                    "id": getattr(block, "id", ""),
                    "type": "function",
                    "function": {
                        "name": getattr(block, "name", ""),
                        "arguments": "",
                    },
                }
                self.tool_calls.append(self._current_tool)
        elif event_type == "content_block_delta":
            delta = getattr(event, "delta", None)
            if delta:
                text = getattr(delta, "text", None)
                if text:
                    self.content_parts.append(text)
                partial = getattr(delta, "partial_json", None)
                if partial and self._current_tool:
                    self._current_tool["function"]["arguments"] += partial
        elif event_type == "content_block_stop":
            self._current_tool = None
        elif event_type == "message_delta":
            usage = getattr(event, "usage", None)
            if usage:
                self.output_tokens = getattr(usage, "output_tokens", 0) or 0

    def result(self) -> dict[str, Any]:
        message: dict[str, Any] = {"role": "assistant"}
        content = "".join(self.content_parts)
        if content:
            message["content"] = content
        if self.tool_calls:
            message["tool_calls"] = self.tool_calls
        total = self.input_tokens + self.output_tokens
        out: dict[str, Any] = {
            "message": message,
            "usage": {
                "input_tokens": self.input_tokens,
                "output_tokens": self.output_tokens,
                "total_tokens": total,
            },
        }
        if self.model:
            out["model"] = self.model
        return out


class _OpenAIResponsesAccumulator:
    """Accumulate streaming events from the OpenAI Responses API."""

    def __init__(self) -> None:
        self.content_parts: list[str] = []
        self.tool_calls: list[dict[str, Any]] = []
        self.model: str | None = None
        self.usage: dict[str, int] | None = None

    def __call__(self, event: Any) -> None:
        event_type = getattr(event, "type", "")

        if event_type == "response.output_text.delta":
            delta = getattr(event, "delta", "")
            if delta:
                self.content_parts.append(delta)

        elif event_type == "response.output_item.done":
            item = getattr(event, "item", None)
            if item and getattr(item, "type", "") == "function_call":
                self.tool_calls.append({
                    "id": getattr(item, "call_id", ""),
                    "type": "function",
                    "function": {
                        "name": getattr(item, "name", ""),
                        "arguments": getattr(item, "arguments", ""),
                    },
                })

        elif event_type == "response.completed":
            resp = getattr(event, "response", None)
            if resp:
                self.model = getattr(resp, "model", None)
                usage = getattr(resp, "usage", None)
                if usage:
                    input_tokens = getattr(usage, "input_tokens", 0) or 0
                    output_tokens = getattr(usage, "output_tokens", 0) or 0
                    self.usage = {
                        "input_tokens": input_tokens,
                        "output_tokens": output_tokens,
                        "total_tokens": input_tokens + output_tokens,
                    }

    def result(self) -> dict[str, Any]:
        message: dict[str, Any] = {"role": "assistant"}
        content = "".join(self.content_parts)
        if content:
            message["content"] = content
        if self.tool_calls:
            message["tool_calls"] = self.tool_calls
        out: dict[str, Any] = {"message": message}
        if self.model:
            out["model"] = self.model
        if self.usage:
            out["usage"] = self.usage
        return out


def _make_accumulator(provider: str, resource_name: str = "") -> Any:
    if provider == "openai":
        if resource_name == "responses":
            return _OpenAIResponsesAccumulator()
        return _OpenAIAccumulator()
    return _AnthropicAccumulator()


# ---------------------------------------------------------------------------
# Response normalization (non-streaming)
# ---------------------------------------------------------------------------


def _normalize_response(provider: str, response: Any, resource_name: str = "") -> dict[str, Any]:
    if provider == "openai":
        if resource_name == "responses":
            return _normalize_openai_responses(response)
        return _normalize_openai(response)
    return _normalize_anthropic(response)


def _normalize_openai(response: Any) -> dict[str, Any]:
    out: dict[str, Any] = {}

    # Completions / chat completions
    choices = getattr(response, "choices", None)
    if choices:
        choice = choices[0]
        msg = getattr(choice, "message", None)
        if msg:
            message: dict[str, Any] = {"role": getattr(msg, "role", "assistant")}
            if msg.content:
                message["content"] = msg.content
            tcs = getattr(msg, "tool_calls", None)
            if tcs:
                message["tool_calls"] = [
                    {
                        "id": tc.id,
                        "type": getattr(tc, "type", "function"),
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments,
                        },
                    }
                    for tc in tcs
                ]
            out["message"] = message

    # Embeddings
    data = getattr(response, "data", None)
    if data and not choices:
        out["embeddings_count"] = len(data)

    if getattr(response, "model", None):
        out["model"] = response.model

    usage = getattr(response, "usage", None)
    if usage:
        out["usage"] = {
            "input_tokens": getattr(usage, "prompt_tokens", 0) or 0,
            "output_tokens": getattr(usage, "completion_tokens", 0) or 0,
            "total_tokens": getattr(usage, "total_tokens", 0) or 0,
        }
    return out


def _normalize_openai_responses(response: Any) -> dict[str, Any]:
    """Normalize an OpenAI Responses API response into our standard trace format."""
    out: dict[str, Any] = {}
    message: dict[str, Any] = {"role": "assistant"}

    output_text = getattr(response, "output_text", None)
    if output_text:
        message["content"] = output_text

    output_items = getattr(response, "output", None) or []
    tool_calls = []
    for item in output_items:
        if getattr(item, "type", "") == "function_call":
            tool_calls.append({
                "id": getattr(item, "call_id", ""),
                "type": "function",
                "function": {
                    "name": getattr(item, "name", ""),
                    "arguments": getattr(item, "arguments", ""),
                },
            })
    if tool_calls:
        message["tool_calls"] = tool_calls

    out["message"] = message

    if getattr(response, "model", None):
        out["model"] = response.model

    usage = getattr(response, "usage", None)
    if usage:
        input_tokens = getattr(usage, "input_tokens", 0) or 0
        output_tokens = getattr(usage, "output_tokens", 0) or 0
        out["usage"] = {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
        }
    return out


def _normalize_anthropic(response: Any) -> dict[str, Any]:
    out: dict[str, Any] = {}
    message: dict[str, Any] = {"role": "assistant"}

    content_blocks = getattr(response, "content", None)
    if content_blocks:
        text_parts: list[str] = []
        tool_calls: list[dict[str, Any]] = []
        for block in content_blocks:
            btype = getattr(block, "type", "")
            if btype == "text":
                text_parts.append(block.text)
            elif btype == "tool_use":
                inp = getattr(block, "input", {})
                tool_calls.append({
                    "id": block.id,
                    "type": "function",
                    "function": {
                        "name": block.name,
                        "arguments": json.dumps(inp) if isinstance(inp, dict) else str(inp),
                    },
                })
        if text_parts:
            message["content"] = "\n".join(text_parts)
        if tool_calls:
            message["tool_calls"] = tool_calls

    out["message"] = message

    if getattr(response, "model", None):
        out["model"] = response.model

    usage = getattr(response, "usage", None)
    if usage:
        input_tokens = getattr(usage, "input_tokens", 0) or 0
        output_tokens = getattr(usage, "output_tokens", 0) or 0
        out["usage"] = {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
        }
    return out


# ---------------------------------------------------------------------------
# Input capture
# ---------------------------------------------------------------------------


def _capture_input(kwargs: dict[str, Any]) -> dict[str, Any]:
    """Extract a JSON-safe copy of the create() keyword arguments."""
    out: dict[str, Any] = {}
    for key, value in kwargs.items():
        if key == "stream":
            continue
        out[key] = _safe_serialize(value)
    return out


def _safe_serialize(obj: Any) -> Any:
    if obj is None or isinstance(obj, (str, int, float, bool)):
        return obj
    if isinstance(obj, (list, tuple)):
        return [_safe_serialize(item) for item in obj]
    if isinstance(obj, dict):
        return {str(k): _safe_serialize(v) for k, v in obj.items()}
    if hasattr(obj, "model_dump"):
        return obj.model_dump()
    if hasattr(obj, "__dict__"):
        return {
            k: _safe_serialize(v)
            for k, v in obj.__dict__.items()
            if not k.startswith("_")
        }
    return str(obj)


# ---------------------------------------------------------------------------
# Trace submission helpers
# ---------------------------------------------------------------------------


def _format_time(ts: float) -> str:
    return datetime.fromtimestamp(ts, tz=timezone.utc).isoformat()


def _submit_success(
    provider: str,
    resource_name: str,
    extra: dict[str, Any],
    trace_input: dict[str, Any],
    output_data: dict[str, Any],
    started_at: float,
    ended_at: float,
) -> None:
    _submit(provider, resource_name, extra, trace_input, output_data, started_at, ended_at, "success")


def _submit_error(
    provider: str,
    resource_name: str,
    extra: dict[str, Any],
    trace_input: dict[str, Any],
    exc: Exception,
    started_at: float,
) -> None:
    output_data = {"error": str(exc), "error_type": type(exc).__name__}
    _submit(provider, resource_name, extra, trace_input, output_data, started_at, time.time(), "error")


def _submit(
    provider: str,
    resource_name: str,
    extra: dict[str, Any],
    trace_input: dict[str, Any],
    output_data: dict[str, Any],
    started_at: float,
    ended_at: float,
    status: str,
) -> None:
    from carrot_ai import get_client

    parent_id = current_trace_id.get()

    metadata: dict[str, Any] = {
        "provider": provider,
        "source": "sdk",
        "latency_ms": round((ended_at - started_at) * 1000),
    }

    usage = output_data.get("usage")
    if usage:
        metadata.update(usage)

    model = output_data.get("model")
    if model:
        metadata["model"] = model

    if parent_id:
        metadata["parent_trace_id"] = parent_id

    # Trace output only includes the message/error — usage & model go in metadata
    output_payload: dict[str, Any] = {}
    for key in ("message", "embeddings_count", "error", "error_type"):
        if key in output_data:
            output_payload[key] = output_data[key]

    trace: dict[str, Any] = {
        "name": extra.get("name", resource_name),
        "input": trace_input,
        "output": output_payload,
        "metadata": metadata,
        "status": status,
        "started_at": _format_time(started_at),
        "ended_at": _format_time(ended_at),
    }

    tags = extra.get("tags")
    if tags:
        trace["tags"] = list(tags)

    try:
        get_client().submit(trace)
    except Exception:
        logger.debug("carrot_ai: failed to submit trace", exc_info=True)
