import gradio as gr
import requests, base64, io
from dotenv import load_dotenv
import os, datetime

load_dotenv()
API_URL = os.environ.get("API_URL", "http://0.0.0.0:8021")
APP_ENV = os.environ.get("APP_ENV", None)


def image_to_base64(image):
    if image is None:
        return None
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()


# Assuming the API returns an image URL in the response
def generate_image(brand_name, primary_color, secondary_color, description, target_audience, font, logo_image, heading_text, sub_heading_text, image_layout, custom_graphic_prompt):

    logo_base64 = image_to_base64(logo_image)

    # Define your payload/data to send to the image generation API
    data = {
        "brand_name": brand_name,
        "primary_color": primary_color,
        "secondary_color": secondary_color,
        "description": description,
        "target_audience": target_audience,
        "font": font,
        "heading_text": heading_text or "",
        "sub_heading_text": sub_heading_text or "",
        "logo_url": logo_base64,
        "user_prompt": custom_graphic_prompt,
        "image_dimension": image_layout
    }

    # Make the API call
    response = requests.post(API_URL, json=data)

    # Ensure the API call was successful
    if response.status_code != 200:
        print(f"Error {response.status_code}: {response.text}")
        return "Error: Unable to fetch image from the external API."

    image_path = f"images/temp_image_{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}.png"

    with open(image_path, "wb") as image_file:
        image_file.write(response.content)

    return image_path


iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.components.Textbox(placeholder="Brand/Company Name", label="Brand Name"),
        gr.components.ColorPicker(value="#FFFFFF", label="Primary Color"),
        gr.components.ColorPicker(value="#05A592", label="Secondary Color"),
        gr.components.Textbox(placeholder="Brief description of your brand and design guidelines", label="Description", lines=2),
        gr.components.Textbox(placeholder="Your Brand's target audience", label="Target Audience"),
        gr.components.Textbox(placeholder="Google Font Label i.e Rubik", label="Font"),
        gr.components.Image(tool="select", type="pil", label="Logo Image"), # Image upload for logo
        gr.components.Textbox(placeholder="Heading text (not required)", label="Heading Text"),
        gr.components.Textbox(placeholder="Paragraph text (not required)", label="Sub Heading Text"),
        gr.components.Dropdown(
            choices=[
                ("Select Layout", "750x1334"),
                ("Mobile Portrait (750x1334)", "750x1334"),
                ("Square (2048x2048)", "2048x2048")
            ],
            value="750x1334",
            label="Image Layout"
        ),
        gr.components.Textbox(lines=4, placeholder="Enter your custom dall-e 2-3 image prompt here (not required)", value="", label="Custom Graphic Prompt")
    ],
    outputs=[
        gr.components.Image(label="Generated Image")
    ],
    title="Slidegen AI - Image generator",
    description="Generate social media creatives from a few prompts",
    live=False
)

# iface.queue()
# Run the interface
#trigger
iface.launch()