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, text_input=None):
    # Encode the uploaded image
    base64_image = encode_image(image)

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

    # Payload with base64 encoded image as a Data URL
    payload = {
        "model": "gpt-4o-mini",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": text_input or "What's in this image?"},
                    {
                        "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 the generated message
    return response_data["choices"][0]["message"]

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

with gr.Blocks(css=css) as demo:
    gr.Markdown("WordLift Product Description Generation - [GPT-4o-mini Demo]")
    with gr.Tab(label="WordLift Product Description Generation"):
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="Input Picture")
                text_input = gr.Textbox(label="Additional Instructions (Optional)")
                submit_btn = gr.Button(value="Submit")
            with gr.Column():
                output_text = gr.Textbox(label="Output Text")

        submit_btn.click(
            generate_product_description, [input_img, text_input], [output_text]
        )

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