# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# top-level folder for each specific model found within the models/ directory at
# the top-level of this source tree.

import textwrap
from pathlib import Path

from PIL import Image as PIL_Image

from ..llama3.api.datatypes import *  # noqa: F403
from ..prompt_format import (
    llama3_1_builtin_tool_call_dialog,
    # llama3_1_builtin_tool_call_with_image_dialog,
    llama3_2_user_assistant_conversation,
    TextCompletionContent,
    UseCase,
)


def usecases():
    this_dir = Path(__file__).parent.parent.resolve()
    with open(this_dir / "scripts/resources/dog.jpg", "rb") as f:
        img = PIL_Image.open(f).convert("RGB")

    return [
        llama3_2_user_assistant_conversation(),
        UseCase(
            title="User and assistant conversation with Images",
            description="This example shows how to pass and image to the model as part of the messages.",
            dialogs=[
                [
                    UserMessage(
                        content=[
                            ImageMedia(image=img),
                            "Describe this image in two sentences",
                        ],
                    )
                ],
            ],
            notes=textwrap.dedent(
                """
                - The `<|image|>` tag is used to indicate presence of the image
                - The model isn't an early fusion model so doesn't actually translate an image into several tokens. Instead the cross-attention layers take input "on the side" from a vision encoder
                ![Image](mm-model.png)
                - Its important to postion the <|image|> tag appropriately in the prompt. Image will only attend to the subsequent text tokens
                - The <|image|> tag is part of the user message body, implying that it should only come after the header `<|start_header_id|>{role}<|end_header_id|>` in the message body
                - We recommend using a single image in one prompt
                """
            ),
        ),
        UseCase(
            title="Builtin and Zero Shot Tool Calling",
            description=textwrap.dedent(
                """
                Llama3.2 vision models follow the same tool calling format as Llama3.1 models when inputs are text only.
                Use `Environment: ipython` to enable tools.
                Add `Tools: {{tool_name1}},{{tool_name2}}` for each of the builtin tools.
                The same builtin tools as Llama3.1 are available,
                - code_interpreter (for executing python code)
                - brave_search (to search the web)
                - wolfram_alpha (for querying wolfram alpha for mathematical questions)
                """,
            ),
            dialogs=[llama3_1_builtin_tool_call_dialog()],
            notes=textwrap.dedent(
                """
                - Note the `<|python_tag|>` before `brave_search` function call.
                - The `<|eom_id|>` tag is used to indicate the end of the message.
                - Similar to Llama3.1, code_interpreter is not explicitly mentioned but is enabled via `Environment: ipython`.
                - Tool Calling does NOT work with images in the prompt as of now.
                """
            ),
        ),
        # UseCase(
        #     title="Tool Calling for vision models",
        #     description=textwrap.dedent(
        #         """
        #         While Llama3.2 vision models follow the same tool calling format as Llama3.1 models when inputs are text only,
        #         they are not able to do tool calling when prompt contains image inputs (along with text).
        #         The recommended way would be to separate out the image understanding from the tool calling in successive prompts.
        #         Here is an example of how that could be done,
        #         """,
        #     ),
        #     dialogs=[llama3_1_builtin_tool_call_with_image_dialog()],
        #     notes=textwrap.dedent(
        #         """
        #         - Instead of a single prompt (image understanding + tool call), we split into two prompts to achieve the same result.
        #         """
        #     ),
        # ),
        UseCase(
            title="Prompt format for base models",
            description=textwrap.dedent(
                """
                For base models (Llama3.2-11B-Vision and Llama3.2-90B-Vision), the prompt format for a simple completion is as follows
                """
            ),
            dialogs=[
                TextCompletionContent(
                    content="The color of the sky is blue but sometimes it can also be"
                ),
            ],
            notes="- Same as Llama3.1",
        ),
        UseCase(
            title="Prompt format for base models with Image",
            description=textwrap.dedent(
                """
                For base models (Llama3.2-11B-Vision and Llama3.2-90B-Vision), here is an example of how the text completion format looks with an image,
                """
            ),
            dialogs=[
                TextCompletionContent(
                    content=[
                        ImageMedia(image=img),
                        "If I had to write a haiku for this one",
                    ]
                ),
            ],
            notes="- Note the placement of the special tags <|begin_of_text|> and <|image|>",
        ),
    ]
