import json
import logging
from typing import Any

from ...contexts import CompletionChoice, GenerationUsage, InputMessage, MessagePart
from ...utils import get, safe_json_loads


def get_finish_reason(response_obj: Any) -> str | None:
    """Extract finish reason from Anthropic response."""
    try:
        if hasattr(response_obj, "stop_reason"):
            return response_obj.stop_reason
        return None
    except Exception as e:
        logging.warning("Error getting finish reason from Anthropic", exc_info=e)
        return None


def get_completion_choices(response_obj: Any) -> list[CompletionChoice]:
    """Extract completion choices from Anthropic response."""
    try:
        choices = []
        if hasattr(response_obj, "content") and response_obj.content:
            for content_block in response_obj.content:
                if hasattr(content_block, "type") and content_block.type == "tool_use":
                    tool_call = {
                        "id": content_block.id,
                        "type": "function",
                        "function": {
                            "name": content_block.name,
                            "arguments": json.dumps(
                                content_block.input, ensure_ascii=False, default=repr
                            ),
                        },
                    }
                    choices.append(
                        CompletionChoice(
                            message={
                                "role": getattr(response_obj, "role", "assistant"),
                                "content": None,
                                "function_call": None,
                                "name": None,
                                "tool_calls": [tool_call],
                            },
                            delta=None,
                            text=None,
                        )
                    )
                elif hasattr(content_block, "text"):
                    choices.append(
                        CompletionChoice(
                            message={
                                "role": getattr(response_obj, "role", "assistant"),
                                "content": content_block.text,
                                "function_call": None,
                                "name": None,
                                "tool_calls": None,
                            },
                            delta=None,
                            text=content_block.text,
                        )
                    )
        return choices
    except Exception as e:
        logging.warning("Error getting completion choices from Anthropic", exc_info=e)
        return []


def get_usage_from_response(response_obj: Any) -> GenerationUsage:
    """Extract usage information from Anthropic response."""
    try:
        usage = getattr(response_obj, "usage", None)
        if usage:
            input_tokens = getattr(usage, "input_tokens", None)
            output_tokens = getattr(usage, "output_tokens", None)
            total_tokens = None
            if input_tokens is not None and output_tokens is not None:
                total_tokens = input_tokens + output_tokens

            return {
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "prompt_tokens": input_tokens,
                "completion_tokens": output_tokens,
                "total_tokens": total_tokens,
                "max_tokens": None,
                "temperature": None,
                "top_p": None,
                "stream": None,
                "tool_calls": None,
            }
        return {}
    except Exception as e:
        logging.warning("Error getting usage from Anthropic response", exc_info=e)
        return {}


def get_request_parameters(kwargs: dict) -> GenerationUsage:
    """Extract request parameters from kwargs."""
    try:
        return {
            "max_tokens": kwargs.get("max_tokens"),
            "temperature": kwargs.get("temperature"),
            "top_p": kwargs.get("top_p"),
            "stream": kwargs.get("stream", False),
            "input_tokens": None,
            "output_tokens": None,
            "prompt_tokens": None,
            "completion_tokens": None,
            "total_tokens": None,
            "tool_calls": None,
        }
    except Exception as e:
        logging.warning("Error getting request parameters from Anthropic", exc_info=e)
        return {}


def format_messages_for_tracing(messages: list[dict]) -> list[InputMessage]:
    """Format Anthropic messages for tracing, converting to InputMessage format with parts."""
    try:
        formatted_messages: list[InputMessage] = []

        def _is_tool_result(block: Any) -> bool:
            return (isinstance(block, dict) and block.get("type") == "tool_result") or (
                hasattr(block, "type") and getattr(block, "type", None) == "tool_result"
            )

        def _create_tool_response_part(block: Any) -> MessagePart:
            tool_result_content = get(block, "content", "")
            tool_use_id = get(block, "tool_use_id", None)

            response = safe_json_loads(
                tool_result_content,
                default=tool_result_content if tool_result_content is not None else {},
            )

            tool_response_part = MessagePart(
                type="tool_call_response",
                response=response,
            )
            if tool_use_id:
                tool_response_part["id"] = tool_use_id

            return tool_response_part

        def _process_content_blocks(content: list, role: str) -> list[InputMessage]:
            messages: list[InputMessage] = []
            text_parts: list[MessagePart] = []

            if role == "user" and all(_is_tool_result(block) for block in content):
                for block in content:
                    tool_response_part = _create_tool_response_part(block)
                    messages.append(
                        InputMessage(role="tool", parts=[tool_response_part])
                    )
                return messages

            for block in content:
                block_type = get(block, "type", None)

                if block_type == "text":
                    text_content = get(block, "text", "")
                    text_parts.append(MessagePart(type="text", content=text_content))
                elif block_type == "tool_result":
                    tool_response_part = _create_tool_response_part(block)
                    messages.append(
                        InputMessage(role="tool", parts=[tool_response_part])
                    )

            if text_parts:
                messages.insert(0, InputMessage(role=role, parts=text_parts))
            elif not any(_is_tool_result(block) for block in content):
                messages.insert(
                    0,
                    InputMessage(
                        role=role, parts=[MessagePart(type="text", content="")]
                    ),
                )

            return messages

        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")

            if isinstance(content, list):
                formatted_messages.extend(_process_content_blocks(content, role))
            else:
                parts = [MessagePart(type="text", content=str(content or ""))]
                formatted_messages.append(InputMessage(role=role, parts=parts))

        return formatted_messages
    except Exception as e:
        logging.warning("Error formatting messages for tracing", exc_info=e)
        return []


def get_stream_text(chunk: Any) -> str:
    """Extract text from streaming chunk."""
    try:
        if hasattr(chunk, "delta"):
            delta = chunk.delta
            if hasattr(delta, "text"):
                return delta.text or ""
        if hasattr(chunk, "content_block"):
            content_block = chunk.content_block
            if hasattr(content_block, "text"):
                return content_block.text or ""
        return ""
    except Exception:
        return ""
