import gradio as gr
import numpy as np
import random
import os
import base64
import requests
import io
from PIL import Image, ImageOps
import pillow_heif  # For HEIF/AVIF support

# --- Constants ---
MAX_SEED = np.iinfo(np.int32).max
API_URL = "https://router.huggingface.co/fal-ai/fal-ai/flux-kontext/dev?_subdomain=queue"

def get_headers():
    """Get headers for API requests"""
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise gr.Error("HF_TOKEN environment variable not found. Please add your Hugging Face token to the Space settings.")
    
    return {
        "Authorization": f"Bearer {hf_token}",
        "X-HF-Bill-To": "huggingface"
    }

def query_api(payload):
    """Send request to the API and return response"""
    headers = get_headers()
    response = requests.post(API_URL, headers=headers, json=payload)
    
    if response.status_code != 200:
        raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
    
    return response.content

# --- Core Inference Function for ChatInterface ---
def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress()):
    """
    Performs image generation or editing based on user input from the chat interface.
    """
    # Register HEIF opener with PIL for AVIF/HEIF support
    pillow_heif.register_heif_opener()
    
    prompt = message["text"]
    files = message["files"]

    if not prompt and not files:
        raise gr.Error("Please provide a prompt and/or upload an image.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Prepare the payload
    payload = {
        "parameters": {
            "prompt": prompt,
            "seed": seed,
            "guidance_scale": guidance_scale,
            "num_inference_steps": steps
        }
    }

    if files:
        print(f"Received image: {files[0]}")
        try:
            # Try to open and convert the image
            input_image = Image.open(files[0])
            # Convert to RGB if needed (handles RGBA, P, etc.)
            if input_image.mode != "RGB":
                input_image = input_image.convert("RGB")
            # Auto-orient the image based on EXIF data
            input_image = ImageOps.exif_transpose(input_image)
            
            # Convert PIL image to base64 for the API
            img_byte_arr = io.BytesIO()
            input_image.save(img_byte_arr, format='PNG')
            img_byte_arr.seek(0)
            image_base64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
            
            # Add image to payload for image-to-image
            payload["inputs"] = image_base64
            
        except Exception as e:
            raise gr.Error(f"Could not process the uploaded image: {str(e)}. Please try uploading a different image format (JPEG, PNG, WebP).")
            
        progress(0.1, desc="Processing image...")
    else:
        print(f"Received prompt for text-to-image: {prompt}")
        # For text-to-image, we don't need the inputs field
        progress(0.1, desc="Generating image...")

    try:
        # Make API request
        image_bytes = query_api(payload)
        
        # Convert response bytes to PIL Image
        image = Image.open(io.BytesIO(image_bytes))
        
        progress(1.0, desc="Complete!")
        return gr.Image(value=image)
        
    except Exception as e:
        raise gr.Error(f"Failed to generate image: {str(e)}")

# --- UI Definition using gr.ChatInterface ---

seed_slider = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
randomize_checkbox = gr.Checkbox(label="Randomize seed", value=False)
guidance_slider = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=2.5)
steps_slider = gr.Slider(label="Steps", minimum=1, maximum=30, value=28, step=1)

demo = gr.ChatInterface(
    fn=chat_fn,
    title="FLUX.1 Kontext [dev] - Direct API",
    description="""<p style='text-align: center;'>
    A simple chat UI for the <b>FLUX.1 Kontext</b> model using direct API calls with requests.
    <br>
    To edit an image, upload it and type your instructions (e.g., "Add a hat").
    <br>
    To generate an image, just type a prompt (e.g., "A photo of an astronaut on a horse").
    <br>
    Find the model on <a href='https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev' target='_blank'>Hugging Face</a>.
    </p>""",
    multimodal=True,
    textbox=gr.MultimodalTextbox(
        file_types=["image"],
        placeholder="Type a prompt and/or upload an image...",
        render=False
    ),
    additional_inputs=[
        seed_slider,
        randomize_checkbox,
        guidance_slider,
        steps_slider
    ],
    theme="soft"
)

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