import os
import gradio as gr
import numpy as np
import random
import spaces
import torch
import tempfile
import os
import requests
import time
from io import BytesIO

from PIL import Image

token = os.environ["API_TOKEN"]
model = "black-forest-labs/FLUX.1-schnell"
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

custom_css = """
.built-with {
    display: none !important;
}
.show-api {
    display: none !important;
}
"""

def text_to_image(prompt, height, width, seed, num_inference_steps,
                  model=model,
                  token=token):
    """
    Generate an image from a text prompt using the Hugging Face Inference API.
    Returns a PIL Image (JPEG) if successful, otherwise returns None after retrying.

    Parameters:
        prompt (str): The text prompt describing the image.
        height (int): Height of the generated image.
        width (int): Width of the generated image.
        seed (int): Random seed for generation reproducibility.
        num_inference_steps (int): Number of inference steps.
        model (str): Hugging Face model identifier.
        token (str): Hugging Face API token.

    Returns:
        PIL.Image.Image or None: The generated image as a PIL Image object or None on failure.
    """
    api_url = f"https://api-inference.huggingface.co/models/{model}"
    headers = {
        "Authorization": f"Bearer {token}"
    }
    payload = {
        "inputs": prompt,
        "parameters": {
            "height": height,
            "width": width,
            "seed": seed,
            "num_inference_steps": num_inference_steps
        }
    }
    for attempt in range(1, 4):
        response = requests.post(api_url, headers=headers, json=payload)
        if response.status_code == 200:
            try:
                image = Image.open(BytesIO(response.content))
                if image.format != 'JPEG':
                    with BytesIO() as output:
                        image.convert("RGB").save(output, format="JPEG")
                        output.seek(0)
                        image = Image.open(output)
                return image
            except Exception as e:
                print(f"Error processing the image data: {e}")
        else:
            print(f"Attempt {attempt}: Request failed with status code {response.status_code}")
        
        if attempt == 1:
            print("Waiting for 3 seconds before retrying...")
            time.sleep(3)
        elif attempt == 2:
            print("Waiting for 5 seconds before retrying...")
            time.sleep(5)
    return None

def infer(prompt, seed=42, randomize_seed=True, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    image = text_to_image(
        prompt,
        height=height,
        width=width,
        seed=seed,
        num_inference_steps=4
    )
    temp_dir = tempfile.gettempdir()
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg", dir=temp_dir)
    image.save(temp_file, format="JPEG")
    temp_file_path = temp_file.name
    temp_file.close()
    return temp_file_path, seed

examples = [
    "A girl and a boy dancing in the forest",
    "Tiny cat in a space suite in the moon",
    "an anime illustration of girl holding book in her hand in a library",
]

with gr.Blocks(css=custom_css) as demo:

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
            with gr.Row():
                
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            with gr.Row():
                
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result, seed],
            cache_examples="lazy"
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps],
        outputs = [result, seed]
    )

demo.launch()#show_api=False)