"""Anthropic client instrumentation wrapper."""

from __future__ import annotations

import json
import time
from collections.abc import AsyncIterator
from typing import Any, Protocol, runtime_checkable

from .session import AALSession
from .types import EventStatus


@runtime_checkable
class MessagesAPI(Protocol):
    """Duck-typed Anthropic messages API."""

    async def create(self, **kwargs: Any) -> Any: ...


@runtime_checkable
class AnthropicClientLike(Protocol):
    """Duck-typed Anthropic client."""

    messages: MessagesAPI


class _WrappedMessages:
    """Proxy for messages API that intercepts create() calls."""

    def __init__(self, original: MessagesAPI, session: AALSession) -> None:
        self._original = original
        self._session = session

    async def create(self, **kwargs: Any) -> Any:
        start = time.perf_counter()
        response = await self._original.create(**kwargs)

        if kwargs.get("stream"):
            return _wrap_stream_response(response, self._session, start)

        latency_ms = round((time.perf_counter() - start) * 1000)
        _log_tool_use_blocks(response, self._session, latency_ms)
        return response

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


class _WrappedClient:
    """Proxy for Anthropic client that wraps messages.create()."""

    def __init__(self, client: AnthropicClientLike, session: AALSession) -> None:
        self._client = client
        self.messages = _WrappedMessages(client.messages, session)

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


def wrap_anthropic_client(
    client: AnthropicClientLike, session: AALSession
) -> AnthropicClientLike:
    """Wrap the Anthropic client to log tool use actions (non-mutating)."""
    return _WrappedClient(client, session)  # type: ignore[return-value]


def _log_tool_use_blocks(response: Any, session: AALSession, latency_ms: float) -> None:
    """Extract and log tool_use blocks from a non-streaming response."""
    content = getattr(response, "content", [])
    usage = getattr(response, "usage", None)
    input_tokens: int | None = getattr(usage, "input_tokens", None)
    output_tokens: int | None = getattr(usage, "output_tokens", None)
    model: str | None = getattr(response, "model", None)

    for block in content:
        if getattr(block, "type", None) == "tool_use":
            action: dict[str, Any] = {
                "action_type": "TOOL_CALL",
                "tool_name": getattr(block, "name", "unknown"),
                "parameters": getattr(block, "input", {}),
                "status": "SUCCESS",
                "latency_ms": latency_ms,
            }
            if input_tokens is not None:
                action["input_tokens"] = input_tokens
            if output_tokens is not None:
                action["output_tokens"] = output_tokens
            if model is not None:
                action["model"] = model
            session.log_action(action)  # type: ignore[arg-type]


class _StreamWrapper:
    """Wraps an async stream to intercept tool_use blocks."""

    def __init__(self, stream: Any, session: AALSession, start_time: float) -> None:
        self._stream = stream
        self._session = session
        self._start_time = start_time
        self._tool_blocks: dict[int, dict[str, str]] = {}
        self._input_tokens: int | None = None
        self._output_tokens: int | None = None
        self._stream_model: str | None = None

    def __aiter__(self) -> AsyncIterator[Any]:
        return self._iterate()

    async def _iterate(self) -> AsyncIterator[Any]:
        error = None
        try:
            async for event in self._stream:
                event_type = getattr(event, "type", "")

                if event_type == "message_start":
                    msg = getattr(event, "message", None)
                    if msg:
                        self._stream_model = getattr(msg, "model", None)
                        usage = getattr(msg, "usage", None)
                        if usage:
                            self._input_tokens = getattr(usage, "input_tokens", None)

                if event_type == "message_delta":
                    usage = getattr(event, "usage", None)
                    if usage:
                        self._output_tokens = getattr(usage, "output_tokens", None)

                if event_type == "content_block_start":
                    cb = getattr(event, "content_block", None)
                    if cb and getattr(cb, "type", None) == "tool_use":
                        idx = getattr(event, "index", -1)
                        self._tool_blocks[idx] = {
                            "name": getattr(cb, "name", "unknown"),
                            "input_json": "",
                        }

                if event_type == "content_block_delta":
                    delta = getattr(event, "delta", None)
                    if delta and getattr(delta, "type", None) == "input_json_delta":
                        idx = getattr(event, "index", -1)
                        block = self._tool_blocks.get(idx)
                        if block:
                            block["input_json"] += getattr(delta, "partial_json", "")

                yield event
        except Exception as exc:
            error = exc
            raise
        finally:
            if self._tool_blocks:
                latency_ms = round((time.perf_counter() - self._start_time) * 1000)
                status: EventStatus = "ERROR" if error else "SUCCESS"
                for block in self._tool_blocks.values():
                    try:
                        input_data = json.loads(block["input_json"] or "{}")
                    except json.JSONDecodeError:
                        input_data = {}
                    action: dict[str, Any] = {
                        "action_type": "TOOL_CALL",
                        "tool_name": block["name"],
                        "parameters": input_data,
                        "status": status,
                        "latency_ms": latency_ms,
                    }
                    if self._input_tokens is not None:
                        action["input_tokens"] = self._input_tokens
                    if self._output_tokens is not None:
                        action["output_tokens"] = self._output_tokens
                    if self._stream_model is not None:
                        action["model"] = self._stream_model
                    self._session.log_action(action)  # type: ignore[arg-type]

    def __getattr__(self, name: str) -> Any:
        """Proxy all other attributes to the original stream."""
        return getattr(self._stream, name)


def _wrap_stream_response(stream: Any, session: AALSession, start_time: float) -> Any:
    """Wrap a streaming response to intercept tool_use blocks."""
    return _StreamWrapper(stream, session, start_time)
