import os
import gc
import time
import threading
import traceback

import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image

MAX_SEED = np.iinfo(np.int32).max
LANCZOS = getattr(Image, "Resampling", Image).LANCZOS
MAX_OUTPUT_DIM = 2048

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

print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"), flush=True)
print("torch.__version__ =", torch.__version__, flush=True)
print("Using device:", device, flush=True)
print(f"CUDA device_count={torch.cuda.device_count()}, is_available={torch.cuda.is_available()}", flush=True)

# TF32 matmul: ~10-15% free speedup on Ampere/Hopper (bfloat16 accumulation paths benefit too)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
print("[startup] TF32 enabled", flush=True)

print("[startup] importing dimensions...", flush=True)
from dimensions import compute_output_dimensions, max_dim_for_mode
print("[startup] importing diffusers...", flush=True)
from diffusers import FlowMatchEulerDiscreteScheduler
print("[startup] importing QwenImageEditPlusPipeline...", flush=True)
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
print("[startup] importing QwenImageTransformer2DModel...", flush=True)
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
print("[startup] importing QwenDoubleStreamAttnProcessorFA3...", flush=True)
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
print("[startup] all imports done", flush=True)

dtype = torch.bfloat16


def _start_heartbeat(label: str) -> threading.Event:
    done = threading.Event()
    t0 = time.perf_counter()
    def _beat():
        while not done.wait(timeout=15):
            print(f"[startup] {label} still loading... ({time.perf_counter()-t0:.0f}s)", flush=True)
    threading.Thread(target=_beat, daemon=True).start()
    return done


_t0_load = time.perf_counter()
print("[startup] loading transformer from_pretrained (prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23)...", flush=True)
_hb = _start_heartbeat("transformer")
_transformer = QwenImageTransformer2DModel.from_pretrained(
    "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23",
    torch_dtype=dtype,
    device_map="cuda",
)
_hb.set()
print(f"[startup] transformer loaded in {time.perf_counter()-_t0_load:.1f}s", flush=True)

_t1_load = time.perf_counter()
print("[startup] loading pipeline from_pretrained (FireRedTeam/FireRed-Image-Edit-1.1)...", flush=True)
_hb = _start_heartbeat("pipeline")
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "FireRedTeam/FireRed-Image-Edit-1.1",
    transformer=_transformer,
    torch_dtype=dtype,
).to(device)
_hb.set()
print(f"[startup] pipeline loaded in {time.perf_counter()-_t1_load:.1f}s", flush=True)

print("[startup] using default attention processor.", flush=True)

negative_prompt = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry, deformed"


def _load_pil_list(*filepaths):
    """Turn one or more optional filepaths (from gr.Image slots) into a list of RGB PIL Images."""
    pil_images = []
    for path in filepaths:
        if not path:
            continue
        try:
            pil_images.append(Image.open(path).convert("RGB"))
        except Exception as e:
            print(f"Error loading image {path}: {e}")
    return pil_images


def update_dimensions_on_upload(image, max_dim):
    if image is None:
        return max_dim, max_dim
    w, h = image.size
    return compute_output_dimensions(w, h, max_dim)


class _InferTimer:
    def __init__(self, cuda_ok: bool) -> None:
        self._cuda_ok = cuda_ok
        self._marks: dict = {}

    def mark(self, name: str) -> None:
        ev = None
        if self._cuda_ok:
            ev = torch.cuda.Event(enable_timing=True)
            ev.record()
        self._marks[name] = (ev, time.perf_counter())

    def elapsed_ms(self, a: str, b: str) -> float:
        ev_a, t_a = self._marks[a]
        ev_b, t_b = self._marks[b]
        if ev_a and ev_b:
            return ev_a.elapsed_time(ev_b)  # true GPU-timeline ms
        return (t_b - t_a) * 1000.0

    def wall_start(self, name: str) -> float:
        return self._marks[name][1]

    def __contains__(self, name: str) -> bool:
        return name in self._marks

    def print_timings(self) -> None:
        if self._cuda_ok:
            try:
                torch.cuda.synchronize()
            except Exception:
                pass
        rows = [
            ("image_load", "load_start", "load_end"),
            ("preprocess", "pipe_start", "first_step"),
            ("inference",  "first_step", "last_step"),
            ("vae_decode", "last_step",  "pipe_end"),
        ]
        total_ms = 0.0
        lines = []
        for label, a, b in rows:
            if a in self._marks and b in self._marks:
                ms = self.elapsed_ms(a, b)
                total_ms += ms
                lines.append(f"[timing]   {label:<14} {ms:8.1f} ms")
        if "load_start" in self._marks and "pipe_end" in self._marks:
            overall_ms = self.elapsed_ms("load_start", "pipe_end")
            lines.append(f"[timing]   {'overhead':<14} {overall_ms - total_ms:8.1f} ms")
            lines.append(f"[timing]   {'── total ──':<14} {overall_ms:8.1f} ms")
        print("[timing] ─────────────────────────────────────")
        print("\n".join(lines))
        print("[timing] ─────────────────────────────────────")


def _gpu_mem_str(cuda_ok: bool, sync: bool = False) -> str:
    if not cuda_ok:
        return "CUDA not available"
    if sync:
        try:
            torch.cuda.synchronize()
        except Exception as se:
            return f"CUDA sync failed: {se}"
    alloc    = torch.cuda.memory_allocated()    / 1024**3
    reserved = torch.cuda.memory_reserved()     / 1024**3
    peak     = torch.cuda.max_memory_allocated() / 1024**3
    return f"alloc={alloc:.2f}GB reserved={reserved:.2f}GB peak={peak:.2f}GB"


def _validate_infer_inputs(pil_images: list, prompt: str) -> None:
    if not pil_images:
        raise gr.Error("Please upload at least one image to edit.")
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter an edit prompt.")


def _resolve_seed(seed: int, randomize_seed: bool) -> int:
    return random.randint(0, MAX_SEED) if randomize_seed else seed


# ── Gradio blocks ──────────────────────────────────────────────────────────────

def infer(image_1, image_2, prompt, seed, randomize_seed, guidance_scale, steps, mode, gpu_duration=20, progress=gr.Progress(track_tqdm=True)):
    # CPU-only preprocessing — GPU not yet allocated
    gc.collect()
    pil_images = _load_pil_list(image_1, image_2)
    _validate_infer_inputs(pil_images, prompt)
    seed = _resolve_seed(seed, randomize_seed)
    width, height = update_dimensions_on_upload(pil_images[0], max_dim_for_mode(mode))
    result_image, seed, _duration = _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, mode, int(gpu_duration))
    return result_image, seed


@spaces.GPU(duration=lambda *a, **kw: int(a[8]) if len(a) > 8 else 60)
def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, mode, gpu_duration=20):
    _cuda_ok = torch.cuda.is_available()
    timer = _InferTimer(_cuda_ok)
    t0 = time.perf_counter()

    print(f"[infer] ===== START =====")
    print(f"[infer] steps={steps}, guidance={guidance_scale}, seed={seed}, gpu_duration={gpu_duration}s, mode={mode}")
    print(f"[infer] prompt={repr(prompt[:120])}")

    if _cuda_ok:
        p = torch.cuda.get_device_properties(0)
        print(f"[infer] GPU: {p.name}, total={p.total_memory/1024**3:.1f}GB, cap={p.major}.{p.minor}")
        torch.cuda.reset_peak_memory_stats()

    print(f"[infer] {_gpu_mem_str(_cuda_ok)} — t={time.perf_counter()-t0:.1f}s")

    torch.cuda.empty_cache()
    print(f"[infer] cache cleared — {_gpu_mem_str(_cuda_ok)}")

    print(f"[infer] {len(pil_images)} image(s) pre-decoded, output={width}x{height}, seed={seed}")

    generator = torch.Generator(device=device).manual_seed(seed)

    _step_t = []
    def _step_cb(pipeline, step_idx, timestep, cb_kwargs):
        now = time.perf_counter()
        _step_t.append(now)
        if step_idx == 0:
            timer.mark("first_step")
        timer.mark("last_step")  # overwritten each step; final value = end of last step
        delta_ms = (now - (_step_t[-2] if len(_step_t) > 1 else t0)) * 1000
        tag = " ← includes compile" if step_idx == 0 else ""
        print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={now-t0:.1f}s")
        return cb_kwargs

    timer.mark("pipe_start")
    print(f"[infer] calling pipe... t={time.perf_counter()-t0:.1f}s")
    try:
        result_image = pipe(
            image=pil_images, prompt=prompt, negative_prompt=negative_prompt,
            height=height, width=width, num_inference_steps=steps,
            generator=generator, true_cfg_scale=guidance_scale,
            callback_on_step_end=_step_cb,
            callback_on_step_end_tensor_inputs=["latents"],
        ).images[0]
        timer.mark("pipe_end")
        print(f"[infer] VAE decode + postprocess done — {_gpu_mem_str(_cuda_ok, sync=True)} | t={time.perf_counter()-t0:.1f}s")
        timer.print_timings()
        duration = timer.elapsed_ms("pipe_start", "pipe_end") / 1000.0
        return result_image, seed, duration
    except Exception as e:
        print(f"[infer] ERROR: {type(e).__name__}: {e} | t={time.perf_counter()-t0:.1f}s")
        print(traceback.format_exc())
        try:
            torch.cuda.synchronize()
        except Exception as cuda_err:
            print(f"[infer] CUDA synchronize after error: {cuda_err}")
        timer.print_timings()
        raise
    finally:
        gc.collect()
        torch.cuda.empty_cache()
        print(f"[infer] ===== END t={time.perf_counter()-t0:.1f}s =====")


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

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("FireRed Image Edit — drop in references and a prompt, get an edit out.")
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    image_1 = gr.Image(label="Image 1", type="filepath", interactive=True)
                    image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True)

                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the edit you want...",
                )
                run_button = gr.Button("Edit", variant="primary")

                with gr.Accordion("Advanced Settings", open=False):
                    mode = gr.Dropdown(
                        label="Mode",
                        choices=["fast", "quality"],
                        value="fast",
                        info="'fast' caps the long side at 1024px; 'quality' allows up to 2048px.",
                    )

                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                    )

                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

                    guidance_scale = gr.Slider(
                        label="True guidance scale",
                        minimum=1.0,
                        maximum=10.0,
                        step=0.1,
                        value=1.0,
                    )

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

                    gpu_duration = gr.Slider(
                        label="GPU duration (seconds)",
                        minimum=10,
                        maximum=120,
                        step=5,
                        value=20,
                    )

            with gr.Column():
                result = gr.Image(label="Result", type="pil", format="png")
                seed_out = gr.Number(label="Seed used", interactive=False)

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[image_1, image_2, prompt, seed, randomize_seed, guidance_scale, steps, mode, gpu_duration],
        outputs=[result, seed_out],
    )


if __name__ == "__main__":
    demo.queue(max_size=30).launch(
        mcp_server=True,
        ssr_mode=False,
        show_error=True,
    )