from openai import OpenAI
import base64
import io
from dotenv import load_dotenv
import os

load_dotenv()

def _encode_image_to_base64(image_path) -> str:
    """
    Helper to read an image file from disk and encode it to base64.
    """

    buffer = io.BytesIO()
    image_path.save(buffer, format="JPEG")
    image_bytes = buffer.getvalue()
    image_string = base64.b64encode(image_bytes).decode("utf-8")
    final_string = f"data:image/jpeg;base64,{image_string}"
    return final_string

def generate_gpt_prompt(
        pos_img_path: str,
        neg_img_path: str,
        category: str,
        story_instruction: str,
        sentiment, image_elements, is_text, text_elements, non_compliant, emotion
):
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    # Prepare base64 images
    base64_pos = _encode_image_to_base64(pos_img_path)

    base64_neg = _encode_image_to_base64(neg_img_path) if neg_img_path else None

    # 1) ANALYSIS
    if base64_neg:
        # Two images
        user_content_analysis = [
            {
                "type": "text",
                "text": f"""These are the images of {category} ad category.
The first image is winning creative and the other is losing creative.
Think and Analyse these images and list down the winning and losing elements.
Make sure you give response in markdown."""
            },
            {
                "type": "image_url",
                "image_url": {"url": f"{base64_pos}"}
            },
            {
                "type": "image_url",
                "image_url": {"url": f"{base64_neg}"}
            },
        ]
    else:
        # Single image
        user_content_analysis = [
            {
                "type": "text",
                "text": f"""This is the winning image of {category} ad category.
Think and Analyse this image and list down the possible winning elements.
Make sure you give response in markdown."""
            },
            {
                "type": "image_url",
                "image_url": {"url": f"{base64_pos}"}
            }
        ]

    response_analyses = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": user_content_analysis}],
    )
    analyses = response_analyses.choices[0].message.content

    story_prompt = f"Based on this analysis, generate me the new image story for the {category} ad whose Sentiment of the new image story MUST be {sentiment}."

    if image_elements:
        story_prompt += f"You MUST creatively add these {image_elements} image elements."

    if is_text is True:
        story_prompt += f"Also, MUST add these '{text_elements}' text in new image story."
    else:
        story_prompt += "DON'T add any text overlays on the image."

    if non_compliant:
        story_prompt += f"DON'T add the mentioned stuff in the image story {non_compliant}."

    if emotion:
        story_prompt += f"The image should reflect {emotion} emotion."

    story_prompt += f"Blend everything with new image story {story_instruction}"

    # 2) STORY
    user_content_story = [
        {
            "type": "text",
            "text": f"""{story_prompt}"""
        }
    ]

    response_story = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": user_content_analysis,
            },
            {
                "role": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": analyses
                    }
                ]
            },
            {
                "role": "user",
                "content": user_content_story
            }
        ],
    )
    story = response_story.choices[0].message.content

    # 3) PROMPT
    final_user_content_prompt = [
        {
            "type": "text",
            "text": """Now convert this image description to the prompt for the image generation through ideogram model.
The prompt should be short and crisp for it to understand really well.
Please also consider the background colors, text colors, text overlays and any other elements in the image to be included in prompt.
Make sure you only return the prompt."""
        }
    ]

    response_prompt = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": user_content_analysis,
            },
            {
                "role": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": analyses
                    }
                ]
            },
            {
                "role": "user",
                "content": user_content_story
            },
            {
                "role": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": story
                    }
                ]
            },
            {
                "role": "user",
                "content": final_user_content_prompt
            }
        ],
    )
    final_prompt = response_prompt.choices[0].message.content

    return final_prompt