# Copyright 2024 Flyto
# Licensed under the Apache License, Version 2.0
"""Anthropic provider (tool use loop)."""
import logging
from typing import Any, Dict, List, Optional, Tuple

from flyto_ai.models import StreamCallback, StreamEvent, StreamEventType
from flyto_ai.providers.base import (
    DispatchFn, LLMProvider, dispatch_and_log_tool, fire_stream as _fire,
)
from flyto_ai.providers.openai import _looks_like_browser_task

logger = logging.getLogger(__name__)


class AnthropicProvider(LLMProvider):
    """Anthropic Claude provider with tool use loop."""

    def __init__(
        self,
        api_key: str,
        model: str = "claude-sonnet-4-5-20250929",
        temperature: float = 0.7,
        max_tokens: int = 4096,
    ) -> None:
        self._api_key = api_key
        self._model = model
        self._temperature = temperature
        self._max_tokens = max_tokens
        self._client = None

    def __repr__(self) -> str:
        key_hint = "{}...".format(self._api_key[:4]) if self._api_key and len(self._api_key) > 4 else "***"
        return "AnthropicProvider(model={!r}, api_key={!r})".format(self._model, key_hint)

    async def chat(
        self,
        messages: List[Dict[str, Any]],
        system_prompt: str,
        tools: List[Dict],
        dispatch_fn: DispatchFn,
        max_rounds: int = 30,
        on_stream: Optional[StreamCallback] = None,
    ) -> Tuple[str, List[Dict[str, Any]], int, Dict[str, int]]:
        if self._client is None:
            import anthropic
            self._client = anthropic.AsyncAnthropic(api_key=self._api_key)
        client = self._client

        # Convert to Anthropic tool format
        anthropic_tools = [
            {
                "name": t["name"],
                "description": t["description"],
                "input_schema": t["inputSchema"],
            }
            for t in tools
        ]

        # Strip system messages (Anthropic uses system param)
        claude_messages = [
            {"role": msg["role"], "content": msg["content"]}
            for msg in messages
            if msg["role"] in ("user", "assistant")
        ]

        tool_call_log: List[Dict[str, Any]] = []
        total_usage = {
            "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0,
            "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0,
        }

        def _accumulate_usage(response):
            if hasattr(response, "usage") and response.usage:
                u = response.usage
                total_usage["prompt_tokens"] += getattr(u, "input_tokens", 0) or 0
                total_usage["completion_tokens"] += getattr(u, "output_tokens", 0) or 0
                total_usage["total_tokens"] += (getattr(u, "input_tokens", 0) or 0) + (getattr(u, "output_tokens", 0) or 0)
                total_usage["cache_creation_input_tokens"] += getattr(u, "cache_creation_input_tokens", 0) or 0
                total_usage["cache_read_input_tokens"] += getattr(u, "cache_read_input_tokens", 0) or 0

        for round_num in range(max_rounds):
            create_kwargs = {
                "model": self._model,
                "system": system_prompt,
                "messages": claude_messages,
                "max_tokens": self._max_tokens,
                "temperature": self._temperature,
            }
            if anthropic_tools:
                create_kwargs["tools"] = anthropic_tools
                if round_num == 0 and _looks_like_browser_task(messages):
                    create_kwargs["tool_choice"] = {"type": "any"}
                else:
                    create_kwargs["tool_choice"] = {"type": "auto"}

            # ── Streaming path ──────────────────────────────────
            if on_stream is not None:
                async with client.messages.stream(**create_kwargs) as stream:
                    async for text in stream.text_stream:
                        _fire(on_stream, StreamEvent(
                            type=StreamEventType.TOKEN,
                            content=text,
                        ))
                    response = await stream.get_final_message()
            else:
                # ── Non-streaming path ──────────────────────────
                response = await client.messages.create(**create_kwargs)

            _accumulate_usage(response)

            has_tool_use = any(block.type == "tool_use" for block in response.content)

            if not has_tool_use:
                text_parts = [block.text for block in response.content if block.type == "text"]
                content = "\n".join(text_parts)
                if response.stop_reason == "max_tokens":
                    content += "\n\n[Note: Response was truncated due to token limit.]"
                _fire(on_stream, StreamEvent(type=StreamEventType.DONE))
                return content, tool_call_log, round_num + 1, total_usage

            # Strip text blocks from intermediate rounds to prevent fabrication.
            # When the LLM generates text alongside tool_use, the text often contains
            # fabricated results. Removing it forces the LLM to generate text AFTER
            # seeing actual tool results.
            tool_only_content = [block for block in response.content if block.type == "tool_use"]
            claude_messages.append({"role": "assistant", "content": tool_only_content})

            tool_results = []
            for block in response.content:
                if block.type != "tool_use":
                    continue

                func_name = block.name
                func_args = block.input if isinstance(block.input, dict) else {}

                result_str, log_entry, images = await dispatch_and_log_tool(
                    func_name, func_args, dispatch_fn, round_num, on_stream,
                )
                tool_call_log.append(log_entry)

                # ask_user marker — break tool loop, user input needed
                if log_entry.get("result", {}).get("__ASK_USER__"):
                    # Still append a tool_result so Anthropic API is happy
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result_str,
                    })
                    claude_messages.append({"role": "user", "content": tool_results})
                    text = "I need some information from you before I can continue."
                    _fire(on_stream, StreamEvent(type=StreamEventType.DONE))
                    return text, tool_call_log, round_num + 1, total_usage

                if images:
                    content_parts = [{"type": "text", "text": result_str}]
                    for img in images:
                        content_parts.append({
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": img.get("media_type", "image/png"),
                                "data": img["base64"],
                            },
                        })
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": content_parts,
                    })
                else:
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result_str,
                    })

            claude_messages.append({"role": "user", "content": tool_results})

        # Force final answer with improved prompt
        completed = [tc["function"] for tc in tool_call_log if tc.get("ok", True)]
        failed = [tc["function"] for tc in tool_call_log if not tc.get("ok", True)]
        summary_parts = [
            "You have used all {} tool rounds.".format(max_rounds),
            "Completed: {}".format(", ".join(completed[-5:]) if completed else "none"),
        ]
        if failed:
            summary_parts.append("Failed: {}".format(", ".join(failed[-3:])))
        summary_parts.append(
            "Please summarize what was accomplished, what remains incomplete, "
            "and suggest next steps."
        )
        summary_text = {"type": "text", "text": " ".join(summary_parts)}

        if claude_messages and claude_messages[-1]["role"] == "user":
            content = claude_messages[-1]["content"]
            if isinstance(content, list):
                content.append(summary_text)
            else:
                claude_messages.append({"role": "assistant", "content": "I'll summarize now."})
                claude_messages.append({"role": "user", "content": [summary_text]})
        else:
            claude_messages.append({"role": "user", "content": [summary_text]})

        response = await client.messages.create(
            model=self._model,
            system=system_prompt,
            messages=claude_messages,
            max_tokens=self._max_tokens,
            temperature=self._temperature,
        )
        _accumulate_usage(response)
        text_parts = [block.text for block in response.content if block.type == "text"]
        _fire(on_stream, StreamEvent(type=StreamEventType.DONE))
        return "\n".join(text_parts), tool_call_log, max_rounds, total_usage
