import os
import gc
import time
import threading
import traceback

import gradio as gr
import numpy as np
import spaces
import torch
import random
import base64
import json
import html as html_lib
from io import BytesIO
from PIL import Image
from logging_utils import LogUploader

_log_uploader = LogUploader(
    token=os.environ.get("HF_TOKEN"),
    repo_id=os.environ.get("LOG_DATASET_REPO"),
    max_files=int(os.environ.get("LOG_MAX_FILES", "5000")),
    batch_interval=int(os.environ.get("LOG_BATCH_INTERVAL", "60")),
)

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)

with open("examples.json") as _f:
    EXAMPLES_CONFIG = json.load(_f)

with open("suggestions.json") as _f:
    SUGGESTIONS_CONFIG = json.load(_f)


def make_thumb_b64(path, max_dim=220):
    if not os.path.exists(path):
        return ""
    try:
        img = Image.open(path).convert("RGB")
        img.thumbnail((max_dim, max_dim), LANCZOS)
        buf = BytesIO()
        img.save(buf, format="JPEG", quality=65)
        return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
    except Exception as e:
        print(f"Thumbnail error for {path}: {e}")
        return ""


def encode_full_image(path):
    if not os.path.exists(path):
        return ""
    try:
        with open(path, "rb") as f:
            data = f.read()
        ext = path.rsplit(".", 1)[-1].lower()
        mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "webp": "image/webp"}.get(ext, "image/jpeg")
        return f"data:{mime};base64,{base64.b64encode(data).decode()}"
    except Exception as e:
        print(f"Encode error for {path}: {e}")
        return ""


def _example_thumbs_html(images):
    html = ""
    for path in images:
        thumb = make_thumb_b64(path)
        if thumb:
            html += f'<img src="{thumb}" alt="">'
        else:
            html += '<div class="example-thumb-placeholder">Preview</div>'
    return html


def _example_card_html(idx, ex):
    thumbs_html = _example_thumbs_html(ex["images"])
    n = len(ex["images"])
    badge = f'{n} image{"s" if n > 1 else ""}'
    prompt_short = html_lib.escape(ex["prompt"][:90])
    if len(ex["prompt"]) > 90:
        prompt_short += "..."
    return f'''<div class="example-card" data-idx="{idx}">
            <div class="example-thumbs">{thumbs_html}</div>
            <div class="example-meta"><span class="example-badge">{badge}</span></div>
            <div class="example-prompt-text">{prompt_short}</div>
        </div>'''


def build_example_cards_html():
    return "".join(_example_card_html(i, ex) for i, ex in enumerate(EXAMPLES_CONFIG))


def _parse_example_idx(idx_str):
    try:
        return int(float(idx_str)) if idx_str and idx_str.strip() else -1
    except (ValueError, TypeError):
        return -1


def load_example_data(idx_str):
    idx = _parse_example_idx(idx_str)
    if idx < 0 or idx >= len(EXAMPLES_CONFIG):
        return json.dumps({"images": [], "prompt": "", "names": [], "status": "error"})
    ex = EXAMPLES_CONFIG[idx]
    b64_list, names = [], []
    for path in ex["images"]:
        b64 = encode_full_image(path)
        if b64:
            b64_list.append(b64)
            names.append(os.path.basename(path))
    return json.dumps({"images": b64_list, "prompt": ex["prompt"], "names": names, "status": "ok"})


def build_suggestion_chips_html():
    chips = []
    for s in SUGGESTIONS_CONFIG:
        prompt_json = html_lib.escape(json.dumps(s["prompt"]))
        label = html_lib.escape(s["label"])
        chips.append(f'<button class="suggestion-chip" onclick="window.__setPrompt({prompt_json})">{label}</button>')
    return "".join(chips)


print("Building example thumbnails...")
EXAMPLE_CARDS_HTML = build_example_cards_html()
print(f"Built {len(EXAMPLES_CONFIG)} example cards.")
SUGGESTION_CHIPS_HTML = build_suggestion_chips_html()
print(f"Built {len(SUGGESTIONS_CONFIG)} suggestion chips.")


def b64_to_pil_list(b64_json_str):
    if not b64_json_str or b64_json_str.strip() in ("", "[]"):
        return []
    try:
        b64_list = json.loads(b64_json_str)
    except Exception:
        return []
    pil_images = []
    for b64_str in b64_list:
        if not b64_str or not isinstance(b64_str, str):
            continue
        try:
            if b64_str.startswith("data:image"):
                _, data = b64_str.split(",", 1)
            else:
                data = b64_str
            image_data = base64.b64decode(data)
            pil_images.append(Image.open(BytesIO(image_data)).convert("RGB"))
        except Exception as e:
            print(f"Error decoding image: {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


def _spawn_log(pil_images, result_image, prompt, seed, steps, guidance_scale,
               width, height, duration, success, error=""):
    threading.Thread(
        target=_log_uploader.log_inference,
        args=(pil_images, result_image, prompt, seed, steps, guidance_scale,
              width, height, duration, success, error),
        daemon=True,
    ).start()


# ── static assets ─────────────────────────────────────────────────────────────

with open("static/app.css") as _f:
    css = _f.read()

with open("static/gallery.js") as _f:
    gallery_js = _f.read()

with open("static/wire_outputs.js") as _f:
    wire_outputs_js = _f.read()

with open("static/run_preprocess.js") as _f:
    run_preprocess_js = _f.read()

with open("static/mode_toggle.js") as _f:
    mode_toggle_js = _f.read()

with open("static/negative_prompt.txt") as _f:
    negative_prompt = _f.read().strip()

# ── HTML template ──────────────────────────────────────────────────────────────

with open("templates/app.html") as _f:
    app_html = _f.read().format(
        example_cards_html=EXAMPLE_CARDS_HTML,
        suggestion_chips_html=SUGGESTION_CHIPS_HTML,
    )

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

def infer(images_b64_json, 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 = b64_to_pil_list(images_b64_json)
    _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))
    t0 = time.perf_counter()
    try:
        result_image, seed, duration = _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, mode, int(gpu_duration))
        # _spawn_log is called here (main process) so the thread survives after _infer_gpu's
        # @spaces.GPU subprocess exits — previously the daemon thread was killed on subprocess exit.
        _spawn_log(pil_images, result_image, prompt, seed, steps, guidance_scale, width, height, duration, True)
        return result_image, seed
    except Exception as e:
        duration = time.perf_counter() - t0
        _spawn_log(pil_images, None, prompt, seed, steps, guidance_scale, width, height, duration, False, str(e))
        raise


@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 =====")


with gr.Blocks() as demo:

    hidden_images_b64 = gr.Textbox(value="[]", elem_id="hidden-images-b64", elem_classes="hidden-input", container=False)
    prompt = gr.Textbox(value="", elem_id="prompt-gradio-input", elem_classes="hidden-input", container=False)
    seed = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=0, elem_id="gradio-seed", elem_classes="hidden-input", container=False)
    randomize_seed = gr.Checkbox(value=True, elem_id="gradio-randomize", elem_classes="hidden-input", container=False)
    guidance_scale = gr.Slider(minimum=1.0, maximum=10.0, step=0.1, value=1.0, elem_id="gradio-guidance", elem_classes="hidden-input", container=False)
    steps = gr.Slider(minimum=1, maximum=50, step=1, value=4, elem_id="gradio-steps", elem_classes="hidden-input", container=False)
    mode = gr.Textbox(value="fast", elem_id="gradio-mode", elem_classes="hidden-input", container=False)
    gpu_duration = gr.Slider(minimum=10, maximum=120, step=5, value=20, elem_id="gradio-gpu-duration", elem_classes="hidden-input", container=False)
    result = gr.Image(elem_id="gradio-result", elem_classes="hidden-input", container=False, format="png")

    example_idx = gr.Textbox(value="", elem_id="example-idx-input", elem_classes="hidden-input", container=False)
    example_result = gr.Textbox(value="", elem_id="example-result-data", elem_classes="hidden-input", container=False)
    example_load_btn = gr.Button("Load Example", elem_id="example-load-btn")

    gr.HTML(app_html)

    run_btn = gr.Button("Run", elem_id="gradio-run-btn")

    demo.load(fn=None, js=gallery_js)
    demo.load(fn=None, js=wire_outputs_js)
    demo.load(fn=None, js=mode_toggle_js)

    run_btn.click(
        fn=infer,
        inputs=[hidden_images_b64, prompt, seed, randomize_seed, guidance_scale, steps, mode, gpu_duration],
        outputs=[result, seed],
        js=run_preprocess_js,
    )

    example_load_btn.click(
        fn=load_example_data,
        inputs=[example_idx],
        outputs=[example_result],
        queue=False,
    )

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