import base64
import requests
import gradio as gr
from PIL import Image
import numpy as np
from datetime import datetime
import os

# OpenAI API Key
api_key = os.getenv("OPENAI_API_KEY")


# Function to encode the image
def encode_image(image_array):
    # Convert numpy array to an image file and encode it in base64
    img = Image.fromarray(np.uint8(image_array))
    img_buffer = os.path.join(
        "/tmp", f"temp_image_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
    )
    img.save(img_buffer, format="JPEG")

    with open(img_buffer, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


# Function to generate product description using OpenAI API
def generate_product_description(image, description_type, custom_instruction=None):
    # Encode the uploaded image
    base64_image = encode_image(image)

    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}

    # Set the description type or custom instruction
    description_prompts = {
        "Short Formal": "Create a compelling and succinct product description from the image.",
        "Bullet Points": "Provide a detailed product description in bullet points based on the image.",
        "Amazon Optimized": "Write an Amazon-style product description that includes key features, benefits, and a call to action.",
        "Fashion": "Generate a stylish and trendy product description suitable for a fashion item based on the image.",
        "Sport": "Create an energetic and engaging product description for a sports-related item based on the image.",
    }

    if description_type == "Other" and custom_instruction:
        instruction = custom_instruction
    else:
        instruction = description_prompts.get(
            description_type, "Create a product description based on the image."
        )

    # Payload with base64 encoded image as a Data URL
    payload = {
        "model": "gpt-4o-mini",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": instruction},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                    },
                ],
            }
        ],
        "max_tokens": 300,
    }

    response = requests.post(
        "https://api.openai.com/v1/chat/completions", headers=headers, json=payload
    )
    response_data = response.json()

    # Handle errors
    if response.status_code != 200:
        raise ValueError(
            f"OpenAI API Error: {response_data.get('error', {}).get('message', 'Unknown Error')}"
        )

    # Extract and return only the generated message content
    return response_data["choices"][0]["message"]["content"]


css = """
  #output {
    height: 500px;
    overflow: auto;
    border: 1px solid #ccc;
  }
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("WordLift Product Description Generation - [FREE]")
    with gr.Tab(label="WordLift Product Description Generation"):
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="Input Picture")
                description_type = gr.Dropdown(
                    label="Select Description Type",
                    choices=[
                        "Short Formal",
                        "Bullet Points",
                        "Amazon Optimized",
                        "Fashion",
                        "Sport",
                        "Other",
                    ],
                    value="Short Formal",
                )
                custom_instruction = gr.Textbox(
                    label="Custom Instruction (Only for 'Other')", visible=False
                )
                submit_btn = gr.Button(value="Submit")
            with gr.Column():
                output_text = gr.Textbox(label="Output Text")

        # Toggle visibility of custom instruction based on selected type
        def toggle_custom_instruction(type_selection):
            return gr.update(visible=(type_selection == "Other"))

        description_type.change(
            toggle_custom_instruction,
            inputs=[description_type],
            outputs=[custom_instruction],
        )

        submit_btn.click(
            generate_product_description,
            [input_img, description_type, custom_instruction],
            [output_text],
        )

# Launch Gradio app
demo.queue(api_open=False)
demo.launch(debug=True)