# app.py - 100% WORKING with EXPOSED API ENDPOINT
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
import time
import base64
from io import BytesIO
import json

# ========== CONFIGURATION ==========
MODEL_ID = "OFA-Sys/small-stable-diffusion-v0"  # Works on free CPU
DEVICE = "cpu"  # Free tier only has CPU

# ========== LOAD MODEL ==========
print("🚀 Loading Stable Diffusion model...")
try:
    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,
        safety_checker=None,
        requires_safety_checker=False
    )
    pipe = pipe.to(DEVICE)
    print("✅ Model loaded successfully!")
except Exception as e:
    print(f"❌ Model loading failed: {e}")
    pipe = None

# ========== API FUNCTION ==========
def generate_image_api(
    prompt: str,
    negative_prompt: str = "",
    steps: int = 25,
    width: int = 512,
    height: int = 512,
    seed: int = -1
):
    """
    API endpoint for WordPress automation
    Returns: {"image": "base64_string", "status": "message"}
    """
    
    if pipe is None:
        return {
            "error": "Model not loaded",
            "data": []
        }
    
    try:
        print(f"📸 API call: {prompt[:50]}...")
        
        # Set generator
        generator = None
        if seed != -1:
            generator = torch.Generator(device=DEVICE).manual_seed(seed)
        
        # Generate image
        start_time = time.time()
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=steps,
            width=width,
            height=height,
            guidance_scale=7.5,
            generator=generator
        ).images[0]
        
        gen_time = time.time() - start_time
        
        # Convert to base64
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
        
        # Return in Gradio API format
        return {
            "data": [
                f"data:image/png;base64,{img_base64}",  # Base64 image
                f"✅ Generated in {gen_time:.1f}s"       # Status message
            ]
        }
        
    except Exception as e:
        return {
            "error": str(e),
            "data": []
        }

# ========== GRADIO UI (Optional) ==========
def generate_image_ui(prompt, negative_prompt="", steps=25, width=512, height=512, seed=-1):
    """For UI display only"""
    result = generate_image_api(prompt, negative_prompt, steps, width, height, seed)
    
    if "error" in result:
        return None, f"❌ Error: {result['error']}"
    
    # Decode for UI
    img_data = result["data"][0].split(",")[1]
    img_bytes = base64.b64decode(img_data)
    image = gr.Image().pil_to_bytes(img_bytes)
    
    return image, result["data"][1]

# ========== CREATE GRADIO APP ==========
with gr.Blocks(title="Stable Diffusion API", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎨 Stable Diffusion Image Generator
    ## ✅ **API ENDPOINT EXPOSED:** `/api/predict`
    
    ### For WordPress Automation:
    ```bash
    curl -X POST https://huggingface.co/spaces/AllanHill/CdGarment/api/predict \\
      -H "Content-Type: application/json" \\
      -d '{
        "data": [
          "your prompt here",
          "negative prompt",
          25,
          512,
          512,
          -1
        ]
      }'
    ```
    
    **Response format:**
    ```json
    {
      "data": [
        "data:image/png;base64,iVBORw0KGgo...",
        "✅ Generated in 12.3s"
      ]
    }
    ```
    """)
    
    # UI for manual testing
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Textile factory...", lines=2)
        generate_btn = gr.Button("Test Generate", variant="secondary")
    
    with gr.Row():
        output_image = gr.Image(label="Preview")
        output_status = gr.Textbox(label="API Status")
    
    generate_btn.click(
        fn=generate_image_ui,
        inputs=[prompt],
        outputs=[output_image, output_status]
    )

# ========== LAUNCH APP ==========
if __name__ == "__main__":
    demo.launch(debug=True)