"""Anthropic client wrapper for automatic logging."""

from __future__ import annotations

import time
import uuid
from functools import wraps
from typing import TYPE_CHECKING, Any, Iterator

if TYPE_CHECKING:
    from anthropic import Anthropic
    from anthropic.types import Message, RawMessageStreamEvent
    from aitracer.client import AITracer


def wrap_anthropic_client(client: "Anthropic", tracer: "AITracer") -> "Anthropic":
    """
    Wrap an Anthropic client to automatically log all API calls.

    Args:
        client: Anthropic client instance.
        tracer: AITracer instance.

    Returns:
        Wrapped Anthropic client (same instance, modified in place).
    """
    # Store original method
    original_create = client.messages.create

    @wraps(original_create)
    def wrapped_create(*args: Any, **kwargs: Any) -> Any:
        """Wrapped messages.create method."""
        start_time = time.time()
        span_id = str(uuid.uuid4())

        # Extract request data
        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        system = kwargs.get("system")
        max_tokens = kwargs.get("max_tokens", 1024)
        stream = kwargs.get("stream", False)

        try:
            response = original_create(*args, **kwargs)

            if stream:
                # Handle streaming response
                return _wrap_stream_response(
                    response=response,
                    tracer=tracer,
                    model=model,
                    messages=messages,
                    system=system,
                    start_time=start_time,
                    span_id=span_id,
                )
            else:
                # Handle non-streaming response
                latency_ms = int((time.time() - start_time) * 1000)
                _log_message(
                    tracer=tracer,
                    model=model,
                    messages=messages,
                    system=system,
                    response=response,
                    latency_ms=latency_ms,
                    span_id=span_id,
                )
                return response

        except Exception as e:
            # Log error
            latency_ms = int((time.time() - start_time) * 1000)
            tracer.log(
                model=model,
                provider="anthropic",
                input_data=_build_input_data(messages, system),
                output_data=None,
                latency_ms=latency_ms,
                status="error",
                error_message=str(e),
                span_id=span_id,
            )
            raise

    # Replace method
    client.messages.create = wrapped_create  # type: ignore

    return client


def _wrap_stream_response(
    response: Iterator["RawMessageStreamEvent"],
    tracer: "AITracer",
    model: str,
    messages: list,
    system: Any,
    start_time: float,
    span_id: str,
) -> Iterator["RawMessageStreamEvent"]:
    """Wrap streaming response to log after completion."""
    content_parts: list[str] = []
    input_tokens = 0
    output_tokens = 0

    try:
        for event in response:
            # Accumulate content from text deltas
            if hasattr(event, "type"):
                if event.type == "content_block_delta":
                    if hasattr(event, "delta") and hasattr(event.delta, "text"):
                        content_parts.append(event.delta.text)
                elif event.type == "message_delta":
                    if hasattr(event, "usage"):
                        output_tokens = getattr(event.usage, "output_tokens", 0)
                elif event.type == "message_start":
                    if hasattr(event, "message") and hasattr(event.message, "usage"):
                        input_tokens = getattr(event.message.usage, "input_tokens", 0)

            yield event

        # Log after stream completes
        latency_ms = int((time.time() - start_time) * 1000)
        full_content = "".join(content_parts)

        tracer.log(
            model=model,
            provider="anthropic",
            input_data=_build_input_data(messages, system),
            output_data={"content": full_content},
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            status="success",
            span_id=span_id,
        )

    except Exception as e:
        latency_ms = int((time.time() - start_time) * 1000)
        tracer.log(
            model=model,
            provider="anthropic",
            input_data=_build_input_data(messages, system),
            output_data=None,
            latency_ms=latency_ms,
            status="error",
            error_message=str(e),
            span_id=span_id,
        )
        raise


def _log_message(
    tracer: "AITracer",
    model: str,
    messages: list,
    system: Any,
    response: "Message",
    latency_ms: int,
    span_id: str,
) -> None:
    """Log a non-streaming message."""
    # Extract response content
    output_content = None
    if response.content:
        text_parts = []
        for block in response.content:
            if hasattr(block, "text"):
                text_parts.append(block.text)
        output_content = "".join(text_parts)

    input_tokens = getattr(response.usage, "input_tokens", 0) if response.usage else 0
    output_tokens = getattr(response.usage, "output_tokens", 0) if response.usage else 0

    tracer.log(
        model=model,
        provider="anthropic",
        input_data=_build_input_data(messages, system),
        output_data={"content": output_content},
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        latency_ms=latency_ms,
        status="success",
        span_id=span_id,
    )


def _build_input_data(messages: list, system: Any) -> dict:
    """Build input data dictionary."""
    input_data: dict[str, Any] = {"messages": _serialize_messages(messages)}
    if system:
        input_data["system"] = system if isinstance(system, str) else str(system)
    return input_data


def _serialize_messages(messages: list) -> list[dict]:
    """Serialize messages to JSON-compatible format."""
    result = []
    for msg in messages:
        if isinstance(msg, dict):
            result.append(msg)
        elif hasattr(msg, "model_dump"):
            result.append(msg.model_dump())
        elif hasattr(msg, "__dict__"):
            result.append(msg.__dict__)
        else:
            result.append({"content": str(msg)})
    return result
