"""Anthropic Claude provider using Messages API with tool use support."""

from __future__ import annotations

import inspect
import json
import logging
from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union

from ...config import Config
from ...exceptions import APIError, ConfigurationError
from ...rate_limit import is_rate_limit_error
from ...types import (
    StreamStepBoundary,
    StreamTextDelta,
    StreamToolArgumentsDelta,
    StreamToolCallEnd,
    StreamToolCallStart,
    StreamUsageSink,
    TokenUsage,
)
from ..tool_utils import (
    extract_error_code,
    normalize_tool_result,
    should_break_loop,
    update_step_tracking,
)

logger = logging.getLogger(__name__)


def _convert_tools_to_anthropic(tools_schema: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """Convert OpenAI-style tools to Anthropic format.

    Supports both:
    - OpenAI nested: {"type": "function", "function": {"name", "description", "parameters"}}
    - Flat (tools registry): {"type": "function", "name", "description", "parameters"}
    Anthropic: {"name", "description", "input_schema"} where input_schema is JSON Schema.
    """
    if not tools_schema:
        return []

    anthropic_tools: List[Dict[str, Any]] = []
    for t in tools_schema:
        if t.get("type") == "function":
            fn = t.get("function") or t
            if not isinstance(fn, dict) or "name" not in fn:
                continue
            params = fn.get("parameters", {})
            anthropic_tools.append(
                {
                    "name": fn.get("name", "unknown"),
                    "description": fn.get("description", ""),
                    "input_schema": params if isinstance(params, dict) else {},
                }
            )
        elif "name" in t and "input_schema" in t:
            anthropic_tools.append(t)
        elif "name" in t and "description" in t:
            anthropic_tools.append(
                {
                    "name": t["name"],
                    "description": t.get("description", ""),
                    "input_schema": t.get("input_schema", t.get("parameters", {})),
                }
            )
    return anthropic_tools


def _output_format_from_response_format(response_format: Any) -> Optional[Dict[str, Any]]:
    """Map unified json_schema response_format to Anthropic Messages ``output_format``."""
    if not response_format or not isinstance(response_format, dict):
        return None
    if response_format.get("type") != "json_schema":
        return None
    js = response_format.get("json_schema") or {}
    schema = js.get("schema")
    if not schema:
        return None
    name = js.get("name", "response_schema")
    return {
        "type": "json_schema",
        "json_schema": {
            "name": name,
            "schema": schema,
            "strict": bool(js.get("strict", True)),
        },
    }


class AnthropicMessagesClient:
    """Anthropic Claude client using Messages API with tool use support."""

    def __init__(self, config: Config, request_timeout: Optional[float] = None) -> None:
        self._api_key = config.require_anthropic_key()

        try:
            from anthropic import AsyncAnthropic
        except ImportError as e:
            raise ConfigurationError(
                "Anthropic package not installed. Install with: pip install anthropic"
            ) from e
        except Exception as e:
            raise ConfigurationError(f"Failed to import Anthropic client: {e}") from e

        self._AsyncAnthropic = AsyncAnthropic
        self._request_timeout = request_timeout

    async def _execute_tool(
        self,
        name: Optional[str],
        args: Dict[str, Any],
        execute_tool_cb: Optional[Callable[[str, Dict[str, Any]], Any]],
    ) -> Dict[str, Any]:
        """Execute a tool call via callback."""
        if execute_tool_cb is None:
            return {"ok": False, "result": {}, "error": "no executor provided"}

        try:
            maybe = execute_tool_cb(name, args)
            if inspect.isawaitable(maybe) or hasattr(maybe, "__await__"):
                result = await maybe
            else:
                result = maybe
            return normalize_tool_result(result)
        except Exception as e:
            logger.error(f"Tool execution failed for {name}: {e}")
            return {"ok": False, "result": {}, "error": str(e)}

    def _build_messages(
        self,
        prompt: str,
        messages: Optional[List[Dict[str, Any]]],
    ) -> List[Dict[str, Any]]:
        """Build message list from history + prompt."""
        out: List[Dict[str, Any]] = []
        if messages:
            for msg in messages:
                role = msg.get("role", "user")
                content = msg.get("content", "")
                if role == "assistant":
                    out.append({"role": "assistant", "content": content})
                else:
                    out.append({"role": "user", "content": content})
        out.append({"role": "user", "content": prompt})
        return out

    def _content_to_text(self, content: Any) -> str:
        """Extract text from Anthropic response content blocks."""
        if isinstance(content, str):
            return content
        if not isinstance(content, list):
            return ""
        texts: List[str] = []
        for block in content:
            if isinstance(block, dict):
                if block.get("type") == "text":
                    texts.append(block.get("text", ""))
            elif hasattr(block, "type") and getattr(block, "type") == "text":
                texts.append(getattr(block, "text", "") or "")
        return "".join(texts)

    def _parse_tool_use(self, block: Any) -> Dict[str, Any]:
        """Parse tool_use block to {id, name, input}."""
        if isinstance(block, dict):
            inp = block.get("input", {})
            if isinstance(inp, str):
                try:
                    inp = json.loads(inp) if inp else {}
                except json.JSONDecodeError:
                    inp = {}
            return {
                "id": block.get("id", ""),
                "name": block.get("name", ""),
                "input": inp,
            }
        inp = getattr(block, "input", {}) or {}
        if isinstance(inp, str):
            try:
                inp = json.loads(inp) if inp else {}
            except json.JSONDecodeError:
                inp = {}
        return {
            "id": getattr(block, "id", ""),
            "name": getattr(block, "name", ""),
            "input": inp if isinstance(inp, dict) else {},
        }

    def _get_tool_error_retry_message(
        self,
        output_payloads: List[Dict[str, Any]],
        tool_error_callback: Optional[
            Callable[[str, Optional[str], Dict[str, Any]], Optional[str]]
        ],
    ) -> Optional[str]:
        """Check tool outputs for errors; return retry message if callback provides one."""
        if not tool_error_callback:
            return None
        for out in output_payloads:
            if out["payload"].get("ok"):
                continue
            msg = tool_error_callback(
                out["name"], extract_error_code(out["payload"]), out["payload"]
            )
            if msg:
                return msg
        return None

    async def _finalize_empty_response(
        self,
        client: Any,
        params: Dict[str, Any],
        base_messages: List[Dict[str, Any]],
    ) -> tuple[str, int, int]:
        """Finalize when response is empty; returns (final_text, input_delta, output_delta)."""
        finalize_params = dict(params)
        finalize_params.pop("tools", None)
        finalize_params["messages"] = base_messages + [
            {
                "role": "user",
                "content": "Finalize now. Return the final answer. No further tool calls.",
            }
        ]
        finalize_resp = await client.messages.create(**finalize_params)
        text = self._content_to_text(getattr(finalize_resp, "content", []) or [])
        fu = getattr(finalize_resp, "usage", None)
        inp_delta = getattr(fu, "input_tokens", 0) or 0 if fu else 0
        out_delta = getattr(fu, "output_tokens", 0) or 0 if fu else 0
        return text, inp_delta, out_delta

    def _blocks_to_api_format(self, content: Any) -> List[Dict[str, Any]]:
        """Convert response content blocks to API request format."""
        if not isinstance(content, list):
            return []
        out: List[Dict[str, Any]] = []
        for block in content:
            if isinstance(block, dict):
                out.append(block)
            elif hasattr(block, "model_dump"):
                out.append(block.model_dump(exclude_none=True))
            else:
                btype = getattr(block, "type", "text")
                if btype == "text":
                    out.append({"type": "text", "text": getattr(block, "text", "") or ""})
                elif btype == "tool_use":
                    inp = getattr(block, "input", {})
                    out.append(
                        {
                            "type": "tool_use",
                            "id": getattr(block, "id", ""),
                            "name": getattr(block, "name", ""),
                            "input": inp if isinstance(inp, dict) else {},
                        }
                    )
        return out

    async def generate(
        self,
        *,
        prompt: str,
        model: str,
        messages: Optional[List[Dict[str, Any]]] = None,
        max_tokens: Optional[int] = None,
        top_p: Optional[float] = None,
        presence_penalty: Optional[float] = None,
        instructions: Optional[str] = None,
        reasoning_effort: Optional[str] = None,
        tools_schema: Optional[List[Dict[str, Any]]] = None,
        response_format: Optional[Any] = None,
        execute_tool_cb: Optional[Callable[[str, Dict[str, Any]], Any]] = None,
        tool_error_callback: Optional[
            Callable[[str, Optional[str], Dict[str, Any]], Optional[str]]
        ] = None,
        max_steps: int = 16,
        max_effective_tool_steps: int = 8,
        force_tool_use: bool = False,
        temperature: Optional[float] = None,
        system_instruct: str = "",
    ) -> tuple[str, TokenUsage]:
        if not prompt or not prompt.strip():
            raise ValueError("Prompt cannot be empty")
        if not model or not model.strip():
            raise ValueError("Model name cannot be empty")

        try:
            client_kwargs: Dict[str, Any] = {"api_key": self._api_key}
            if self._request_timeout is not None:
                client_kwargs["timeout"] = self._request_timeout
            client = self._AsyncAnthropic(**client_kwargs)
        except Exception as e:
            raise APIError(f"Failed to initialize Anthropic client: {e}") from e

        anthropic_tools = _convert_tools_to_anthropic(tools_schema or [])
        base_messages = self._build_messages(prompt, messages)

        params: Dict[str, Any] = {
            "model": model,
            "max_tokens": max_tokens or 4096,
            "messages": base_messages,
        }
        if instructions:
            params["system"] = instructions
        if top_p is not None:
            params["top_p"] = top_p
        if anthropic_tools:
            params["tools"] = anthropic_tools
            if force_tool_use:
                params["tool_choice"] = {"type": "any"}
                logger.info("[ANTHROPIC] tool_choice=any (force tool use)")

        out_fmt = _output_format_from_response_format(response_format)
        if out_fmt:
            params["output_format"] = out_fmt

        logger.info(
            "[ANTHROPIC] Request params keys: %s (tool_choice=%s)",
            list(params.keys()),
            params.get("tool_choice"),
        )

        last_resp: Optional[Any] = None
        last_nonempty_output = ""
        effective_steps = 0
        consecutive_reasoning_only = 0
        pending_resp: Optional[Any] = None
        total_input = 0
        total_output = 0

        for step in range(max_steps):
            try:
                if pending_resp is not None:
                    resp = pending_resp
                    pending_resp = None
                else:
                    resp = await client.messages.create(**params)
                last_resp = resp
            except Exception as e:
                if is_rate_limit_error(e):
                    logger.warning(f"Anthropic API rate limit hit at step {step + 1}: {e}")
                    raise
                logger.error(f"Anthropic API call failed at step {step + 1}: {e}")
                raise APIError(f"Anthropic API request failed: {e}") from e

            content = getattr(resp, "content", []) or []
            text = self._content_to_text(content)
            if text.strip():
                last_nonempty_output = text

            stop_reason = getattr(resp, "stop_reason", None) or "end_turn"
            block_types = [
                b.get("type", getattr(b, "type", "?"))
                if isinstance(b, dict)
                else getattr(b, "type", "?")
                for b in content
            ]
            usage = getattr(resp, "usage", None)
            usage_str = ""
            if usage:
                inp = getattr(usage, "input_tokens", None)
                out = getattr(usage, "output_tokens", None)
                if inp is not None:
                    total_input += inp
                if out is not None:
                    total_output += out
                if inp is not None or out is not None:
                    usage_str = f", usage=input={inp or 0} output={out or 0}"

            logger.info(
                "[ANTHROPIC] step=%d stop_reason=%s block_types=%s%s",
                step + 1,
                stop_reason,
                block_types,
                usage_str,
            )

            if stop_reason != "tool_use" and anthropic_tools and force_tool_use:
                logger.warning(
                    "Expected tool_use but got stop_reason=%s (force_tool_use=True). "
                    "Model may have hit max_tokens or returned text-only.",
                    stop_reason,
                )

            if stop_reason == "tool_use":
                tool_uses = [
                    self._parse_tool_use(block)
                    for block in content
                    if (isinstance(block, dict) and block.get("type") == "tool_use")
                    or (getattr(block, "type", None) == "tool_use")
                ]

                if not tool_uses:
                    logger.warning(
                        "[ANTHROPIC] stop_reason=tool_use but no tool_use blocks in content "
                        "(block_types=%s). Breaking.",
                        block_types,
                    )
                    break

                tool_names = [t["name"] for t in tool_uses]
                logger.info("[ANTHROPIC] Executing tools: %s", tool_names)

                output_payloads: List[Dict[str, Any]] = []
                had_non_reasoning_tool = False

                for tu in tool_uses:
                    tu_id = tu["id"]
                    name = tu["name"]
                    inp = tu["input"]
                    result_payload = await self._execute_tool(name, inp, execute_tool_cb)
                    output_payloads.append(
                        {"tool_use_id": tu_id, "name": name, "payload": result_payload}
                    )
                    if name and "reasoning" not in name.lower():
                        had_non_reasoning_tool = True

                retry_message = self._get_tool_error_retry_message(
                    output_payloads, tool_error_callback
                )
                if retry_message is not None:
                    logger.info("[TOOL_ERROR] New session (callback returned retry message)")
                    new_messages = base_messages + [{"role": "user", "content": retry_message}]
                    params["messages"] = new_messages
                    pending_resp = await client.messages.create(**params)
                    continue

                tool_results: List[Dict[str, Any]] = []
                for out in output_payloads:
                    tool_results.append(
                        {
                            "type": "tool_result",
                            "tool_use_id": out["tool_use_id"],
                            "content": json.dumps(out["payload"]),
                            "is_error": not out["payload"].get("ok", False),
                        }
                    )

                new_messages = list(base_messages)
                assistant_content = self._blocks_to_api_format(content)
                new_messages.append({"role": "assistant", "content": assistant_content})
                new_messages.append({"role": "user", "content": tool_results})
                params["messages"] = new_messages
                base_messages = new_messages

                effective_steps, consecutive_reasoning_only = update_step_tracking(
                    had_non_reasoning_tool,
                    effective_steps,
                    consecutive_reasoning_only,
                    max_effective_tool_steps,
                )

                if should_break_loop(
                    effective_steps, consecutive_reasoning_only, max_effective_tool_steps
                ):
                    logger.info(
                        "[ANTHROPIC] Breaking loop: effective_steps=%d consecutive_reasoning=%d",
                        effective_steps,
                        consecutive_reasoning_only,
                    )
                    break
            else:
                logger.info(
                    "[ANTHROPIC] Breaking: stop_reason=%s (not tool_use)",
                    stop_reason,
                )
                break

        final_text = (
            self._content_to_text(getattr(last_resp, "content", []) or []) if last_resp else ""
        )
        if not final_text.strip():
            final_text = last_nonempty_output or ""

        if not final_text.strip():
            try:
                final_text, inp_delta, out_delta = await self._finalize_empty_response(
                    client, params, base_messages
                )
                total_input += inp_delta
                total_output += out_delta
            except Exception as e:
                if is_rate_limit_error(e):
                    raise
                logger.warning(f"Failed to finalize response: {e}")
                if not final_text:
                    raise APIError(f"Failed to generate final response: {e}") from e

        if not final_text.strip():
            usage = TokenUsage(
                input_tokens=total_input,
                output_tokens=total_output,
                total_tokens=total_input + total_output,
                cached_tokens=None,
            )
            return "", usage

        usage = TokenUsage(
            input_tokens=total_input,
            output_tokens=total_output,
            total_tokens=total_input + total_output,
            cached_tokens=None,
        )
        return final_text, usage

    async def generate_stream(
        self,
        *,
        prompt: str,
        model: str,
        messages: Optional[List[Dict[str, Any]]] = None,
        max_tokens: Optional[int] = None,
        top_p: Optional[float] = None,
        temperature: Optional[float] = None,
        instructions: Optional[str] = None,
        system_instruct: str = "",
        presence_penalty: Optional[float] = None,
        reasoning_effort: Optional[str] = None,
        tools_schema: Optional[List[Dict[str, Any]]] = None,
        response_format: Optional[Any] = None,
        execute_tool_cb: Optional[Callable[[str, Dict[str, Any]], Any]] = None,
        tool_error_callback: Optional[
            Callable[[str, Optional[str], Dict[str, Any]], Optional[str]]
        ] = None,
        max_steps: int = 16,
        max_effective_tool_steps: int = 8,
        force_tool_use: bool = False,
        usage_sink: Optional[StreamUsageSink] = None,
    ) -> AsyncIterator[Union[object, TokenUsage]]:
        del presence_penalty, reasoning_effort  # not used on Anthropic Messages

        if not prompt or not prompt.strip():
            raise ValueError("Prompt cannot be empty")
        if not model or not model.strip():
            raise ValueError("Model name cannot be empty")

        try:
            client_kwargs: Dict[str, Any] = {"api_key": self._api_key}
            if self._request_timeout is not None:
                client_kwargs["timeout"] = self._request_timeout
            client = self._AsyncAnthropic(**client_kwargs)
        except Exception as e:
            raise APIError(f"Failed to initialize Anthropic client: {e}") from e

        anthropic_tools = _convert_tools_to_anthropic(tools_schema or [])
        base_messages = self._build_messages(prompt, messages)
        use_tools = bool(anthropic_tools and execute_tool_cb)

        total_input = 0
        total_output = 0
        effective_steps = 0
        consecutive_reasoning_only = 0
        pending_resp: Optional[Any] = None

        def apply_usage_from_message(message: Any) -> None:
            """Accumulate Anthropic message.usage into totals and usage_sink."""
            nonlocal total_input, total_output
            usage = getattr(message, "usage", None)
            if usage is None:
                return
            inp = getattr(usage, "input_tokens", 0) or 0
            out = getattr(usage, "output_tokens", 0) or 0
            total_input += inp
            total_output += out
            if usage_sink is not None:
                usage_sink.replace_with(
                    TokenUsage(
                        total_input,
                        total_output,
                        total_input + total_output,
                        None,
                    )
                )

        try:
            for step in range(max_steps):
                if use_tools and step > 0:
                    yield StreamStepBoundary(step)

                params: Dict[str, Any] = {
                    "model": model,
                    "max_tokens": max_tokens or 4096,
                    "messages": base_messages,
                }
                system = (instructions or system_instruct or "").strip() or None
                if system:
                    params["system"] = system
                if temperature is not None:
                    params["temperature"] = temperature
                if top_p is not None:
                    params["top_p"] = top_p
                if anthropic_tools:
                    params["tools"] = anthropic_tools
                    if force_tool_use:
                        params["tool_choice"] = {"type": "any"}
                out_fmt = _output_format_from_response_format(response_format)
                if out_fmt:
                    params["output_format"] = out_fmt

                if pending_resp is not None:
                    resp = pending_resp
                    pending_resp = None
                    apply_usage_from_message(resp)
                else:
                    current_tool_id: Optional[str] = None
                    resp_msg: Optional[Any] = None
                    async with client.messages.stream(**params) as stream:
                        try:
                            async for event in stream:
                                et = getattr(event, "type", "") or ""
                                if et == "content_block_start":
                                    block = getattr(event, "content_block", None)
                                    btype = None
                                    if block is not None:
                                        btype = getattr(block, "type", None)
                                    if btype == "tool_use":
                                        current_tool_id = str(getattr(block, "id", "") or "")
                                        yield StreamToolCallStart(
                                            current_tool_id,
                                            str(getattr(block, "name", "") or ""),
                                        )
                                elif et == "content_block_delta":
                                    delta = getattr(event, "delta", None)
                                    if delta is None:
                                        continue
                                    dt = getattr(delta, "type", None)
                                    if dt == "text_delta":
                                        t = getattr(delta, "text", "") or ""
                                        if t:
                                            yield StreamTextDelta(t)
                                    elif dt == "input_json_delta" and current_tool_id:
                                        frag = getattr(delta, "partial_json", "") or ""
                                        if frag:
                                            yield StreamToolArgumentsDelta(current_tool_id, frag)
                                elif et == "message_delta" and usage_sink is not None:
                                    u = getattr(event, "usage", None)
                                    if u is not None:
                                        usage_sink.merge(
                                            getattr(u, "input_tokens", 0) or 0,
                                            getattr(u, "output_tokens", 0) or 0,
                                            None,
                                        )
                        finally:
                            try:
                                resp_msg = await stream.get_final_message()
                            except Exception as e:
                                logger.debug(
                                    "Anthropic get_final_message after stream close: %s",
                                    e,
                                    exc_info=True,
                                )
                                resp_msg = None
                            if resp_msg is not None:
                                apply_usage_from_message(resp_msg)
                    resp = resp_msg
                    if resp is None:
                        break

                content = getattr(resp, "content", []) or []
                stop_reason = getattr(resp, "stop_reason", None) or "end_turn"

                if not use_tools:
                    break

                if stop_reason != "tool_use":
                    break

                tool_uses = [
                    self._parse_tool_use(block)
                    for block in content
                    if (isinstance(block, dict) and block.get("type") == "tool_use")
                    or (getattr(block, "type", None) == "tool_use")
                ]
                if not tool_uses:
                    break

                for tu in tool_uses:
                    yield StreamToolCallEnd(
                        id=str(tu["id"]),
                        name=str(tu["name"]),
                        arguments=dict(tu["input"]),
                    )

                output_payloads: List[Dict[str, Any]] = []
                had_non_reasoning_tool = False
                for tu in tool_uses:
                    result_payload = await self._execute_tool(
                        tu["name"], tu["input"], execute_tool_cb
                    )
                    output_payloads.append(
                        {
                            "tool_use_id": tu["id"],
                            "name": tu["name"],
                            "payload": result_payload,
                        }
                    )
                    if tu["name"] and "reasoning" not in tu["name"].lower():
                        had_non_reasoning_tool = True

                retry_message = self._get_tool_error_retry_message(
                    output_payloads, tool_error_callback
                )
                if retry_message is not None:
                    base_messages = base_messages + [{"role": "user", "content": retry_message}]
                    params["messages"] = base_messages
                    pending_resp = await client.messages.create(**params)
                    continue

                tool_results: List[Dict[str, Any]] = []
                for out in output_payloads:
                    tool_results.append(
                        {
                            "type": "tool_result",
                            "tool_use_id": out["tool_use_id"],
                            "content": json.dumps(out["payload"]),
                            "is_error": not out["payload"].get("ok", False),
                        }
                    )

                new_messages = list(base_messages)
                assistant_content = self._blocks_to_api_format(content)
                new_messages.append({"role": "assistant", "content": assistant_content})
                new_messages.append({"role": "user", "content": tool_results})
                base_messages = new_messages

                effective_steps, consecutive_reasoning_only = update_step_tracking(
                    had_non_reasoning_tool,
                    effective_steps,
                    consecutive_reasoning_only,
                    max_effective_tool_steps,
                )
                if should_break_loop(
                    effective_steps, consecutive_reasoning_only, max_effective_tool_steps
                ):
                    break

            yield TokenUsage(
                input_tokens=total_input,
                output_tokens=total_output,
                total_tokens=total_input + total_output,
                cached_tokens=None,
            )
        except Exception as e:
            if is_rate_limit_error(e):
                raise
            raise APIError(f"Anthropic streaming failed: {e}") from e
