"""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", [])
    for block in content:
        if getattr(block, "type", None) == "tool_use":
            session.log_action(
                {
                    "action_type": "TOOL_CALL",
                    "tool_name": getattr(block, "name", "unknown"),
                    "parameters": getattr(block, "input", {}),
                    "status": "SUCCESS",
                    "latency_ms": latency_ms,
                }
            )


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]] = {}

    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 == "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 = {}
                    self._session.log_action(
                        {
                            "action_type": "TOOL_CALL",
                            "tool_name": block["name"],
                            "parameters": input_data,
                            "status": status,
                            "latency_ms": latency_ms,
                        }
                    )

    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)
