import gradio as gr
import requests
import io
import os
import json
from PIL import Image

# Load Hugging Face API Key from Space Secrets
HF_API_KEY = os.getenv("HF_TOKEN")
API_URL = "https://api-inference.huggingface.co/models/sameedhayat/corgy_dog_LoRA"
headers = {"Authorization": f"Bearer {HF_API_KEY}"}

def query(payload):
    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=30)
        response.raise_for_status()
        
        # Check if the response is an image
        content_type = response.headers.get("Content-Type", "")
        if "image" in content_type:
            return response.content
        else:
            return f"API Error: Unexpected response format - {response.text}"
    except requests.exceptions.RequestException as e:
        return f"API Error: {str(e)}"

def generate_image(user_text):
    prompt = f"{user_text} in the style of TOK"
    
    try:
        image_bytes = query({"inputs": prompt})
        
        if isinstance(image_bytes, str) and "API Error" in image_bytes:
            return image_bytes  # Return error message
        
        image = Image.open(io.BytesIO(image_bytes))
        return image
    except Exception as e:
        return f"Error generating image: {str(e)}"

def fetch_showcase_images():
    """Fetch images from the local 'images' directory and use prompts from a JSON file."""
    images_dir = "images"
    json_file = os.path.join(images_dir, "prompts.json")
    showcase_images = []
    
    if os.path.exists(json_file):
        with open(json_file, "r") as f:
            image_prompts = json.load(f)
        
        for filename, prompt in image_prompts.items():
            image_path = os.path.join(images_dir, filename)
            if os.path.exists(image_path):
                showcase_images.append((image_path, prompt))
    
    return showcase_images

def create_interface():
    showcase_images = fetch_showcase_images()
    
    with gr.Blocks() as demo:
        gr.Markdown("## 🧵 EmbroideryAI - AI-Powered Embroidery Design")
        
        with gr.Row():
            user_input = gr.Textbox(label="Describe your embroidery design")
            btn = gr.Button("Generate")
        
        output_image = gr.Image(label="Generated Image")
        btn.click(generate_image, inputs=user_input, outputs=output_image)
        
        gr.Markdown("### 🎨 Showcase of AI-Generated Designs")
        for img_path, prompt in showcase_images:
            gr.Markdown(f"**Prompt:** {prompt}")
            gr.Image(value=img_path, show_label=False)
    
    demo.launch()

if __name__ == "__main__":
    create_interface()