"""
Anthropic tool_use loop integration for Amplitude AI SDK.

Provides :class:`AmplitudeToolLoop` — a high-level wrapper that manages
Anthropic's multi-turn tool_use agentic loop and automatically tracks
each turn in Amplitude.

Usage::

    from anthropic import Anthropic
    from amplitude import Amplitude
    from amplitude_ai.integrations.anthropic_tools import AmplitudeToolLoop

    amplitude = Amplitude("YOUR_API_KEY")
    client = Anthropic()

    loop = AmplitudeToolLoop(
        amplitude=amplitude,
        client=client,
        user_id="user-1",
        tool_handlers={
            "get_weather": get_weather_fn,
            "search": search_fn,
        },
    )

    result = loop.run(
        model="claude-sonnet-4-20250514",
        messages=[{"role": "user", "content": "What's the weather in SF?"}],
        tools=[...],
    )
"""

from __future__ import annotations

import time
import uuid
from typing import Any, Callable, Dict, Optional

try:
    from anthropic import Anthropic as _AnthropicClient

    ANTHROPIC_AVAILABLE = True
except ImportError:
    ANTHROPIC_AVAILABLE = False
    _AnthropicClient = object  # type: ignore[assignment,misc,unused-ignore]

from amplitude import Amplitude

from ..context import get_active_context
from ..core.privacy import PrivacyConfig
from ..core.tracking import (
    track_ai_message,
    track_tool_call,
    track_user_message,
)
from ..utils.costs import calculate_cost
from ..utils.logger import get_logger


class AmplitudeToolLoop:
    """Manages Anthropic's multi-turn tool_use loop with Amplitude tracking.

    Each iteration of the loop:

    1. Calls ``client.messages.create(...)``
    2. Emits ``[Agent] AI Response`` with tool_calls (if any)
    3. For each ``tool_use`` block, calls the registered handler and
       emits ``[Agent] Tool Call``
    4. Feeds ``tool_result`` messages back and repeats

    The loop stops when the model returns ``stop_reason != "tool_use"``
    or ``max_iterations`` is reached.
    """

    def __init__(
        self,
        amplitude: Amplitude,
        client: Any,
        user_id: str,
        tool_handlers: Dict[str, Callable[..., Any]],
        privacy_config: Optional[PrivacyConfig] = None,
        session_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        parent_agent_id: Optional[str] = None,
        customer_org_id: Optional[str] = None,
        agent_version: Optional[str] = None,
        context: Optional[Dict[str, Any]] = None,
        env: Optional[str] = None,
        groups: Optional[Dict[str, Any]] = None,
        max_iterations: int = 10,
    ) -> None:
        if not ANTHROPIC_AVAILABLE:
            from ..exceptions import ProviderError

            raise ProviderError(
                "Anthropic SDK not installed. Install with: pip install anthropic"
            )

        ctx = get_active_context()

        self.amplitude = amplitude
        self.client = client
        self.user_id = user_id or (ctx.user_id if ctx else None) or ""
        self.tool_handlers = tool_handlers
        self.privacy_config = privacy_config or PrivacyConfig()
        self.session_id = session_id or (ctx.session_id if ctx else None) or str(uuid.uuid4())
        self.trace_id = trace_id or (ctx.trace_id if ctx else None) or str(uuid.uuid4())
        self.agent_id = agent_id or (ctx.agent_id if ctx else None)
        self.parent_agent_id = parent_agent_id or (ctx.parent_agent_id if ctx else None)
        self.customer_org_id = customer_org_id or (ctx.customer_org_id if ctx else None)
        self.agent_version = agent_version or (ctx.agent_version if ctx else None)
        self.context = context or (ctx.context if ctx else None)
        self.env = env or (ctx.env if ctx else None)
        self.groups = groups or (ctx.groups if ctx else None)
        self.max_iterations = max_iterations

        self._turn_counter = 1

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def run(self, **kwargs: Any) -> Any:
        """Run the synchronous tool_use loop.

        Accepts all kwargs that ``anthropic.Anthropic().messages.create()``
        supports (``model``, ``messages``, ``tools``, ``system``, etc.).

        Returns the final ``anthropic.types.Message`` (the one where
        ``stop_reason != "tool_use"``).
        """
        messages = list(kwargs.pop("messages", []))
        model = kwargs.get("model", "unknown")

        self._track_initial_user_messages(messages)

        response = None
        for _ in range(self.max_iterations):
            start_time = time.time()
            response = self.client.messages.create(messages=messages, **kwargs)
            latency_ms = (time.time() - start_time) * 1000

            self._track_ai_response(response, model, latency_ms)

            if getattr(response, "stop_reason", None) != "tool_use":
                return response

            tool_results = self._execute_tool_calls(response, messages)
            messages.append({"role": "assistant", "content": response.content})
            messages.append({"role": "user", "content": tool_results})

        return response

    async def arun(self, **kwargs: Any) -> Any:
        """Run the async tool_use loop.

        Same as :meth:`run` but uses ``client.messages.create()`` with
        ``await``.  Requires ``anthropic.AsyncAnthropic`` as the client.

        Note: tool handlers registered in ``tool_handlers`` are invoked
        synchronously.  If your handlers perform async I/O, wrap them
        or use :meth:`run` in a thread instead.
        """
        messages = list(kwargs.pop("messages", []))
        model = kwargs.get("model", "unknown")

        self._track_initial_user_messages(messages)

        response = None
        for _ in range(self.max_iterations):
            start_time = time.time()
            response = await self.client.messages.create(messages=messages, **kwargs)
            latency_ms = (time.time() - start_time) * 1000

            self._track_ai_response(response, model, latency_ms)

            if getattr(response, "stop_reason", None) != "tool_use":
                return response

            tool_results = self._execute_tool_calls(response, messages)
            messages.append({"role": "assistant", "content": response.content})
            messages.append({"role": "user", "content": tool_results})

        return response

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _track_initial_user_messages(self, messages: list[Any]) -> None:
        """Track user messages from the initial message list."""
        for msg in messages:
            role = msg.get("role", "") if isinstance(msg, dict) else getattr(msg, "role", "")
            content = msg.get("content", "") if isinstance(msg, dict) else getattr(msg, "content", "")
            if role == "user" and isinstance(content, str) and content:
                track_user_message(
                    amplitude=self.amplitude,
                    user_id=self.user_id,
                    message_content=content,
                    session_id=self.session_id,
                    trace_id=self.trace_id,
                    turn_id=self._turn_counter,
                    agent_id=self.agent_id,
                    parent_agent_id=self.parent_agent_id,
                    customer_org_id=self.customer_org_id,
                    agent_version=self.agent_version,
                    context=self.context,
                    env=self.env,
                    groups=self.groups,
                    privacy_config=self.privacy_config,
                )
                self._turn_counter += 1

    def _track_ai_response(
        self,
        response: Any,
        model: str,
        latency_ms: float,
    ) -> None:
        """Track a [Agent] AI Response event from an Anthropic Message."""
        _logger = get_logger(self.amplitude)

        try:
            response_text = _extract_text_content(response)
            tool_calls = _extract_tool_calls(response)

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

            cache_read = getattr(usage, "cache_read_input_tokens", None) if usage else None
            cache_creation = getattr(usage, "cache_creation_input_tokens", None) if usage else None

            total_tokens = None
            if input_tokens is not None and output_tokens is not None:
                total_tokens = input_tokens + output_tokens

            total_cost_usd = None
            if input_tokens is not None and output_tokens is not None:
                try:
                    total_cost_usd = calculate_cost(
                        model_name=model,
                        input_tokens=input_tokens,
                        output_tokens=output_tokens,
                        cache_read_input_tokens=cache_read or 0,
                        cache_creation_input_tokens=cache_creation or 0,
                    )
                except Exception:
                    pass

            finish_reason = getattr(response, "stop_reason", None)

            track_ai_message(
                amplitude=self.amplitude,
                user_id=self.user_id,
                model_name=model,
                provider="anthropic",
                response_content=response_text,
                latency_ms=latency_ms,
                session_id=self.session_id,
                trace_id=self.trace_id,
                turn_id=self._turn_counter,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                total_cost_usd=total_cost_usd,
                finish_reason=finish_reason,
                tool_calls=tool_calls,
                cache_read_input_tokens=cache_read,
                cache_creation_input_tokens=cache_creation,
                agent_id=self.agent_id,
                parent_agent_id=self.parent_agent_id,
                customer_org_id=self.customer_org_id,
                agent_version=self.agent_version,
                context=self.context,
                env=self.env,
                groups=self.groups,
                privacy_config=self.privacy_config,
            )
            self._turn_counter += 1
        except Exception as exc:
            _logger.warning("Anthropic tool loop AI response tracking failed: %s", exc)

    def _execute_tool_calls(
        self,
        response: Any,
        messages: list[Any],
    ) -> list[Dict[str, Any]]:
        """Execute tool_use blocks from the response and track each call.

        Returns a list of ``tool_result`` content blocks to feed back
        to the API.
        """
        _logger = get_logger(self.amplitude)
        tool_results: list[Dict[str, Any]] = []

        content_blocks = getattr(response, "content", []) or []
        for block in content_blocks:
            block_type = getattr(block, "type", None)
            if block_type != "tool_use":
                continue

            tool_name = getattr(block, "name", "unknown")
            tool_id = getattr(block, "id", "")
            tool_input = getattr(block, "input", {})

            error_msg: Optional[str]
            handler = self.tool_handlers.get(tool_name)
            if handler is None:
                error_msg = f"No handler registered for tool '{tool_name}'"
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": tool_id,
                    "content": error_msg,
                    "is_error": True,
                })
                try:
                    track_tool_call(
                        amplitude=self.amplitude,
                        user_id=self.user_id,
                        tool_name=tool_name,
                        success=False,
                        latency_ms=0,
                        session_id=self.session_id,
                        trace_id=self.trace_id,
                        turn_id=self._turn_counter,
                        tool_input=tool_input if isinstance(tool_input, dict) else None,
                        error_message=error_msg,
                        agent_id=self.agent_id,
                        parent_agent_id=self.parent_agent_id,
                        customer_org_id=self.customer_org_id,
                        agent_version=self.agent_version,
                        context=self.context,
                        env=self.env,
                        groups=self.groups,
                        privacy_config=self.privacy_config,
                    )
                except Exception as exc:
                    _logger.warning("Tool call tracking failed: %s", exc)
                continue

            start_time = time.time()
            error_msg = None
            result_content = ""
            success = True

            try:
                if isinstance(tool_input, dict):
                    result = handler(**tool_input)
                else:
                    result = handler(tool_input)
                result_content = str(result)
            except Exception as exc:
                success = False
                error_msg = str(exc)
                result_content = f"Error: {exc}"

            tool_latency_ms = (time.time() - start_time) * 1000

            tool_results.append({
                "type": "tool_result",
                "tool_use_id": tool_id,
                "content": result_content,
                **({"is_error": True} if not success else {}),
            })

            try:
                track_tool_call(
                    amplitude=self.amplitude,
                    user_id=self.user_id,
                    tool_name=tool_name,
                    success=success,
                    latency_ms=tool_latency_ms,
                    session_id=self.session_id,
                    trace_id=self.trace_id,
                    turn_id=self._turn_counter,
                    tool_input=tool_input if isinstance(tool_input, dict) else None,
                    tool_output=result_content,
                    error_message=error_msg,
                    agent_id=self.agent_id,
                    parent_agent_id=self.parent_agent_id,
                    customer_org_id=self.customer_org_id,
                    agent_version=self.agent_version,
                    context=self.context,
                    env=self.env,
                    groups=self.groups,
                    privacy_config=self.privacy_config,
                )
            except Exception as exc:
                _logger.warning("Tool call tracking failed: %s", exc)

        return tool_results


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


def _extract_text_content(response: Any) -> str:
    """Extract text content from an Anthropic Message response."""
    content_blocks = getattr(response, "content", []) or []
    parts: list[str] = []
    for block in content_blocks:
        if getattr(block, "type", None) == "text":
            parts.append(getattr(block, "text", ""))
    return "".join(parts)


def _extract_tool_calls(response: Any) -> Optional[list[Dict[str, Any]]]:
    """Extract tool_use blocks as normalized tool_calls."""
    content_blocks = getattr(response, "content", []) or []
    tool_calls: list[Dict[str, Any]] = []
    for block in content_blocks:
        if getattr(block, "type", None) == "tool_use":
            tool_calls.append({
                "type": "function",
                "id": getattr(block, "id", ""),
                "function": {
                    "name": getattr(block, "name", ""),
                    "arguments": str(getattr(block, "input", {})),
                },
            })
    return tool_calls if tool_calls else None
