"""Z-Anime 6B Image Generation (CPU) via sd-cli binary"""

import os, time, subprocess, tempfile, threading
from PIL import Image
import gradio as gr

# ---------------------------------------------------------------------------
# Model paths (downloaded at build time)
# ---------------------------------------------------------------------------
DIFFUSION = "/app/models/z-anime-8steps-q5_0.gguf"
LLM = "/app/models/qwen3_4b_q8_0.gguf"
VAE = "/app/models/ae.safetensors"

RESOLUTIONS = ["512x512", "768x512", "512x768"]
STEPS = 8
CFG = 1.0
TIMEOUT = 10800

_active_proc = None
_proc_lock = threading.Lock()

# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def generate(prompt, negative_prompt, resolution, seed):
    global _active_proc

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

    prompt = prompt.strip()[:500]
    w, h = (int(x) for x in resolution.split("x"))
    seed = int(seed or -1) if seed is not None else -1

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
        output_path = f.name

    cmd = [
        "/app/sd-cli",
        "--diffusion-model", DIFFUSION,
        "--llm", LLM,
        "--vae", VAE,
        "-p", prompt,
        "-n", negative_prompt or "",
        "-W", str(w),
        "-H", str(h),
        "--steps", str(STEPS),
        "--cfg-scale", str(CFG),
        "--sampling-method", "euler_a",
        "--schedule", "beta",
        "-o", output_path,
        "--diffusion-fa",
        "--vae-tiling",
        "-v",
    ]
    if seed >= 0:
        cmd += ["-s", str(seed)]

    print(f"[gen] {w}x{h} steps={STEPS} seed={seed} prompt={prompt[:80]}")
    t0 = time.time()

    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        with _proc_lock:
            _active_proc = proc

        try:
            stdout, stderr = proc.communicate(timeout=TIMEOUT)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
            raise

        elapsed = time.time() - t0

        with _proc_lock:
            _active_proc = None

        if proc.returncode != 0:
            err = stderr.decode(errors="replace")[-500:] if stderr else "Unknown error"
            if proc.returncode == -9:
                raise gr.Error("Out of memory (killed by OS). Try 512x512.")
            raise gr.Error(f"sd-cli failed (code {proc.returncode}): {err}")

        if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:
            raise gr.Error("No output image generated")

        img = Image.open(output_path)
        status = f"Generated in {elapsed:.1f}s ({w}x{h}, {STEPS} steps)"
        print(f"[gen] {status}")
        return img, status

    except subprocess.TimeoutExpired:
        with _proc_lock:
            _active_proc = None
        raise gr.Error(f"Generation timed out ({TIMEOUT//60} min limit)")
    except gr.Error:
        with _proc_lock:
            _active_proc = None
        raise
    except Exception as e:
        with _proc_lock:
            _active_proc = None
        raise gr.Error(f"Error: {e}")

# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Z-Anime (CPU)") as demo:
    gr.Markdown(
        "**[Z-Anime 6B](https://huggingface.co/SeeSee21/Z-Anime)** S3-DiT Q5_0 GGUF "
        "(distill 8-step) via [sd.cpp](https://github.com/leejet/stable-diffusion.cpp) | "
        "Free CPU inference"
    )
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Prompt", lines=3,
                placeholder="anime girl with silver hair, fantasy armor, dramatic lighting")
            neg_input = gr.Textbox(label="Negative Prompt", lines=2,
                value="lowres, bad anatomy, bad hands, text, error, worst quality, blurry")
            with gr.Row():
                res_input = gr.Dropdown(choices=RESOLUTIONS, value="512x512",
                    label="Resolution")
                seed_input = gr.Number(value=-1, label="Seed (-1=random)", precision=0)
            gen_btn = gr.Button("Generate (8 steps, CFG 1)", variant="primary", size="lg")
        with gr.Column():
            output_img = gr.Image(type="pil", label="Output")
            status_box = gr.Textbox(label="Status", interactive=False)

    gen_btn.click(fn=generate,
        inputs=[prompt_input, neg_input, res_input, seed_input],
        outputs=[output_img, status_box],
        concurrency_limit=1)

    def _on_unload():
        with _proc_lock:
            proc = _active_proc
        if proc and proc.poll() is None:
            print("[cleanup] User disconnected, killing sd-cli process")
            proc.kill()

    demo.unload(_on_unload)

demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)