import requests
import gradio as gr
import os
import json
import base64

# 📝 Form-like Configuration Section - Easily customize your script settings here!
# -----------------------------------------------------------
# Server Configuration - Change these to match your server!
fooocus_host = 'http://192.168.50.136:7878'  # 🌐 Host address for the Fooocus server

# 📁 File Paths - Change these paths if your images are stored in a different location!
imgs_base_path = "./saved_images"  # Define the base path for saving images

# Headers for HTTP Requests - Modify if needed!
headers = {
    "accept": "application/json"
}

# -----------------------------------------------------------
# 🛠️ Utility Functions - These functions perform the core operations of the script.

# Function to encode images to base64 - This helps in preparing images for API requests.
def encode_image_to_base64(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

# Function to save images from API response - This saves the processed images locally.
def save_image_from_response(response, base_filename="generated-image"):
    if not os.path.exists(imgs_base_path):
        os.makedirs(imgs_base_path)  # Create the directory if it does not exist
    output_path = os.path.join(imgs_base_path, base_filename)
    for i, item in enumerate(response, start=1):
        if item.get('url'):
            modified_url = item['url'].replace('http://127.0.0.1:7878', fooocus_host)
            img_data = requests.get(modified_url).content
            with open(f"{output_path}{i}.png", 'wb') as handler:
                handler.write(img_data)
                print(f"🖼️ Image saved as {output_path}{i}.png")

# Enhanced Gradio app for text and image prompt generation
def generate_image_from_prompt(prompt, image_file=None):
    api_url = f"{fooocus_host}/v1/generation/text-to-image"
    request_payload = {"prompt": prompt}
    if image_file:
        image_base64 = encode_image_to_base64(image_file.name)
        request_payload["image_prompt"] = image_base64  # Add image prompt if provided
    response = requests.post(api_url, json=request_payload, headers=headers)
    if response.status_code == 200:
        result = response.json()
        if isinstance(result, list) and len(result) > 0:
            save_image_from_response(result, "generated-image")
            image_url = result[0].get("url")  # Assuming the first item in the list contains the URL
            if image_url:
                modified_url = image_url.replace('http://127.0.0.1:7878', fooocus_host)
                return f"<img src='{modified_url}'/>"  # Return an HTML image tag with the modified URL
            else:
                return "Error: URL not found in response."
        elif isinstance(result, dict):
            save_image_from_response([result], "generated-image-single")  # Save single image response
            image_url = result.get("url")
            if image_url:
                modified_url = image_url.replace('http://127.0.0.1:7878', fooocus_host)
                return f"<img src='{modified_url}'/>"  # Return an HTML image tag with the modified URL
            else:
                return "Error: URL not found in response."
        else:
            return "Error: Unexpected response format. Please try again."
    else:
        return f"Error: Could not generate image. Status code: {response.status_code}"

iface = gr.Interface(
    fn=generate_image_from_prompt,
    inputs=[gr.Textbox(label="Enter your prompt"), gr.File(label="Upload an image prompt (optional)")],
    outputs=gr.HTML(label="Generated Image"),  # Adjusted to handle image URLs via HTML
    title="Enhanced Image Generator",
    description="Enter a prompt and optionally upload an image to generate an image using the API. Images are saved locally."
)

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