"""Anthropic SDK monkey-patches for auto-instrumentation."""

from __future__ import annotations

import json
import time
from datetime import UTC, datetime
from typing import Any

import anthropic

from caliper._annotation import generate_request_id, set_last_request_id
from caliper._features import get_linked_request_id
from caliper._patch_helpers import (
    capture_tools_enabled,
    get_sdk_version,
    make_tracked_async_call,
    make_tracked_call,
    merge_metadata,
    pop_caliper_metadata,
    push_record,
)
from caliper._patch_registry import get_original, save_and_patch
from caliper._record import Record

# ---------------------------------------------------------------------------
# Anthropic-specific helpers
# ---------------------------------------------------------------------------


def _classify_error(exc: Exception) -> tuple[str, int | None, str | None]:
    """Map Anthropic exception types to (status, http_status, error_code)."""
    if isinstance(exc, anthropic.RateLimitError):
        return "rate_limited", 429, getattr(exc, "code", None) or "rate_limit_error"
    if isinstance(exc, anthropic.APITimeoutError):
        return "timeout", None, None
    if isinstance(exc, anthropic.APIConnectionError):
        return "timeout", None, None
    if isinstance(exc, anthropic.APIStatusError):
        return "error", exc.status_code, getattr(exc, "code", None)
    return "error", None, None


def _extract_provider_request_id(response: Any) -> str | None:
    """Extract provider's correlation ID from response."""
    return getattr(response, "_request_id", None)


def _extract_tools_provided(kwargs: dict) -> list[str] | None:
    """Extract tool names from the request kwargs."""
    tools = kwargs.get("tools")
    if not tools:
        return None
    names = [t.get("name", "") for t in tools if isinstance(t, dict)]
    return names if names else None


def _extract_tool_calls(response: Any) -> str | None:
    """Extract tool_use blocks from the response as a JSON string."""
    content = getattr(response, "content", None)
    if not content:
        return None
    tool_uses = []
    for block in content:
        if getattr(block, "type", None) == "tool_use":
            tool_uses.append(
                {
                    "name": getattr(block, "name", ""),
                    "id": getattr(block, "id", ""),
                    "input": getattr(block, "input", {}),
                }
            )
    return json.dumps(tool_uses) if tool_uses else None


# ---------------------------------------------------------------------------
# Record builders (Anthropic-specific)
# ---------------------------------------------------------------------------


def _build_record_from_response(
    response: Any,
    kwargs: dict,
    duration_ms: int,
    request_meta: dict[str, str],
    request_id: str,
) -> Record:
    """Build a Record from a successful non-streaming Anthropic response."""
    usage = response.usage
    input_tokens = usage.input_tokens if usage else 0
    output_tokens = usage.output_tokens if usage else 0
    cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0

    tools_provided = _extract_tools_provided(kwargs) if capture_tools_enabled() else None
    tool_calls = _extract_tool_calls(response) if capture_tools_enabled() else None

    return Record(
        request_id=request_id,
        timestamp=datetime.now(UTC).isoformat(),
        provider="anthropic",
        model=response.model,
        endpoint="chat",
        duration_ms=duration_ms,
        tokens_input=input_tokens,
        tokens_output=output_tokens,
        tokens_total=input_tokens + output_tokens,
        status="success",
        http_status=200,
        cached=cache_read > 0,
        provider_request_id=_extract_provider_request_id(response),
        linked_request_id=get_linked_request_id(),
        tools_provided=tools_provided,
        tool_calls=tool_calls,
        metadata=merge_metadata(request_meta),
        sdk_version=get_sdk_version(),
    )


def _build_record_from_error(
    exc: Exception,
    kwargs: dict,
    duration_ms: int,
    request_meta: dict[str, str],
    request_id: str,
) -> Record:
    """Build a Record from a failed Anthropic call."""
    status, http_status, error_code = _classify_error(exc)

    return Record(
        request_id=request_id,
        timestamp=datetime.now(UTC).isoformat(),
        provider="anthropic",
        model=kwargs.get("model", "unknown"),
        endpoint="chat",
        duration_ms=duration_ms,
        tokens_input=0,
        tokens_output=0,
        tokens_total=0,
        status=status,
        http_status=http_status,
        error_code=error_code,
        cached=False,
        linked_request_id=get_linked_request_id(),
        metadata=merge_metadata(request_meta),
        sdk_version=get_sdk_version(),
    )


# ---------------------------------------------------------------------------
# Non-streaming wrappers (via factory)
# ---------------------------------------------------------------------------

_tracked_create = make_tracked_call(
    anthropic.resources.Messages,
    "create",
    _build_record_from_response,
    _build_record_from_error,
)

_tracked_async_create = make_tracked_async_call(
    anthropic.resources.AsyncMessages,
    "create",
    _build_record_from_response,
    _build_record_from_error,
)


# ---------------------------------------------------------------------------
# Streaming wrappers (Anthropic-specific)
# ---------------------------------------------------------------------------


def _tracked_stream(self, **kwargs):
    request_meta = pop_caliper_metadata(kwargs)
    request_id = generate_request_id()
    original_fn = get_original(anthropic.resources.Messages, "stream")
    stream = original_fn(self, **kwargs)
    return _PatchedStream(stream, kwargs, request_meta, request_id)


async def _tracked_async_stream(self, **kwargs):
    request_meta = pop_caliper_metadata(kwargs)
    request_id = generate_request_id()
    original_fn = get_original(anthropic.resources.AsyncMessages, "stream")
    stream = original_fn(self, **kwargs)
    return _AsyncPatchedStream(stream, kwargs, request_meta, request_id)


# ---------------------------------------------------------------------------
# Streaming event processing
# ---------------------------------------------------------------------------


def _process_event(event: Any, state: dict) -> str | None:
    """Process a streaming event, update state, return text delta if any."""
    match event.type:
        case "message_start":
            msg = event.message
            state["model"] = msg.model
            state["provider_request_id"] = _extract_provider_request_id(msg)
            if msg.usage:
                state["input_tokens"] = msg.usage.input_tokens
                state["cache_read"] = getattr(msg.usage, "cache_read_input_tokens", 0) or 0
        case "content_block_start":
            if capture_tools_enabled():
                block = getattr(event, "content_block", None)
                if block and getattr(block, "type", None) == "tool_use":
                    state.setdefault("_tool_blocks", []).append(
                        {
                            "name": getattr(block, "name", ""),
                            "id": getattr(block, "id", ""),
                            "_json_parts": [],
                        }
                    )
        case "content_block_delta":
            if hasattr(event.delta, "text"):
                if state.get("first_token_time") is None:
                    state["first_token_time"] = time.monotonic()
                return event.delta.text
            if capture_tools_enabled() and hasattr(event.delta, "partial_json"):
                tool_blocks = state.get("_tool_blocks")
                if tool_blocks:
                    tool_blocks[-1]["_json_parts"].append(event.delta.partial_json)
        case "content_block_stop":
            if capture_tools_enabled():
                tool_blocks = state.get("_tool_blocks")
                if tool_blocks and tool_blocks[-1].get("_json_parts"):
                    raw = "".join(tool_blocks[-1].pop("_json_parts"))
                    try:
                        tool_blocks[-1]["input"] = json.loads(raw)
                    except (json.JSONDecodeError, TypeError):
                        tool_blocks[-1]["input"] = {}
                elif tool_blocks:
                    tool_blocks[-1].pop("_json_parts", None)
                    tool_blocks[-1].setdefault("input", {})
        case "message_delta":
            if hasattr(event, "usage") and event.usage:
                state["output_tokens"] = event.usage.output_tokens
            state["stop_reason"] = event.delta.stop_reason
    return None


# ---------------------------------------------------------------------------
# _PatchedStream — wraps sync MessageStreamManager
# ---------------------------------------------------------------------------


class _PatchedStream:
    """Wraps an Anthropic MessageStreamManager to capture streaming metrics."""

    def __init__(self, stream_manager: Any, kwargs: dict, request_meta: dict[str, str], request_id: str) -> None:
        self._stream_manager = stream_manager
        self._kwargs = kwargs
        self._request_meta = request_meta
        self._request_id = request_id
        self._stream: Any = None
        self._state: dict = {}
        self._start: float = 0.0

    def __enter__(self):
        self._start = time.monotonic()
        self._state = {}
        self._stream = self._stream_manager.__enter__()
        set_last_request_id(self._request_id)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            self._finalize()
        except Exception:
            pass
        return self._stream_manager.__exit__(exc_type, exc_val, exc_tb)

    def __iter__(self):
        for event in self._stream:
            _process_event(event, self._state)
            yield event

    @property
    def text_stream(self):
        """Yield text deltas from content_block_delta events."""
        for event in self._stream:
            delta = _process_event(event, self._state)
            if delta is not None:
                yield delta

    def __getattr__(self, name: str):
        return getattr(self._stream, name)

    def _finalize(self) -> None:
        elapsed_ms = int((time.monotonic() - self._start) * 1000)
        first_token_time = self._state.get("first_token_time")
        ttft_ms = int((first_token_time - self._start) * 1000) if first_token_time else None
        input_tokens = self._state.get("input_tokens", 0)
        output_tokens = self._state.get("output_tokens", 0)
        cache_read = self._state.get("cache_read", 0)

        tools_provided = _extract_tools_provided(self._kwargs) if capture_tools_enabled() else None
        tool_blocks = self._state.get("_tool_blocks")
        tool_calls = json.dumps(tool_blocks) if capture_tools_enabled() and tool_blocks else None

        record = Record(
            request_id=self._request_id,
            timestamp=datetime.now(UTC).isoformat(),
            provider="anthropic",
            model=self._state.get("model", self._kwargs.get("model", "unknown")),
            endpoint="chat",
            duration_ms=elapsed_ms,
            ttft_ms=ttft_ms,
            tokens_input=input_tokens,
            tokens_output=output_tokens,
            tokens_total=input_tokens + output_tokens,
            status="success",
            http_status=200,
            cached=cache_read > 0,
            provider_request_id=self._state.get("provider_request_id"),
            linked_request_id=get_linked_request_id(),
            tools_provided=tools_provided,
            tool_calls=tool_calls,
            metadata=merge_metadata(self._request_meta),
            sdk_version=get_sdk_version(),
        )
        push_record(record)


# ---------------------------------------------------------------------------
# _AsyncPatchedStream — wraps async AsyncMessageStreamManager
# ---------------------------------------------------------------------------


class _AsyncPatchedStream:
    """Wraps an Anthropic AsyncMessageStreamManager to capture streaming metrics."""

    def __init__(self, stream_manager: Any, kwargs: dict, request_meta: dict[str, str], request_id: str) -> None:
        self._stream_manager = stream_manager
        self._kwargs = kwargs
        self._request_meta = request_meta
        self._request_id = request_id
        self._stream: Any = None
        self._state: dict = {}
        self._start: float = 0.0

    async def __aenter__(self):
        self._start = time.monotonic()
        self._state = {}
        self._stream = await self._stream_manager.__aenter__()
        set_last_request_id(self._request_id)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        try:
            self._finalize()
        except Exception:
            pass
        return await self._stream_manager.__aexit__(exc_type, exc_val, exc_tb)

    async def __aiter__(self):
        async for event in self._stream:
            _process_event(event, self._state)
            yield event

    @property
    async def text_stream(self):
        """Yield text deltas from content_block_delta events."""
        async for event in self._stream:
            delta = _process_event(event, self._state)
            if delta is not None:
                yield delta

    def __getattr__(self, name: str):
        return getattr(self._stream, name)

    def _finalize(self) -> None:
        elapsed_ms = int((time.monotonic() - self._start) * 1000)
        first_token_time = self._state.get("first_token_time")
        ttft_ms = int((first_token_time - self._start) * 1000) if first_token_time else None
        input_tokens = self._state.get("input_tokens", 0)
        output_tokens = self._state.get("output_tokens", 0)
        cache_read = self._state.get("cache_read", 0)

        tools_provided = _extract_tools_provided(self._kwargs) if capture_tools_enabled() else None
        tool_blocks = self._state.get("_tool_blocks")
        tool_calls = json.dumps(tool_blocks) if capture_tools_enabled() and tool_blocks else None

        record = Record(
            request_id=self._request_id,
            timestamp=datetime.now(UTC).isoformat(),
            provider="anthropic",
            model=self._state.get("model", self._kwargs.get("model", "unknown")),
            endpoint="chat",
            duration_ms=elapsed_ms,
            ttft_ms=ttft_ms,
            tokens_input=input_tokens,
            tokens_output=output_tokens,
            tokens_total=input_tokens + output_tokens,
            status="success",
            http_status=200,
            cached=cache_read > 0,
            provider_request_id=self._state.get("provider_request_id"),
            linked_request_id=get_linked_request_id(),
            tools_provided=tools_provided,
            tool_calls=tool_calls,
            metadata=merge_metadata(self._request_meta),
            sdk_version=get_sdk_version(),
        )
        push_record(record)


# ---------------------------------------------------------------------------
# Patch entry point
# ---------------------------------------------------------------------------


def patch() -> None:
    """Patch Anthropic Messages.create and Messages.stream (sync + async)."""
    save_and_patch(anthropic.resources.Messages, "create", _tracked_create)
    save_and_patch(anthropic.resources.Messages, "stream", _tracked_stream)
    save_and_patch(anthropic.resources.AsyncMessages, "create", _tracked_async_create)
    save_and_patch(anthropic.resources.AsyncMessages, "stream", _tracked_async_stream)
