"""Processing of Anthropic API responses into litemind Message objects."""

from typing import Any, List, Optional, Union

from pydantic import BaseModel

from litemind.agent.messages.message import Message
from litemind.utils.json_to_object import json_to_object


def process_response_from_anthropic(
    anthropic_response: Any,
    response_format: Optional[Union[BaseModel, str]] = None,
) -> Message:
    """Convert an Anthropic API response into a litemind Message.

    Handles all content block types: text (with citations), tool use/results,
    server tool use (web search, MCP), thinking/redacted thinking, and
    structured output conversion.

    Parameters
    ----------
    anthropic_response : Any
        The raw response object from the Anthropic API.
    response_format : Optional[Union[BaseModel, str]]
        If provided, parses the text response into this Pydantic model.

    Returns
    -------
    Message
        Processed response as a litemind Message with role "assistant".
    """

    # Initialize the processed response:
    processed_response = Message(role="assistant")

    # Check the stop reason and handle special cases:
    stop_reason = anthropic_response.stop_reason
    if stop_reason == "refusal":
        # If the stop reason is refusal, return a message indicating refusal:
        processed_response.append_text("I cannot assist with that request (refusal).")
    elif stop_reason == "max_tokens":
        # Response was truncated due to token limit - add metadata
        processed_response.attributes["truncated"] = True
        processed_response.attributes["stop_reason"] = "max_tokens"
    # "end_turn", "tool_use", "pause_turn", "stop_sequence" are normal completions

    # Text content:
    text_content = ""

    # Variable to hold whether the response is a tool use:
    is_tool_use = False

    # Create a litemind response message from Anthropic's response that includes tool calls:
    for block in anthropic_response.content:
        if block.type == "text":
            # Handle regular text blocks (may include citations)
            text = block.text
            processed_response.append_text(text)
            text_content += text

            # Handle citations if present in the block
            if hasattr(block, "citations") and block.citations:
                # Add citation information as metadata to the last text block
                last_block = processed_response.blocks[-1]
                if not hasattr(last_block, "attributes"):
                    last_block.attributes = {}

                citation_info = []
                for citation in block.citations:
                    if (
                        hasattr(citation, "type")
                        and citation.type == "web_search_result_location"
                    ):
                        cit_dict = {
                            "type": "web_search_citation",
                            "url": getattr(citation, "url", ""),
                            "title": getattr(citation, "title", ""),
                            "cited_text": getattr(citation, "cited_text", ""),
                        }
                        # Preserve encrypted_index for round-tripping
                        encrypted_index = getattr(citation, "encrypted_index", None)
                        if encrypted_index is not None:
                            cit_dict["encrypted_index"] = encrypted_index
                        citation_info.append(cit_dict)

                if citation_info:
                    last_block.attributes["citations"] = citation_info

        elif block.type == "tool_use":
            # Handle regular function tool calls
            processed_response.append_tool_call(
                tool_name=block.name, arguments=block.input, id=block.id
            )
            is_tool_use = True

        elif block.type == "server_tool_use":
            # Handle built-in server tools (web search, MCP, code execution)
            # These are automatically executed by Anthropic, but we track them
            # Mark as server tool so convert_messages can round-trip correctly
            block_ref = processed_response.append_tool_call(
                tool_name=block.name, arguments=block.input, id=block.id
            )
            block_ref.attributes["server_tool_type"] = "server_tool_use"
            is_tool_use = True

        elif block.type == "mcp_tool_use":
            # Handle MCP tool use (similar to server_tool_use)
            # Mark as server tool so convert_messages can round-trip correctly
            block_ref = processed_response.append_tool_call(
                tool_name=block.name, arguments=block.input, id=block.id
            )
            block_ref.attributes["server_tool_type"] = "mcp_tool_use"
            is_tool_use = True

        elif block.type == "web_search_tool_result":
            # Handle web search results.
            # content can be either:
            # - A list of WebSearchResultBlock (success)
            # - A single error object with type "web_search_tool_result_error"
            block_content = block.content

            # Normalize: if content is not a list, wrap it in one
            if not isinstance(block_content, list):
                block_content = [block_content]

            # Convert results to a human-readable string
            result_lines = []
            raw_content = []
            for sr in block_content:
                if hasattr(sr, "type") and sr.type == "web_search_result":
                    title = getattr(sr, "title", "")
                    url = getattr(sr, "url", "")
                    result_lines.append(f"- {title}: {url}")
                    # Store raw content as serializable dicts for round-tripping
                    sr_dict = {"type": "web_search_result"}
                    for attr in [
                        "url",
                        "title",
                        "encrypted_content",
                        "page_age",
                    ]:
                        val = getattr(sr, attr, None)
                        if val is not None:
                            sr_dict[attr] = val
                    raw_content.append(sr_dict)
                elif hasattr(sr, "type") and sr.type == "web_search_tool_result_error":
                    # Handle error results from web search
                    error_code = getattr(sr, "error_code", "unknown")
                    result_lines.append(f"(web search error: {error_code})")
                    raw_content.append(
                        {
                            "type": "web_search_tool_result_error",
                            "error_code": error_code,
                        }
                    )
                else:
                    result_lines.append(str(sr))
                    # Store unknown types as-is for round-tripping if possible
                    if hasattr(sr, "model_dump"):
                        raw_content.append(sr.model_dump())
                    elif hasattr(sr, "__dict__"):
                        raw_content.append(sr.__dict__)

            combined_results = (
                "\n".join(result_lines) if result_lines else "(no results)"
            )

            block_ref = processed_response.append_tool_use(
                tool_name="web_search",
                arguments={},
                result=combined_results,
                id=block.tool_use_id,
            )
            # Store raw block data for round-tripping (preserves encrypted_content)
            block_ref.attributes["raw_server_tool_result"] = {
                "type": "web_search_tool_result",
                "tool_use_id": block.tool_use_id,
                "content": raw_content,
            }

        elif block.type == "mcp_tool_result":
            # Handle MCP tool results
            from anthropic.types.beta import BetaTextBlock

            mcp_tool_results: List[BetaTextBlock] = block.content

            # Combine all text blocks into a single string:
            combined_text = "\n".join(
                [result.text for result in mcp_tool_results if result.type == "text"]
            )

            # Append the combined text to the processed response
            block_ref = processed_response.append_tool_use(
                tool_name="mcp_tool",
                arguments={},
                result=combined_text,
                id=block.tool_use_id,
            )
            # Store raw block data for round-tripping
            raw_mcp_content = [
                {"type": "text", "text": r.text}
                for r in mcp_tool_results
                if r.type == "text"
            ]
            block_ref.attributes["raw_server_tool_result"] = {
                "type": "mcp_tool_result",
                "tool_use_id": block.tool_use_id,
                "content": raw_mcp_content,
            }

        elif block.type == "thinking":
            # Handle thinking blocks (existing functionality)
            processed_response.append_thinking(
                block.thinking, signature=block.signature
            )

        elif block.type == "redacted_thinking":
            # Handle redacted thinking blocks (existing functionality)
            processed_response.append_thinking(block.data, redacted=True)

        else:
            raise ValueError(f"Unexpected block type: {block.type}")

    # Handle structured output conversion if requested
    if response_format and not is_tool_use:
        processed_response = json_to_object(
            processed_response, response_format, text_content
        )

    return processed_response
