import base64
import io

import gradio as gr
import numpy as np
from PIL import Image
from openai import OpenAI

client = OpenAI(api_key='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')

default_prompt_template = """Replace the parameters you see in the prompt below by the material you are provided in the image, remember to be precise with the material description and colors. If you don't see the related parameter in the material just try to remove the part with the parameter, do not make examples or hypothetical values. At last return the completed prompt only:
-----
Generate a high-resolution image of a [clothing_type] that is [style_and_fit]. The [clothing_type] should be made from a fabric with [texture_qualities] and colored [color_palette]. The [key_feature] should be especially prominent, conveying the material's quality. For positioning, the garment should be displayed in a [display_style] to enhance its form and showcase the [texture_qualities]. The background should be [background_style], with lighting that highlights the [light_interaction_feature] of the material. The overall presentation should evoke a sense of [desired_emotion], suitable for a premium e-commerce experience.
-----
"""


def generate_image(prompt, material_image: np.ndarray, clothing_type, temperature=0.6):
    if prompt is None:
        prompt = default_prompt_template
    if clothing_type is None:
        clothing_type = "jacket"
    if temperature is None:
        temperature = 0.6
    if material_image is None:
        material_image = np.zeros((256, 256, 3), dtype=np.uint8)

    # Save the image to a BytesIO object
    buffer = io.BytesIO()
    Image.fromarray(material_image).save(buffer, format='PNG')

    byte_data = buffer.getvalue()

    # Encode to base64
    base64_str = base64.b64encode(byte_data).decode('utf-8')
    # Constants
    fabric_analysis = client.chat.completions.create(
        model="gpt-4-vision-preview",
        messages=[
            {
                "role": "system",
                "content": [
                    {"type": "text",
                     "text": "As professional fabric analyst, you are tasked to analyze the material in the picture "
                             "user will provide you and describe the specific details user needs to. Remember to be "
                             "precise with the material description, color palette, plot, scheme, and texture. The "
                             "texture and the scheme on the fabric should be precisely described and different colors "
                             "should be clearly mentioned and described. The goal is to provide a detailed and accurate"
                             " description of the material"
                     }
                ]
            },
            {
                "role": "user",
                "content": [
                    {"type": "text",
                     "text": prompt

                     }
                ]
            },
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": "data:image/png;base64," + base64_str},
                    {"type": "text", "text": f"Clothing type is: {clothing_type}"}
                ]
            }
        ],
        temperature=temperature
    )
    dalle_prompt = fabric_analysis.choices[0].message.content
    design_image_response = client.images.generate(
        model="dall-e-3",
        prompt=dalle_prompt,
        size="1024x1024",
        quality="standard",
        n=1,
        style="natural"
    )

    # Extract the URL of the generated design image
    design_image_url = design_image_response.data[0].url
    return [design_image_url, dalle_prompt]


# Gradio Interface

clothing_options = ["jacket", "pants", "t-shirt", "dress", "skirt", "blouse", "coat", "sweater", "suit", "men shirt",
                    "pyjamas"]

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Text(label="Prompt template", value=default_prompt_template),
        gr.Image(type="numpy", label="Upload Fabric Image"),
        gr.Dropdown(choices=clothing_options, label="Select Clothing Item"),
        gr.Slider(minimum=0.0, maximum=1.0, step=.1, label="Temperature", value=.4)
    ],
    outputs=[gr.Image(label="Generated Fashion Design"), gr.Textbox(label="Prompt")],
    title="Fashion Design Generator",
    description="Upload an image of fabric and select a clothing item to generate a fashion design."
)

if __name__ == "__main__":
    iface.launch(share=True)