import os
import gc
import gradio as gr
import numpy as np
import torch
import json
import spaces
import config
import utils
import logging
from PIL import Image, PngImagePlugin
from datetime import datetime
from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

DESCRIPTION = "RealVis Studio"
if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU. </p>"
IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
HF_TOKEN = os.getenv("HF_TOKEN")
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", "512"))
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")

DEFAULT_NEGATIVE = (
    "(worst quality, low quality, normal quality, lowres, low details), "
    "(watermark, signature, text, logo, words, letters), blurry, jpeg artifacts"
)

MODEL = os.getenv(
    "MODEL",
    "https://huggingface.co/SG161222/RealVisXL_V4.0/blob/main/RealVisXL_V4.0.safetensors",
)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def load_pipeline(model_name):
    vae = AutoencoderKL.from_pretrained(
        "madebyollin/sdxl-vae-fp16-fix",
        torch_dtype=torch.float16,
    )
    pipeline = (
        StableDiffusionXLPipeline.from_single_file
        if MODEL.endswith(".safetensors")
        else StableDiffusionXLPipeline.from_pretrained
    )

    pipe = pipeline(
        model_name,
        vae=vae,
        torch_dtype=torch.float16,
        custom_pipeline="lpw_stable_diffusion_xl",
        use_safetensors=True,
        add_watermarker=False,
        use_auth_token=HF_TOKEN,
        variant="fp16",
    )

    pipe.to(device)
    return pipe


@spaces.GPU
def generate(
    prompt: str,
    negative_prompt: str = "",
    seed: int = 0,
    custom_width: int = 1024,
    custom_height: int = 1024,
    guidance_scale: float = 7.0,
    num_inference_steps: int = 30,
    sampler: str = "DPM++ 2M SDE Karras",
    aspect_ratio_selector: str = "1024 x 1024",
    use_upscaler: bool = False,
    upscaler_strength: float = 0.55,
    upscale_by: float = 1.5,
    progress=gr.Progress(track_tqdm=True),
) -> Image:
    generator = utils.seed_everything(seed)

    width, height = utils.aspect_ratio_handler(
        aspect_ratio_selector,
        custom_width,
        custom_height,
    )

    width, height = utils.preprocess_image_dimensions(width, height)

    backup_scheduler = pipe.scheduler
    pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)

    if use_upscaler:
        upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
    metadata = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "resolution": f"{width} x {height}",
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps,
        "seed": seed,
        "sampler": sampler,
    }

    if use_upscaler:
        new_width = int(width * upscale_by)
        new_height = int(height * upscale_by)
        metadata["use_upscaler"] = {
            "upscale_method": "nearest-exact",
            "upscaler_strength": upscaler_strength,
            "upscale_by": upscale_by,
            "new_resolution": f"{new_width} x {new_height}",
        }
    else:
        metadata["use_upscaler"] = None
    logger.info(json.dumps(metadata, indent=4))

    try:
        if use_upscaler:
            latents = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                generator=generator,
                output_type="latent",
            ).images
            upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
            images = upscaler_pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                image=upscaled_latents,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                strength=upscaler_strength,
                generator=generator,
                output_type="pil",
            ).images
        else:
            images = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                generator=generator,
                output_type="pil",
            ).images

        if images and IS_COLAB:
            for image in images:
                filepath = utils.save_image(image, metadata, OUTPUT_DIR)
                logger.info(f"Image saved as {filepath} with metadata")

        return images, metadata
    except Exception as e:
        logger.exception(f"An error occurred: {e}")
        raise
    finally:
        if use_upscaler:
            del upscaler_pipe
        pipe.scheduler = backup_scheduler
        utils.free_memory()


if torch.cuda.is_available():
    pipe = load_pipeline(MODEL)
    logger.info("Loaded on Device!")
else:
    pipe = None

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700;800&family=Fira+Code:wght@400;500&display=swap');

:root {
    --bg:      #080a0e;
    --surf:    #0d1017;
    --card:    #111520;
    --border:  #1c2133;
    --border2: #252d45;
    --amber:   #f59e0b;
    --gold:    #fbbf24;
    --cream:   #fef3c7;
    --text:    #e2e8f0;
    --muted:   #4a5578;
    --r:       14px;
    --r-sm:    8px;
}

*, *::before, *::after { box-sizing: border-box; }

body, .gradio-container {
    background: var(--bg) !important;
    font-family: 'Outfit', sans-serif !important;
    color: var(--text) !important;
}

.gradio-container::before {
    content: '';
    position: fixed; inset: 0; pointer-events: none; z-index: 0;
    background:
        radial-gradient(ellipse 70% 50% at 50% -10%, rgba(245,158,11,0.07) 0%, transparent 65%),
        radial-gradient(ellipse 40% 30% at 90% 90%, rgba(251,191,36,0.04) 0%, transparent 60%);
}

.app-hero { padding: 52px 0 28px; text-align: center; }

.app-hero h1 {
    font-size: 3rem; font-weight: 800; letter-spacing: -0.05em;
    line-height: 1; margin: 0 0 12px;
    background: linear-gradient(135deg, var(--cream) 0%, var(--gold) 40%, var(--amber) 100%);
    -webkit-background-clip: text; -webkit-text-fill-color: transparent; background-clip: text;
}

.app-hero .tagline {
    color: var(--muted); font-size: 0.88rem; font-weight: 300;
    letter-spacing: 0.06em; text-transform: uppercase; margin: 0 0 20px;
}

.app-hero .pills { display: flex; justify-content: center; gap: 8px; flex-wrap: wrap; }

.app-hero .pill {
    background: var(--card); border: 1px solid var(--border2); border-radius: 100px;
    padding: 4px 14px; font-size: 0.74rem; font-weight: 500; color: var(--muted);
    font-family: 'Fira Code', monospace;
}

.app-hero .pill.gold { color: var(--amber); border-color: rgba(245,158,11,0.3); }

.sec-label {
    font-size: 0.62rem !important; font-weight: 700 !important;
    letter-spacing: 0.15em !important; text-transform: uppercase !important;
    color: var(--amber) !important; margin: 0 0 8px !important; display: block;
}

label > span {
    font-family: 'Outfit', sans-serif !important; font-size: 0.72rem !important;
    font-weight: 500 !important; color: var(--muted) !important;
    text-transform: uppercase; letter-spacing: 0.08em;
}

textarea, input[type="text"] {
    background: var(--surf) !important; border: 1px solid var(--border) !important;
    border-radius: var(--r-sm) !important; color: var(--text) !important;
    font-family: 'Outfit', sans-serif !important; font-size: 0.95rem !important;
    transition: border-color 0.2s, box-shadow 0.2s;
}

textarea:focus, input[type="text"]:focus {
    border-color: var(--amber) !important;
    box-shadow: 0 0 0 3px rgba(245,158,11,0.12) !important;
    outline: none !important;
}

.gen-btn {
    background: linear-gradient(135deg, var(--amber), #d97706) !important;
    border: none !important; border-radius: var(--r) !important;
    color: #000 !important; font-family: 'Outfit', sans-serif !important;
    font-weight: 700 !important; font-size: 1rem !important;
    height: 54px !important; width: 100% !important;
    letter-spacing: 0.02em !important; cursor: pointer !important;
    transition: opacity 0.18s, transform 0.15s, box-shadow 0.2s !important;
    box-shadow: 0 4px 20px rgba(245,158,11,0.28) !important;
}

.gen-btn:hover {
    opacity: 0.88 !important; transform: translateY(-1px) !important;
    box-shadow: 0 8px 30px rgba(245,158,11,0.48) !important;
}

.gen-btn:active { transform: translateY(0) !important; }

.result-gallery .grid-wrap {
    background: var(--surf) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--r) !important;
}

.result-gallery img { border-radius: 10px !important; }

.gr-accordion {
    background: var(--card) !important; border: 1px solid var(--border) !important;
    border-radius: var(--r) !important; margin-top: 10px !important;
}

.aspect-radio .wrap {
    display: flex !important; flex-wrap: wrap !important; gap: 6px !important;
    background: transparent !important; padding: 0 !important; border: none !important;
}

.aspect-radio label {
    background: var(--surf) !important; border: 1px solid var(--border) !important;
    border-radius: 100px !important; padding: 5px 12px !important;
    font-size: 0.78rem !important; font-family: 'Fira Code', monospace !important;
    color: var(--muted) !important; cursor: pointer !important;
    transition: all 0.15s !important; white-space: nowrap !important;
}

.aspect-radio label:hover { border-color: var(--amber) !important; color: var(--text) !important; }

.aspect-radio label:has(input:checked) {
    background: var(--amber) !important; border-color: var(--amber) !important;
    color: #000 !important; font-weight: 600 !important;
}

.ghost-btn button {
    background: transparent !important; border: 1px solid var(--border2) !important;
    border-radius: var(--r-sm) !important; color: var(--muted) !important;
    font-family: 'Outfit', sans-serif !important; font-size: 0.82rem !important;
    height: 36px !important; transition: border-color 0.18s, color 0.18s !important;
}

.ghost-btn button:hover { border-color: var(--amber) !important; color: var(--text) !important; }

#duplicate-button {
    margin: auto;
    color: #fff;
    background: linear-gradient(135deg, var(--amber), #d97706);
    border-radius: 100vh;
    font-family: 'Outfit', sans-serif;
}

::-webkit-scrollbar { width: 5px; }
::-webkit-scrollbar-track { background: var(--surf); }
::-webkit-scrollbar-thumb { background: var(--border2); border-radius: 3px; }
::-webkit-scrollbar-thumb:hover { background: var(--amber); }
"""

with gr.Blocks(theme=gr.themes.Base(), css=CSS) as demo:

    gr.HTML("""
    <div class="app-hero">
        <h1>RealVis Studio</h1>
        <p class="tagline">Photorealistic · High-fidelity · Text-to-image</p>
        <div class="pills">
            <span class="pill gold">ZeroGPU ⚡</span>
            <span class="pill">RealVisXL V4.0</span>
            <span class="pill">SDXL · fp16</span>
            <span class="pill">safetensors</span>
        </div>
    </div>
    """)

    gr.DuplicateButton(
        value="Duplicate Space for private use",
        elem_id="duplicate-button",
        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
    )

    with gr.Row():

        with gr.Column(scale=2):
            gr.HTML('<span class="sec-label">① Prompt</span>')
            prompt = gr.Textbox(
                label="",
                lines=4,
                placeholder="high quality, portrait photo of 30 y.o woman, perfect detailed eyes, natural skin, cinematic shot...",
                container=False,
            )
            gr.HTML('<div style="height:8px"></div>')
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                lines=2,
                value=DEFAULT_NEGATIVE
            )

            gr.HTML('<div style="height:10px"></div>')
            run_btn = gr.Button("▶  Generate", variant="primary", elem_classes=["gen-btn"])

            gr.HTML('<div style="height:14px"></div>')
            gr.HTML('<span class="sec-label">② Result</span>')
            result = gr.Gallery(
                label="", columns=1, preview=True,
                container=False, elem_classes=["result-gallery"], height=600,
            )

        with gr.Column(scale=1, min_width=280):
            gr.HTML('<span class="sec-label">③ Settings</span>')
            
            use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
            with gr.Row(visible=False) as upscaler_row:
                upscaler_strength = gr.Slider(
                    label="Strength",
                    minimum=0,
                    maximum=1,
                    step=0.05,
                    value=0.55,
                )
                upscale_by = gr.Slider(
                    label="Upscale by",
                    minimum=1,
                    maximum=1.5,
                    step=0.1,
                    value=1.5,
                )

            gr.HTML('<div style="height:8px"></div>')
            sampler = gr.Dropdown(
                label="Sampler",
                choices=config.sampler_list,
                value="DPM++ 2M SDE Karras",
            )
            
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=1,
                    maximum=12,
                    step=0.1,
                    value=7.0,
                )
                num_inference_steps = gr.Slider(
                    label="Steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )

            seed = gr.Slider(
                label="Seed", minimum=0, maximum=utils.MAX_SEED, step=1, value=0
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

    with gr.Accordion("⚙  Aspect ratio", open=False):
        aspect_ratio_selector = gr.Radio(
            label="", choices=config.aspect_ratios, value="1024 x 1024",
            container=False, elem_classes=["aspect-radio"],
        )
        with gr.Group(visible=False) as custom_resolution:
            with gr.Row():
                custom_width = gr.Slider(
                    label="Width",
                    minimum=MIN_IMAGE_SIZE,
                    maximum=MAX_IMAGE_SIZE,
                    step=8,
                    value=1024,
                )
                custom_height = gr.Slider(
                    label="Height",
                    minimum=MIN_IMAGE_SIZE,
                    maximum=MAX_IMAGE_SIZE,
                    step=8,
                    value=1024,
                )

    with gr.Accordion("📋  Generation metadata", open=False):
        gr_metadata = gr.JSON(label="", show_label=False)

    gr.HTML('<div style="height:12px"></div>')
    gr.HTML('<span class="sec-label">Examples</span>')
    gr.Examples(
        examples=config.examples,
        inputs=prompt,
        outputs=[result, gr_metadata],
        fn=lambda *args, **kwargs: generate(*args, use_upscaler=False, **kwargs),
        cache_examples=CACHE_EXAMPLES,
    )

    use_upscaler.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_upscaler,
        outputs=upscaler_row,
        queue=False,
        api_name=False,
    )

    aspect_ratio_selector.change(
        fn=lambda x: gr.update(visible=x == "Custom"),
        inputs=aspect_ratio_selector,
        outputs=custom_resolution,
        queue=False,
        api_name=False,
    )

    inputs = [
        prompt,
        negative_prompt,
        seed,
        custom_width,
        custom_height,
        guidance_scale,
        num_inference_steps,
        sampler,
        aspect_ratio_selector,
        use_upscaler,
        upscaler_strength,
        upscale_by,
    ]

    prompt.submit(
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result,
        api_name="run",
    )
    negative_prompt.submit(
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result,
        api_name=False,
    )
    run_btn.click(
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=[result, gr_metadata],
        api_name=False,
    )

demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)