# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import base64
import dataclasses
from typing import TYPE_CHECKING, Any, Literal, Union

from marimo import _loggers
from marimo._ai._types import (
    ChatMessage,
    ChatPart,
    DataReasoningPart,
    FilePart,
    ReasoningPart,
    TextPart,
    ToolInvocationPart,
)

if TYPE_CHECKING:
    from anthropic.types import (  # type: ignore[import-not-found]
        ContentBlock,
        DocumentBlockParam,
        ImageBlockParam,
        MessageParam,
        RedactedThinkingBlockParam,
        ServerToolUseBlockParam,
        TextBlockParam,
        ThinkingBlockParam,
        ToolResultBlockParam,
        ToolUseBlockParam,
        WebSearchToolResultBlockParam,
    )
    from google.genai.types import (  # type: ignore[import-not-found]
        BlobDict,
        ContentDict,
        ContentUnionDict,
        PartDict,
    )
    from openai.types.chat import (  # type: ignore[import-not-found]
        ChatCompletionAssistantMessageParam,
        ChatCompletionContentPartImageParam,
        ChatCompletionContentPartParam,
        ChatCompletionContentPartTextParam,
        ChatCompletionMessageParam,
        ChatCompletionMessageToolCallParam,
        ChatCompletionToolMessageParam,
    )
    from openai.types.chat.chat_completion_assistant_message_param import (
        ContentArrayOfContentPart,
    )

LOGGER = _loggers.marimo_logger()


# Message conversions
def get_openai_messages_from_parts(
    role: Literal["system", "user", "assistant"],
    parts: list[ChatPart],
) -> list[dict[str, Any]]:
    messages: list[dict[str, Any]] = []
    for part in parts:
        if isinstance(part, TextPart):
            message = {"role": role, "content": part.text}
            messages.append(message)
        elif isinstance(part, ToolInvocationPart):
            # Create two messages for the tool result
            assistant_message = {
                "role": role,
                "content": None,
                "tool_calls": [
                    {
                        "id": part.tool_call_id,
                        "type": "function",
                        "function": {
                            "name": part.tool_name,
                            "arguments": str(part.input)
                            if part.input
                            else "{}",
                        },
                    }
                ],
            }
            messages.append(assistant_message)
            tool_result_message = {
                "role": "tool",
                "tool_call_id": part.tool_call_id,
                "name": part.tool_name,
                "content": str(part.output),
            }
            messages.append(tool_result_message)
    return messages


def convert_to_openai_messages(
    messages: list[ChatMessage],
) -> list[ChatCompletionMessageParam]:
    openai_messages: list[ChatCompletionMessageParam] = []

    for message in messages:
        if not message.parts or len(message.parts) == 0:
            content_part: ChatCompletionContentPartTextParam = {
                "type": "text",
                "text": message.content,
            }
            openai_messages.append(
                {"role": message.role, "content": [content_part]}  # type: ignore
            )
            continue

        current_parts: list[
            ChatCompletionContentPartTextParam
            | ChatCompletionContentPartParam
            | ContentArrayOfContentPart
        ] = []
        for part in message.parts:
            if isinstance(part, TextPart):
                text_part: ChatCompletionContentPartTextParam = {
                    "type": "text",
                    "text": part.text,
                }
                current_parts.append(text_part)
            elif isinstance(part, FilePart):
                media_type = part.media_type.lstrip()
                if media_type.startswith("image"):
                    image_part: ChatCompletionContentPartImageParam = {
                        "type": "image_url",
                        "image_url": {"url": part.url},
                    }
                    current_parts.append(image_part)
                elif media_type.startswith("text"):
                    file_text_part: ChatCompletionContentPartTextParam = {
                        "type": "text",
                        "text": extract_text(part.url),
                    }
                    current_parts.append(file_text_part)
                else:
                    raise ValueError(f"Unsupported content type {media_type}")
            elif isinstance(part, ToolInvocationPart):
                # For ToolInvocationPart, we need to create separate messages for tool call and tool result
                tool_calls: ChatCompletionMessageToolCallParam = {
                    "id": part.tool_call_id,
                    "type": "function",
                    "function": {
                        "name": part.tool_name,
                        "arguments": str(part.input) if part.input else "{}",
                    },
                }

                tool_invocation_message: ChatCompletionAssistantMessageParam = {
                    "role": "assistant",
                    "tool_calls": [tool_calls],
                }
                openai_messages.append(tool_invocation_message)

                # Create tool result message
                tool_invocation_result: ChatCompletionToolMessageParam = {
                    "tool_call_id": part.tool_call_id,
                    "role": "tool",
                    "content": str(part.output),
                }
                openai_messages.append(tool_invocation_result)

                # Reset parts since we've added the messages
                current_parts = []
            elif dataclasses.is_dataclass(part) and not isinstance(part, type):
                current_parts.append(dataclasses.asdict(part))  # type: ignore
            else:
                LOGGER.debug(
                    "Dropping unsupported part %s during OpenAI conversion",
                    type(part).__name__,
                )

        if current_parts:
            openai_messages.append(
                {"role": message.role, "content": current_parts}  # type: ignore
            )

    return openai_messages


AnthropicParts = Union[
    "ServerToolUseBlockParam",
    "WebSearchToolResultBlockParam",
    "TextBlockParam",
    "ImageBlockParam",
    "ToolUseBlockParam",
    "ToolResultBlockParam",
    "DocumentBlockParam",
    "ThinkingBlockParam",
    "RedactedThinkingBlockParam",
    "ContentBlock",
]


def convert_to_anthropic_messages(
    messages: list[ChatMessage],
) -> list[MessageParam]:
    anthropic_messages: list[MessageParam] = []

    for message in messages:
        anthropic_role = (
            "assistant" if message.role == "system" else message.role
        )

        if not message.parts:
            if not message.content:
                continue

            # Convert content to string
            text_block: TextBlockParam = {
                "type": "text",
                "text": str(message.content),
            }
            anthropic_messages.append(
                {"role": anthropic_role, "content": [text_block]}
            )
            continue

        current_parts: list[AnthropicParts] = []
        data_reasoning_parts = [
            part
            for part in message.parts
            if isinstance(part, DataReasoningPart)
        ]

        for part in message.parts:
            if isinstance(part, TextPart):
                if not part.text:
                    continue

                text_part: TextBlockParam = {
                    "type": "text",
                    "text": part.text,
                }
                current_parts.append(text_part)
            elif isinstance(part, ReasoningPart):
                if not part.text:
                    continue

                signature = ""
                if part.details and len(part.details) > 0:
                    signature = part.details[0].signature or ""

                # Use the first data reasoning part if there is no signature
                # And remove it from the list
                # We can optimize this
                if not signature and data_reasoning_parts:
                    signature = data_reasoning_parts[0].data.signature
                    data_reasoning_parts.pop(0)

                thinking_message: ThinkingBlockParam = {
                    "type": "thinking",
                    "thinking": part.text,
                    "signature": signature,
                }
                current_parts.append(thinking_message)
            elif isinstance(part, ToolInvocationPart):
                # For ToolInvocationPart, we need to create separate messages for tool use and tool result
                tool_use_block: ToolUseBlockParam = {
                    "type": "tool_use",
                    "id": part.tool_call_id,
                    "name": part.tool_name,
                    "input": part.input,
                }

                # Add the tool use to an assistant message
                assistant_message_parts = current_parts + [tool_use_block]
                anthropic_messages.append(
                    {"role": "assistant", "content": assistant_message_parts}
                )

                # Then create the tool result as a user message
                tool_result_message: ToolResultBlockParam = {
                    "tool_use_id": part.tool_call_id,
                    "type": "tool_result",
                    "content": [
                        {
                            "type": "text",
                            "text": str(part.output),
                        }
                    ],
                }
                anthropic_messages.append(
                    {"role": "user", "content": [tool_result_message]}
                )

                # Reset parts since we've added the messages
                current_parts = []
            elif isinstance(part, FilePart):
                media_type = part.media_type.lstrip()
                # Only these types are supported by Anthropic
                if media_type in [
                    "image/jpeg",
                    "image/png",
                    "image/gif",
                    "image/webp",
                ]:
                    file_block: ImageBlockParam = {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": media_type,  # type: ignore
                            "data": _extract_data(part.url),
                        },
                    }
                    current_parts.append(file_block)
                elif media_type.startswith("text"):
                    file_text_part: TextBlockParam = {
                        "type": "text",
                        "text": extract_text(part.url),
                    }
                    current_parts.append(file_text_part)
                else:
                    raise ValueError(f"Unsupported content type {media_type}")

        if current_parts:
            anthropic_messages.append(
                {"role": anthropic_role, "content": current_parts}
            )

    return anthropic_messages


def convert_to_groq_messages(
    messages: list[ChatMessage],
) -> list[dict[Any, Any]]:
    groq_messages: list[dict[Any, Any]] = []

    for message in messages:
        if message.parts:
            # Currently only supports text content (Llava is deprecated now)
            # See here - https://console.groq.com/docs/deprecations
            file_parts = [
                part for part in message.parts if isinstance(part, FilePart)
            ]
            # Convert attachments to text if possible
            text_content = str(message.content)  # Explicitly convert to string
            for file in file_parts:
                if file.media_type.startswith("text"):
                    text_content += "\n" + extract_text(file.url)
                else:
                    raise ValueError(
                        f"Unsupported content type {file.media_type}. Only text content is supported."
                    )

            groq_messages.append(
                {"role": message.role, "content": text_content}
            )
        else:
            groq_messages.append(
                {
                    "role": message.role,
                    "content": str(
                        message.content
                    ),  # Explicitly convert to string
                }
            )

    return groq_messages


def get_google_messages_from_parts(
    role: Literal["system", "user", "assistant"],
    parts: list[ChatPart],
) -> list[ContentDict]:
    messages: list[ContentDict] = []

    for part in parts:
        if isinstance(part, TextPart):
            # Create a message with text content
            text_message: ContentDict = {
                "role": "user" if role == "user" else "model",
                "parts": [{"text": part.text}],
            }
            messages.append(text_message)
        elif isinstance(part, ReasoningPart):
            # Google uses the "thought" field for reasoning content
            # According to Google's thinking models documentation
            reasoning_message: ContentDict = {
                "role": "user" if role == "user" else "model",
                "parts": [{"text": part.text, "thought": True}],
            }
            messages.append(reasoning_message)
        elif isinstance(part, FilePart):
            media_type = part.media_type
            if not media_type.startswith(("image", "text")):
                raise ValueError(f"Unsupported content type {media_type}")
            inline_data: BlobDict = {
                "mime_type": media_type,
                "data": base64.b64decode(_extract_data(part.url)),
            }
            messages.append(
                {
                    "role": "user" if role == "user" else "model",
                    "parts": [{"inline_data": inline_data}],
                }
            )
        elif isinstance(part, ToolInvocationPart):
            # Create function call message for Google
            function_call_message: ContentDict = {
                "role": "model",
                "parts": [
                    {
                        "function_call": {
                            "name": part.tool_name,
                            "args": part.input or {},
                        }
                    }
                ],
            }
            messages.append(function_call_message)

            # Create function response message
            function_response_message: ContentDict = {
                "role": "user",
                "parts": [
                    {
                        "function_response": {
                            "name": part.tool_name,
                            "response": {"result": str(part.output)},
                        }
                    }
                ],
            }
            messages.append(function_response_message)

    return messages


def convert_to_google_messages(
    messages: list[ChatMessage],
) -> list[ContentUnionDict]:
    google_messages: list[ContentUnionDict] = []

    for message in messages:
        parts: list[PartDict] = []
        if not message.parts or len(message.parts) == 0:
            parts.append({"text": str(message.content)})
        else:
            # Convert internal parts to Google parts format
            for parts_message in get_google_messages_from_parts(
                message.role, message.parts
            ):
                if "parts" in parts_message:
                    parts.extend(parts_message["parts"] or [])

        google_messages.append(
            {
                "role": "user" if message.role == "user" else "model",
                "parts": parts,
            }
        )

    return google_messages


def extract_text(url: str) -> str:
    if url.startswith("data:"):
        # extract base64 encoding from url
        data = url.split(",")[1]
        raw = base64.b64decode(data)
        try:
            return raw.decode("utf-8")
        except UnicodeDecodeError:
            # fallback: try latin1
            return raw.decode("latin1")
    else:
        return url


def _extract_data(url: str) -> str:
    if url.startswith("data:"):
        return url.split(",")[1]
    else:
        return url
