import gradio as gr
import numpy as np
import random
import json

from PIL import Image

import spaces
from http import HTTPStatus
from urllib.parse import urlparse, unquote
from pathlib import PurePosixPath
import requests
import os

from diffusers import DiffusionPipeline
import torch

model_name = "Qwen/Qwen-Image"

pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.bfloat16)
pipe.to('cuda')

MAX_SEED = np.iinfo(np.int32).max
#MAX_IMAGE_SIZE = 1440

examples = json.loads(open("examples.json").read())

aspect_ratios = {
    "FHD 1080, aspect 1:1": (1080, 1080),
    "FHD 1080, aspect 16:9": (1920, 1080),
    "FHD 1080, aspect 9:16": (1080, 1920),
    "FHD 1080, aspect 4:3": (1440, 1080),
    "FHD 1080, aspect 3:4": (1080, 1440),

    "HD 720, aspect 1:1": (720, 720),
    "HD 720, aspect 16:9": (1280, 720),
    "HD 720, aspect 9:16": (720, 1280),
    "HD 720, aspect 4:3": (960, 720),
    "HD 720, aspect 3:4": (720, 960),

    "SD 480, aspect 1:1": (480, 480),
    "SD 480, aspect 16:9": (854, 480),
    "SD 480, aspect 9:16": (480, 854),
    "SD 480, aspect 4:3": (640, 480),
    "SD 480, aspect 3:4": (480, 640),
}



def sanitize_seed(seed):
    """
    Validate and clamp a seed to int32 max. Returns 0 if invalid.

    Rules:
    - Accept int-like values (ints, numeric strings).
    - Must be an integer >= 0 and <= MAX_SEED.
    - Otherwise return 0.
    """
    # Try to coerce from strings/floats that represent integers
    try:
        # Handle strings or floats that are integer-valued
        if isinstance(seed, str):
            seed = seed.strip()
            if seed == "":
                return -1
            seed_int = int(seed, 10)
        elif isinstance(seed, (int, np.integer)):
            seed_int = int(seed)
        elif isinstance(seed, float) and seed.is_integer():
            seed_int = int(seed)
        else:
            return -1
    except (ValueError, TypeError):
        return -1

    if 0 <= seed_int <= MAX_SEED:
        return seed_int
    return -1

def polish_prompt_en(original_prompt):
    
    SYSTEM_PROMPT = open("improve_prompt.txt").read()
    
    original_prompt = original_prompt.strip()
    prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {original_prompt}\n\n Rewritten Prompt:"
    success=False
    
    while not success:
        try:
            polished_prompt = api(prompt, model='qwen-plus')
            polished_prompt = polished_prompt.strip()
            polished_prompt = polished_prompt.replace("\n", " ")
            success = True
            
        except Exception as e:
            print(f"Error during API call: {e}")
            
    return polished_prompt


@spaces.GPU(duration=90)
def infer(
    prompt,
    negative_prompt=" ",
    seed=42,
    aspect_ratio="SD 480, aspect 3:4",
    guidance_scale=4,
    num_inference_steps=50,
    progress=gr.Progress(track_tqdm=True),
):

    print(f"Generating for prompt: \n\t{prompt}\n\t{seed}\n\t{aspect_ratio}\n\t{num_inference_steps}")
  
    seed = sanitize_seed(seed)
    
    if seed == -1:
        seed = random.randint(0, MAX_SEED)

    try:
        width, height = aspect_ratios[aspect_ratio]
    except:
        width, height = (640, 480)
    
    
    
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        true_cfg_scale=guidance_scale,
        generator=torch.Generator(device="cuda").manual_seed(seed)
    ).images[0]

    return image, seed




css = """
#col-container {
    margin: 0 auto;
    max-width: 1920px;
}
"""


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

    prompt = gr.Text(
                label="Prompt",
                show_label=False,
                placeholder="Enter your prompt",
                container=False,
                render=False,
            )

    result = gr.Image(label="Result", render=False)
    seed_output = gr.Textbox(label="Used seed", lines=1, render=False)
    
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            gr.Markdown("HINT: Use smaller image size for testing, will consume less of your free GPU time!")

        with gr.Row():
            gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed_output], fn=infer, examples_per_page=25, cache_examples=False, cache_mode="lazy")
        
        with gr.Row():
            prompt.render()
            run_button = gr.Button("Generate", scale=0, variant="primary")

        result.render()
        seed_output.render()

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
                visible=True,
            )

            seed = gr.Textbox(
                lines=1,
                label="Manual seed",
                info="Manual seed, otherwise random."
            )

            with gr.Row():
                aspect_ratio = gr.Dropdown(
                    label="Image size (aprox.)",
                    choices=list(aspect_ratios.keys()),
                    value="SD 480, aspect 3:4",
                )


            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=7.5,
                    step=0.1,
                    value=4.5,
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=30, 
                )


        
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            aspect_ratio,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed_output],
    )

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