"""
Anthropic integration — wraps anthropic.Anthropic to auto-trace message
completions (streaming and non-streaming).

For streaming, the wrapper yields events from Anthropic's ``Stream`` iterator
while silently buffering ``content_block_delta`` text.  When the stream ends
the aggregated trace is dispatched.

Usage:
    import anthropic
    from syntropy import Syntropy
    from syntropy.integrations.anthropic import wrap_anthropic

    syn = Syntropy(api_key="syn_...", agent_id="my-agent")
    client = wrap_anthropic(anthropic.Anthropic(), syntropy=syn)

    # Non-streaming (auto-traced):
    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        messages=[{"role": "user", "content": "Hello!"}],
    )

    # Streaming (auto-traced on completion):
    with client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        messages=[{"role": "user", "content": "Hello!"}],
        stream=True,
    ) as stream:
        for event in stream:
            ...
"""

from __future__ import annotations

import time
from typing import Any, Optional

from syntropy.client import Syntropy
from syntropy.trace_context import is_inside_trace


# ----------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------

def _extract_prompt(kwargs: dict[str, Any]) -> str:
    messages = kwargs.get("messages", [])
    last_msg = messages[-1] if messages else {}
    content = last_msg.get("content", "") if isinstance(last_msg, dict) else str(last_msg)
    return content if isinstance(content, str) else str(content)


def _wrap_anthropic_stream(
    stream: Any,
    kwargs: dict[str, Any],
    start: float,
    syntropy: Syntropy,
    agent_id: str,
) -> Any:
    """Yield events from an Anthropic Stream, buffer text, trace on end."""
    text_parts: list[str] = []
    model = kwargs.get("model", "unknown")
    tokens_in = 0
    tokens_out = 0
    stop_reason: str | None = None
    stream_error: BaseException | None = None

    try:
        for event in stream:
            event_type = getattr(event, "type", "")

            if event_type == "content_block_delta":
                delta = getattr(event, "delta", None)
                if delta and hasattr(delta, "text"):
                    text_parts.append(delta.text)

            elif event_type == "message_start":
                msg = getattr(event, "message", None)
                if msg:
                    model = getattr(msg, "model", model)
                    usage = getattr(msg, "usage", None)
                    if usage:
                        tokens_in = getattr(usage, "input_tokens", 0)

            elif event_type == "message_delta":
                delta = getattr(event, "delta", None)
                if delta:
                    stop_reason = getattr(delta, "stop_reason", stop_reason)
                usage = getattr(event, "usage", None)
                if usage:
                    tokens_out = getattr(usage, "output_tokens", tokens_out)

            yield event
    except BaseException as exc:
        stream_error = exc
        raise
    finally:
        elapsed_ms = int((time.perf_counter() - start) * 1000)
        syntropy.trace(
            prompt=_extract_prompt(kwargs),
            response="".join(text_parts),
            model=model,
            agent_id=agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            status="error" if stream_error else "success",
            metadata={
                "provider": "anthropic",
                "stream": True,
                "stop_reason": stop_reason,
                "max_tokens": kwargs.get("max_tokens"),
                **({
                    "error": str(stream_error),
                    "partialResponse": True,
                } if stream_error else {}),
            },
        )


# ----------------------------------------------------------------
# Traced proxy classes
# ----------------------------------------------------------------


class _TracedMessages:
    """Proxy around anthropic.Messages that records traces."""

    def __init__(self, original: Any, syntropy: Syntropy, agent_id: str):
        self._original = original
        self._syntropy = syntropy
        self._agent_id = agent_id

    def create(self, **kwargs: Any) -> Any:
        # Skip if already inside an orchestration-level trace (e.g. LangChain)
        if is_inside_trace():
            return self._original.create(**kwargs)

        start = time.perf_counter()
        is_stream = kwargs.get("stream", False)

        response = self._original.create(**kwargs)

        # ---- Streaming path ----
        if is_stream:
            return _wrap_anthropic_stream(
                response, kwargs, start, self._syntropy, self._agent_id,
            )

        # ---- Non-streaming path ----
        elapsed_ms = int((time.perf_counter() - start) * 1000)

        prompt = _extract_prompt(kwargs)

        output = ""
        if hasattr(response, "content") and response.content:
            parts = [
                block.text
                for block in response.content
                if hasattr(block, "text")
            ]
            output = "\n".join(parts)

        usage = getattr(response, "usage", None)
        tokens_in = getattr(usage, "input_tokens", 0) if usage else 0
        tokens_out = getattr(usage, "output_tokens", 0) if usage else 0

        self._syntropy.trace(
            prompt=prompt,
            response=output,
            model=getattr(response, "model", kwargs.get("model", "unknown")),
            agent_id=self._agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "anthropic",
                "stream": False,
                "stop_reason": getattr(response, "stop_reason", None),
                "max_tokens": kwargs.get("max_tokens"),
            },
        )

        return response

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


class _TracedAnthropicMessages:
    """Proxy that injects traced messages namespace."""

    def __init__(self, original: Any, syntropy: Syntropy, agent_id: str):
        self._original = original
        self._syntropy = syntropy
        self._agent_id = agent_id
        self.messages = _TracedMessages(
            original.messages, syntropy, agent_id
        )

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


def wrap_anthropic(
    client: Any,
    *,
    syntropy: Syntropy,
    agent_id: str = "",
) -> Any:
    """
    Wrap an Anthropic client to auto-trace all messages.create() calls.
    Supports both streaming and non-streaming.

    Args:
        client: An anthropic.Anthropic() instance.
        syntropy: A Syntropy client instance.
        agent_id: Override agent_id (defaults to syntropy config).

    Returns:
        A drop-in replacement with tracing enabled.
    """
    return _TracedAnthropicMessages(
        client, syntropy, agent_id or syntropy._config.agent_id
    )
