from judy_core.types import InlineBlobPart, TextPart
from judy_core.utils.tool import ContentUnion
from litellm.types.completion import (
    ChatCompletionContentPartImageParam,
    ChatCompletionContentPartParam,
    ChatCompletionContentPartTextParam,
    ImageURL,
)
from litellm.types.llms.anthropic import (
    AnthropicContentParamSource,
    AnthropicMessagesDocumentParam,
)


def convert_blob_to_image_part(
    blob: InlineBlobPart,
) -> ChatCompletionContentPartImageParam:
    image_url = "data:{};{},{}".format(blob.mime_type, blob.encode_type, blob.data)
    image_part = ChatCompletionContentPartImageParam(
        type="image_url",
        image_url=ImageURL(url=image_url),
    )
    return image_part


def convert_blob_to_document_part(
    blob: InlineBlobPart,
) -> AnthropicMessagesDocumentParam:
    doc_part = AnthropicMessagesDocumentParam(
        type="document",
        source=AnthropicContentParamSource(
            type="base64", media_type=blob.mime_type, data=blob.data
        ),
    )
    return doc_part


def convert_blob_to_part(
    blob: InlineBlobPart, model: str | None = None
) -> ChatCompletionContentPartImageParam | AnthropicMessagesDocumentParam:
    if model:
        if "claude" in model:
            if blob.mime_type.startswith("image/"):
                return convert_blob_to_image_part(blob)
            else:
                return convert_blob_to_document_part(blob)
    return convert_blob_to_image_part(blob)


def convert_content_to_openai_format(
    content: ContentUnion,
    model: str | None = None,
) -> str | list[ChatCompletionContentPartParam]:
    if isinstance(content, str):
        return content
    else:
        message_content = []
        for part in content.parts:
            if isinstance(part, TextPart):
                message_content.append(
                    ChatCompletionContentPartTextParam(
                        {"type": "text", "text": part.text}
                    )
                )
            else:
                blob: InlineBlobPart = part
                part = convert_blob_to_part(blob, model)
                message_content.append(part)
        return message_content
