import gc
import math
import os
import random
import shutil
import time
from threading import Lock

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1")
os.environ.setdefault("device", "cuda:0")

import gradio as gr
import spaces
import torch
from huggingface_hub import snapshot_download
from PIL import ImageOps

from boogu.models.transformers.transformer_boogu import BooguImageTransformer2DModel
from boogu.pipelines.boogu.pipeline_boogu import BooguImagePipeline
from boogu.pipelines.boogu.pipeline_boogu_turbo import BooguImageTurboPipeline


MAX_SEED = 2**32 - 1
DEVICE = "cuda:0"

MODEL_CONFIGS = {
    "Turbo": {
        "repo": "Boogu/Boogu-Image-0.1-Turbo-fp8",
        "pipeline": "turbo",
        "steps": 4,
        "min_steps": 1,
        "max_steps": 8,
        "text_guidance": 1.0,
        "image_guidance": 1.0,
        "height": 1024,
        "width": 1024,
        "max_images": 4,
        "description": "Fast text-to-image. Uses Boogu's 4-step DMD Turbo path.",
    },
    "Base": {
        "repo": "Boogu/Boogu-Image-0.1-Base-fp8",
        "pipeline": "base",
        "steps": 30,
        "min_steps": 10,
        "max_steps": 50,
        "text_guidance": 4.0,
        "image_guidance": 1.0,
        "height": 1024,
        "width": 1024,
        "max_images": 2,
        "description": "Quality text-to-image foundation model.",
    },
    "Edit": {
        "repo": "Boogu/Boogu-Image-0.1-Edit-fp8",
        "pipeline": "edit",
        "steps": 30,
        "min_steps": 10,
        "max_steps": 50,
        "text_guidance": 4.0,
        "image_guidance": 1.0,
        "height": 1024,
        "width": 1024,
        "max_images": 2,
        "description": "Instruction-based image editing with one input image.",
    },
}

current_pipe = None
current_mode = None
pipe_lock = Lock()


def normalize_seed(seed):
    value = int(seed) % MAX_SEED
    if value < 0:
        value += MAX_SEED
    return value


def cleanup_pipe():
    global current_pipe, current_mode
    if current_pipe is not None:
        try:
            current_pipe.to("cpu")
        except Exception:
            pass
    current_pipe = None
    current_mode = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def load_fp8_transformer(repo_id, dtype):
    repo_dir = prepare_repo_dir(repo_id)
    return BooguImageTransformer2DModel.from_pretrained(
        os.path.join(repo_dir, "transformer"),
        torch_dtype=dtype,
        use_safetensors=False,
    ), repo_dir


def prepare_repo_dir(repo_id):
    repo_dir = snapshot_download(
        repo_id,
        token=os.environ.get("HF_TOKEN") or None,
    )
    transformer_dir = os.path.join(repo_dir, "transformer")
    source = os.path.join(transformer_dir, "transformer_boogu.py")
    compat = os.path.join(
        transformer_dir,
        "boogu.models.transformers.transformer_boogu.py",
    )
    if os.path.exists(source) and not os.path.exists(compat):
        shutil.copyfile(source, compat)
    return repo_dir


def get_pipe(mode):
    global current_pipe, current_mode

    if current_pipe is not None and current_mode == mode:
        return current_pipe

    with pipe_lock:
        if current_pipe is not None and current_mode == mode:
            return current_pipe

        cleanup_pipe()
        cfg = MODEL_CONFIGS[mode]
        repo_id = cfg["repo"]
        dtype = torch.bfloat16
        pipeline_cls = (
            BooguImageTurboPipeline
            if cfg["pipeline"] == "turbo"
            else BooguImagePipeline
        )

        print(f"Loading {repo_id}...", flush=True)
        started = time.perf_counter()
        transformer, repo_dir = load_fp8_transformer(repo_id, dtype)
        pipe = pipeline_cls.from_pretrained(
            repo_dir,
            torch_dtype=dtype,
            trust_remote_code=True,
            transformer=transformer,
        )

        # The official Boogu guide recommends CPU offload for constrained GPUs.
        pipe.enable_model_cpu_offload_flag = True
        pipe.enable_model_cpu_offload(device=DEVICE)

        current_pipe = pipe
        current_mode = mode
        print(f"Loaded {repo_id} in {time.perf_counter() - started:.1f}s.", flush=True)
        return current_pipe


def build_prompt(prompt):
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a prompt.")
    return prompt


def build_negative_prompt(negative_prompt):
    return (negative_prompt or "").strip()


def aligned_pixels(height, width):
    h = max(256, int(height))
    w = max(256, int(width))
    h = int(round(h / 16) * 16)
    w = int(round(w / 16) * 16)
    return h, w


def gpu_duration(
    mode,
    prompt,
    negative_prompt,
    input_image,
    preserve_input_size,
    height,
    width,
    num_inference_steps,
    text_guidance_scale,
    image_guidance_scale,
    image_count,
    seed,
    progress=None,
):
    h, w = aligned_pixels(height, width)
    megapixels = max(1.0, (h * w) / (1024 * 1024))
    count = max(1, int(image_count))
    steps = max(1, int(num_inference_steps))
    load_budget = 420
    if mode == "Turbo":
        run_budget = 35 * megapixels * count + 12 * steps * count
    else:
        run_budget = 70 * megapixels * count + 10 * steps * count
    return min(1800, max(180, int(math.ceil(load_budget + run_budget))))


@spaces.GPU(duration=gpu_duration, size="xlarge")
def gpu_generate(
    mode,
    prompt,
    negative_prompt,
    input_image,
    preserve_input_size,
    height,
    width,
    num_inference_steps,
    text_guidance_scale,
    image_guidance_scale,
    image_count,
    seed,
    progress=gr.Progress(track_tqdm=True),
):
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available inside the ZeroGPU worker.")

    cfg = MODEL_CONFIGS[mode]
    pipe = get_pipe(mode)
    instruction = build_prompt(prompt)
    negative_instruction = build_negative_prompt(negative_prompt)
    h, w = aligned_pixels(height, width)
    count = max(1, int(image_count))
    seeds = [(normalize_seed(seed) + i) % MAX_SEED for i in range(count)]

    if mode == "Edit":
        if input_image is None:
            raise gr.Error("Edit mode needs an input image.")
        input_image = ImageOps.exif_transpose(input_image).convert("RGB")

    images = []
    for index, current_seed in enumerate(seeds, start=1):
        progress(0.0, desc=f"{mode}: generating {index}/{len(seeds)}")
        generator = torch.Generator(device=DEVICE).manual_seed(int(current_seed))

        if mode == "Turbo":
            result = pipe(
                instruction=[instruction],
                negative_instruction="",
                empty_instruction="",
                height=h,
                width=w,
                num_inference_steps=int(num_inference_steps),
                text_guidance_scale=1.0,
                image_guidance_scale=1.0,
                empty_instruction_guidance_scale=0.0,
                use_dmd_student_inference=True,
                dmd_conditioning_sigma=0.001,
                generator=generator,
                device=DEVICE,
            )
        elif mode == "Edit":
            edit_height = None if preserve_input_size else h
            edit_width = None if preserve_input_size else w
            max_pixels = 2048 * 2048 if preserve_input_size else h * w
            max_side = 2048 * 2 if preserve_input_size else 2 * max(h, w)
            result = pipe(
                instruction=[instruction],
                input_images=[[input_image]],
                input_image_paths=None,
                negative_instruction=negative_instruction,
                height=edit_height,
                width=edit_width,
                max_input_image_pixels=max_pixels,
                max_input_image_side_length=max_side,
                align_res=bool(preserve_input_size),
                num_inference_steps=int(num_inference_steps),
                text_guidance_scale=float(text_guidance_scale),
                image_guidance_scale=float(image_guidance_scale),
                generator=generator,
                device=DEVICE,
            )
        else:
            result = pipe(
                instruction=instruction,
                negative_instruction=negative_instruction,
                height=h,
                width=w,
                max_input_image_pixels=h * w,
                max_input_image_side_length=2 * max(h, w),
                num_inference_steps=int(num_inference_steps),
                text_guidance_scale=float(text_guidance_scale),
                image_guidance_scale=1.0,
                generator=generator,
                device=DEVICE,
            )

        images.append(result.images[0])

    run_info = {
        "mode": mode,
        "model": cfg["repo"],
        "seeds": seeds,
        "steps": int(num_inference_steps),
        "height": None if mode == "Edit" and preserve_input_size else h,
        "width": None if mode == "Edit" and preserve_input_size else w,
        "text_guidance_scale": 1.0 if mode == "Turbo" else float(text_guidance_scale),
        "image_guidance_scale": 1.0 if mode != "Edit" else float(image_guidance_scale),
        "fp8": True,
    }
    return images, ", ".join(str(s) for s in seeds), run_info


def generate_image(
    mode,
    prompt,
    negative_prompt,
    input_image,
    preserve_input_size,
    height,
    width,
    num_inference_steps,
    text_guidance_scale,
    image_guidance_scale,
    image_count,
    seed,
    randomize_seed,
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return gpu_generate(
        mode,
        prompt,
        negative_prompt,
        input_image,
        bool(preserve_input_size),
        int(height),
        int(width),
        int(num_inference_steps),
        float(text_guidance_scale),
        float(image_guidance_scale),
        int(image_count),
        normalize_seed(seed),
    )


def sync_mode_defaults(mode):
    cfg = MODEL_CONFIGS[mode]
    is_edit = mode == "Edit"
    is_turbo = mode == "Turbo"
    return (
        gr.update(visible=is_edit),
        gr.update(visible=is_edit, value=is_edit),
        gr.update(value=cfg["height"]),
        gr.update(value=cfg["width"]),
        gr.update(
            minimum=cfg["min_steps"],
            maximum=cfg["max_steps"],
            value=cfg["steps"],
            info="Turbo uses 4 steps. Base/Edit usually use 25-50 steps.",
        ),
        gr.update(value=cfg["text_guidance"], interactive=not is_turbo),
        gr.update(value=cfg["image_guidance"], visible=is_edit),
        gr.update(visible=not is_turbo),
        gr.update(maximum=cfg["max_images"], value=1),
        gr.update(value=cfg["description"]),
    )


examples = [
    "A cinematic product photo of a translucent orange headphone on a matte black table, softbox lighting, shallow depth of field, premium advertising style",
    "A Chinese New Year poster, exact text \"BOOGU IMAGE\", red paper-cut design, gold foil typography, crisp bilingual layout",
    "A realistic editorial photo of a ceramic artist in a sunlit workshop, clay dust, handmade bowls, natural skin texture",
    "A clean app icon for a weather app, exact text \"SKYCAST\", glassmorphism, blue and yellow palette, sharp edges",
]


with gr.Blocks(title="Boogu Image 0.1 Demo") as demo:
    gr.Markdown(
        """
        # Boogu Image 0.1 Demo

        Generate and edit images with [Boogu-Image-0.1](https://huggingface.co/Boogu) on ZeroGPU.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            mode = gr.Radio(
                choices=["Turbo", "Base", "Edit"],
                value="Turbo",
                label="Model Mode",
            )
            mode_note = gr.Textbox(
                value=MODEL_CONFIGS["Turbo"]["description"],
                label="Mode",
                interactive=False,
            )

            input_image = gr.Image(
                label="Input Image",
                type="pil",
                image_mode="RGB",
                sources=["upload", "clipboard"],
                visible=False,
            )
            preserve_input_size = gr.Checkbox(
                label="Preserve input image aspect/size",
                value=False,
                visible=False,
            )

            prompt = gr.Textbox(
                label="Prompt / Edit Instruction",
                placeholder="Describe the image, or tell Edit mode what to change...",
                lines=4,
            )

            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="Things you want the model to avoid...",
                lines=3,
                visible=False,
            )

            with gr.Row():
                height = gr.Slider(
                    minimum=512,
                    maximum=2048,
                    value=1024,
                    step=64,
                    label="Height",
                )
                width = gr.Slider(
                    minimum=512,
                    maximum=2048,
                    value=1024,
                    step=64,
                    label="Width",
                )

            with gr.Row():
                num_inference_steps = gr.Slider(
                    minimum=1,
                    maximum=50,
                    value=4,
                    step=1,
                    label="Inference Steps",
                    info="Turbo uses 4 steps. Base/Edit usually use 25-50 steps.",
                )
                image_count = gr.Slider(
                    minimum=1,
                    maximum=4,
                    value=1,
                    step=1,
                    label="Images",
                )

            text_guidance_scale = gr.Slider(
                minimum=1.0,
                maximum=8.0,
                value=1.0,
                step=0.1,
                label="Text Guidance Scale",
                info="For Boogu, 1.0 disables CFG. Turbo keeps this locked at 1.0.",
                interactive=False,
            )

            image_guidance_scale = gr.Slider(
                minimum=1.0,
                maximum=5.0,
                value=1.0,
                step=0.1,
                label="Image Guidance Scale",
                info="Used by Edit mode for the reference image.",
                visible=False,
            )

            with gr.Row():
                seed = gr.Number(label="Seed", value=42, precision=0)
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)

            generate_btn = gr.Button("Generate", variant="primary", size="lg")

        with gr.Column(scale=1):
            output_images = gr.Gallery(
                label="Generated Images",
                columns=2,
                rows=2,
                preview=True,
            )
            used_seeds = gr.Textbox(label="Seeds Used", interactive=False)
            run_info = gr.JSON(label="Run Info")

    gr.Markdown("### Example Prompts")
    gr.Examples(examples=examples, inputs=[prompt], cache_examples=False)

    gr.Markdown(
        "Models by Boogu. This demo uses the official FP8 checkpoints for ZeroGPU compatibility."
    )

    inputs = [
        mode,
        prompt,
        negative_prompt,
        input_image,
        preserve_input_size,
        height,
        width,
        num_inference_steps,
        text_guidance_scale,
        image_guidance_scale,
        image_count,
        seed,
        randomize_seed,
    ]
    outputs = [output_images, used_seeds, run_info]

    mode.change(
        fn=sync_mode_defaults,
        inputs=[mode],
        outputs=[
            input_image,
            preserve_input_size,
            height,
            width,
            num_inference_steps,
            text_guidance_scale,
            image_guidance_scale,
            negative_prompt,
            image_count,
            mode_note,
        ],
    )

    generate_btn.click(fn=generate_image, inputs=inputs, outputs=outputs)
    prompt.submit(fn=generate_image, inputs=inputs, outputs=outputs)


if __name__ == "__main__":
    demo.launch(mcp_server=True, show_error=True)