from __future__ import annotations

from typing import Any


def wrap_anthropic(client: Any) -> Any:
    """Wrap an Anthropic client to automatically trace messages.create() calls."""
    try:
        original_create = client.messages.create
    except AttributeError:
        print("[ZappyBee] Could not wrap client - no messages.create method found.")
        return client

    async def traced_create(*args: Any, **kwargs: Any) -> Any:
        from ..client import TokenCat
        from ..trace import Trace

        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        trace_input = {"model": model, "messages": messages[:2]} if messages else {"model": model}

        trace = await TokenCat.start_trace(
            f"anthropic/{model}",
            input=trace_input,
        )

        step = await trace.start_step(
            "messages.create",
            "llm_call",
            input=trace_input,
            model=model,
        )

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

            prompt_tokens = getattr(getattr(response, "usage", None), "input_tokens", None)
            completion_tokens = getattr(getattr(response, "usage", None), "output_tokens", None)

            output = None
            if hasattr(response, "content") and response.content:
                output = [
                    {"type": getattr(b, "type", "text"), "text": getattr(b, "text", "")}
                    for b in response.content
                ]

            await step.end(
                output=output,
                status="success",
                model=model,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
            )
            await trace.end(output=output, status="success")
            return response

        except Exception as e:
            await step.end(status="error", error_message=str(e), model=model)
            await trace.end(status="error", error_message=str(e))
            raise

    # Monkey-patch
    client.messages.create = traced_create
    return client
