#
# SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
# SPDX-License-Identifier: Apache-2.0
#

import json  # Import JSON module for encoding and decoding JSON data
from src.tools.image import ImageGeneration  # Import ImageGeneration class to handle image creation

# Asynchronous handler for image generation command
async def image_integration(
    input,  # User input containing the /image command and instruction
    new_history,  # Conversation history in message format
    session_id,  # Session ID for conversation context
    selected_model,  # Selected AI model for generation
    jarvis,  # AI backend function for generating responses
    mode,  # Mode for AI response generation
    temperature,  # Temperature parameter for AI
    top_k,  # Top-k parameter for AI
    min_p,  # Min-p parameter for AI
    top_p,  # Top-p parameter for AI
    repetition_penalty  # Repetition penalty for AI
):
    # Extract the image generation instruction after the '/image' command prefix and strip whitespace
    generate_image_instruction = input[6:].strip()  # Get instruction after /image

    # If no instruction text is provided after the command, yield empty and exit early
    if not generate_image_instruction:  # Check if instruction is empty
        yield []  # Yield empty list for missing instruction
        return  # Exit function

    try:  # Try block for image generation
        # Asynchronously create image content based on the instruction using ImageGeneration class
        image = await ImageGeneration.create_image(generate_image_instruction)  # Generate image

        # Serialize the image data and instruction into a JSON formatted string for processing
        image_generation_content = json.dumps({
            "image": image,  # Image content or URL
            "generate_image_instruction": generate_image_instruction  # Instruction for image generation
        })

        # Construct the conversation history including the image generation result and formatting instructions
        image_generation_result = (
            new_history
            + [
                {
                    "role": "system",
                    "content": (
                        "Image generation result:\n\n" + image_generation_content + "\n\n\n"
                        "Show the generated image using the following markdown syntax format, where '{image_link}' is the URL of the image:\n\n"
                        "![Generated Image]({image_link})\n\n"
                        "Please replace '{image_link}' with the actual image URL provided in the context.\n\n"
                        "Then, describe the generated image based on the above information.\n\n\n"
                        "Use the same language as the previous user input or user request.\n"
                        "For example, if the previous user input or user request is in Indonesian, explain in Indonesian.\n"
                        "If it is in English, explain in English. This also applies to other languages.\n\n\n"
                    )
                }
            ]
        )

        # Use async generator to get descriptive text about the generated image from AI
        async for image_description in jarvis(
            session_id=session_id,  # Session ID
            model=selected_model,  # Selected model
            history=image_generation_result,  # Updated history with image result
            user_message=input,  # User input
            mode=mode,  # Mode for AI response
            temperature=temperature,  # temperature parameter
            top_k=top_k,  # top_k parameter
            min_p=min_p,  # min_p parameter
            top_p=top_p,  # top_p parameter
            repetition_penalty=repetition_penalty  # repetition_penalty parameter
        ):
            yield [{"role": "tool", "content": image_description}]  # Yield image description in tool role
        return  # Exit after handling image

    except Exception:  # Exception handling for image generation failure
        # If image generation fails, let AI generate a contextual error message
        generation_failed = (
            new_history
            + [
                {
                    "role": "system",
                    "content": (
                        "Image generation failed for the user's request. The user tried to generate an image with the instruction: '"
                        + generate_image_instruction + "'\n\n\n"
                        "Please explain to the user that the image generation has failed.\n"
                        "Also, explain all possible reasons why it might have failed.\n"
                        "Additionally, advise the user not to include, input, or attach any sensitive content, "
                        "as this may also cause the image generation process to fail.\n\n"
                        "Be helpful and empathetic in your response.\n\n\n"
                        "Use the same language as the previous user input or user request.\n"
                        "For example, if the previous user input or user request is in Indonesian, explain in Indonesian.\n"
                        "If it is in English, explain in English. This also applies to other languages.\n\n\n"
                    )
                }
            ]
        )

        # Use AI to generate a contextual error message
        async for error_response in jarvis(
            session_id=session_id,  # Session ID
            model=selected_model,  # Selected model
            history=generation_failed,  # History with error context
            user_message=input,  # User input
            mode="/no_think",  # Use non-reasoning mode for error handling
            temperature=0.7,  # Fixed temperature for more consistent error messages
            top_k=20,  # Limit token sampling
            min_p=0,  # Minimum probability threshold
            top_p=0.8,  # Nucleus sampling threshold
            repetition_penalty=1  # No repetition penalty
        ):
            yield [{"role": "tool", "content": error_response}]  # Yield error response in tool role
        return  # Exit after error handling