"""Anima 2B Image Generation (CPU) via sd-cli binary"""

import os, time, subprocess, tempfile, shutil
from pathlib import Path
from PIL import Image
from huggingface_hub import hf_hub_download
import gradio as gr

# ---------------------------------------------------------------------------
# Download models
# ---------------------------------------------------------------------------
MODELS_DIR = "/tmp/anima_models"
LORA_DIR = "/tmp/loras"
os.makedirs(MODELS_DIR, exist_ok=True)
os.makedirs(LORA_DIR, exist_ok=True)

def ensure_model(repo_id, filename, subdir=""):
    path = os.path.join(MODELS_DIR, filename)
    if os.path.exists(path):
        return path
    print(f"[init] Downloading {repo_id}/{subdir}/{filename}...")
    src = hf_hub_download(
        repo_id=repo_id,
        filename=f"{subdir}/{filename}" if subdir else filename,
    )
    shutil.copy2(src, path)
    return path

print("[init] Ensuring model files...")
t0 = time.time()
diffusion_path = ensure_model("JusteLeo/Anima2-GGUF", "anima-preview2_q4_K_M.gguf")
llm_path = ensure_model("circlestone-labs/Anima", "qwen_3_06b_base.safetensors", "split_files/text_encoders")
vae_path = ensure_model("circlestone-labs/Anima", "qwen_image_vae.safetensors", "split_files/vae")

# Turbo LoRA (8-step)
lora_src = hf_hub_download("Einhorn/Anima-Preview2-Turbo-LoRA", "anima_preview2_turbo_8step.safetensors")
lora_path = os.path.join(LORA_DIR, "anima_turbo_8step.safetensors")
if not os.path.exists(lora_path):
    shutil.copy2(lora_src, lora_path)

print(f"[init] Models ready in {time.time()-t0:.1f}s")

# ---------------------------------------------------------------------------
# Inference via sd-cli binary
# ---------------------------------------------------------------------------
RESOLUTIONS = ["512x512", "768x768", "1024x1024", "1024x768", "768x1024"]

def generate(prompt, negative_prompt, resolution, steps, cfg_scale, seed):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    w, h = (int(x) for x in resolution.split("x"))
    seed = int(seed) if int(seed) >= 0 else -1

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

    cmd = [
        "sd-cli",
        "--diffusion-model", diffusion_path,
        "--llm", llm_path,
        "--vae", vae_path,
        "--lora-model-dir", LORA_DIR,
        "-p", f"<lora:anima_turbo_8step:1.0> {prompt}",
        "-n", negative_prompt or "",
        "-W", str(w),
        "-H", str(h),
        "--steps", str(int(steps)),
        "--cfg-scale", str(float(cfg_scale)),
        "--sampling-method", "euler",
        "-o", output_path,
        "--diffusion-fa",
        "--vae-tiling",
        "-v",
    ]
    if seed >= 0:
        cmd += ["-s", str(seed)]

    print(f"[gen] {w}x{h} steps={steps} cfg={cfg_scale} seed={seed}")
    t0 = time.time()

    try:
        result = subprocess.run(
            cmd, capture_output=True, text=True, timeout=1800,
        )
        elapsed = time.time() - t0

        if result.returncode != 0:
            err = result.stderr[-500:] if result.stderr else "Unknown error"
            raise gr.Error(f"sd-cli failed (code {result.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, cfg {cfg_scale})"
        print(f"[gen] {status}")
        return img, status

    except subprocess.TimeoutExpired:
        raise gr.Error("Generation timed out (30 min limit)")
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Error: {e}")

# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Anima 2B (CPU)") as demo:
    gr.Markdown(
        "# Anima 2B Image Generation (CPU)\n"
        "Generate anime images with [Anima 2B](https://huggingface.co/circlestone-labs/Anima) "
        "+ [Turbo LoRA](https://huggingface.co/Einhorn/Anima-Preview2-Turbo-LoRA) (8 steps). "
        "Expect **~11 min at 512x512** on free CPU."
    )
    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,
                placeholder="lowres, bad anatomy, bad hands, text, error, worst quality",
                value="lowres, bad anatomy, bad hands, text, error, worst quality, blurry, censored")
            res_input = gr.Dropdown(choices=RESOLUTIONS, value="512x512", label="Resolution")
            with gr.Row():
                steps_input = gr.Slider(minimum=4, maximum=30, value=8, step=1, label="Steps")
                cfg_input = gr.Slider(minimum=1.0, maximum=10.0, value=1.0, step=0.5, label="CFG Scale")
                seed_input = gr.Number(value=-1, label="Seed", precision=0)
            gen_btn = gr.Button("Generate", 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, steps_input, cfg_input, seed_input],
        outputs=[output_img, status_box])

    gr.Markdown("---\nAnima 2B Q4_K_M GGUF + Turbo LoRA (8 steps) | "
                "[Model](https://huggingface.co/circlestone-labs/Anima) | "
                "[sd.cpp](https://github.com/leejet/stable-diffusion.cpp)")

demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, theme="NoCrypt/miku")