"""Anima Base v1.0 + Turbo (AIO GGUF) Image Generation (CPU) via sd-cli binary"""

from __future__ import annotations

import argparse
import ctypes
import gc
import os
import shutil
import subprocess
import sys
import tempfile
import threading
import time
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from PIL import Image

MODELS_DIR = os.environ.get("ANIMA_MODELS_DIR", "/tmp/anima_models")
os.makedirs(MODELS_DIR, exist_ok=True)


def _ensure_file(repo_id, filename, dest_dir):
    path = os.path.join(dest_dir, filename)
    if os.path.exists(path):
        return path
    print(f"[init] Downloading {repo_id}/{filename}...")
    from huggingface_hub import hf_hub_download

    return hf_hub_download(repo_id=repo_id, filename=filename, local_dir=dest_dir)


def _ensure_sdcli():
    """Find sd-cli in PATH (Docker provides it) or local build."""
    for candidate in [shutil.which("sd-cli"), shutil.which("sd-cli.exe")]:
        if candidate and os.path.isfile(candidate):
            return candidate
    print("[init] sd-cli not found. In Docker it's pre-installed. Locally, build stable-diffusion.cpp.")
    sys.exit(1)


SDCLI_PATH = None
aio_path = None
_active_proc = None
_active_t0 = None
_user_killed = False
_proc_lock = threading.Lock()

DEFAULT_NEGATIVE = "worst quality, low quality, score_1, score_2, score_3, artist name"
DEFAULT_RESOLUTION = "512x512"
RESOLUTIONS = ["512x512", "768x512", "512x768"]
CFG = 1.0
SAMPLER = "er_sde"
SCHEDULER = "smoothstep"
DEFAULT_STEPS = 8
TIMEOUT = 10800
EST_MIN_PER_512_8STEPS = 15

_CHILD_ENV_OVERRIDES = {
    "OPENBLAS_NUM_THREADS": "2",
    "OMP_NUM_THREADS": "2",
    "MKL_NUM_THREADS": "1",
    "NUMEXPR_NUM_THREADS": "1",
    # Best-effort guard against allocator arena growth inside the native child.
    "MALLOC_ARENA_MAX": "2",
}


def _build_cmd(prompt, negative_prompt, w, h, steps, seed, output_path):
    cmd = [
        SDCLI_PATH,
        "-m", aio_path,
        "-p", prompt,
        "-n", negative_prompt or "",
        "-W", str(w),
        "-H", str(h),
        "--steps", str(steps),
        "--cfg-scale", str(CFG),
        "--sampling-method", SAMPLER,
        "--scheduler", SCHEDULER,
        "--cache-mode", "spectrum",
        "--vae-tiling",
        "--fa",
        "--offload-to-cpu",
        "-o", output_path,
        "-v",
    ]
    if seed >= 0:
        cmd += ["-s", str(seed)]
    return cmd


def _trim_process_memory():
    """Best-effort release of free Python/native heap pages before sd-cli peaks."""
    gc.collect()
    if not sys.platform.startswith("linux"):
        return
    try:
        ctypes.CDLL("libc.so.6").malloc_trim(0)
    except (AttributeError, OSError):
        pass


def _read_temp_file(file_obj):
    file_obj.flush()
    file_obj.seek(0)
    return file_obj.read()


def _run_sdcli(cmd, timeout=TIMEOUT):
    global _active_t0, _user_killed
    env = os.environ.copy()
    env.update(_CHILD_ENV_OVERRIDES)
    _trim_process_memory()

    # Disk-backed logs avoid growing the Gradio host RSS throughout a long verbose run.
    with tempfile.TemporaryFile() as stdout_file, tempfile.TemporaryFile() as stderr_file:
        proc = subprocess.Popen(
            cmd,
            stdin=subprocess.DEVNULL,
            stdout=stdout_file,
            stderr=stderr_file,
            env=env,
            close_fds=True,
        )
        with _proc_lock:
            global _active_proc
            _active_proc = proc
            _active_t0 = time.time()
            _user_killed = False
        try:
            proc.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
            raise
        finally:
            with _proc_lock:
                _active_proc = None
                _active_t0 = None
        stdout = _read_temp_file(stdout_file)
        stderr = _read_temp_file(stderr_file)
    return proc.returncode, stdout, stderr


def generate(
    prompt: str,
    negative_prompt: str,
    resolution: str,
    steps: int,
    seed: int,
) -> tuple["Image.Image", str]:
    """Generate an anime-style image with Anima Base v1.0 + Turbo LoRA (AIO Q4_K GGUF) via sd.cpp on CPU.

    Args:
        prompt (str): Text description of the desired image (max 500 chars).
        negative_prompt (str): Concepts/styles to avoid. Empty string allowed.
        resolution (str): "WIDTHxHEIGHT". One of "512x512", "768x512", "512x768", "832x1216".
        steps (int): Sampling steps (8-16). 8 recommended for Turbo LoRA.
        seed (int): Random seed. -1 for random.

    Returns:
        tuple[PIL.Image.Image, str]: (generated image, status message).
    """
    import gradio as gr

    prompt = (prompt or "").strip()[:500]
    if not prompt:
        raise gr.Error("Please enter a prompt.")
    steps = int(steps)
    w, h = (int(x) for x in resolution.split("x"))
    pixels = w * h
    est_min = int(pixels / 512 / 512 * steps / 8 * EST_MIN_PER_512_8STEPS)

    if est_min > 180:
        raise gr.Error(f"Estimated ~{est_min} min. Too slow for CPU. Use 512x512 or fewer steps.")

    seed = int(seed) if seed is not None else -1

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

    cmd = _build_cmd(prompt, negative_prompt, w, h, steps, seed, output_path)
    print(f"[gen] {w}x{h} steps={steps} seed={seed} est~{est_min}min prompt={prompt[:80]}")
    t0 = time.time()

    try:
        rc, stdout, stderr = _run_sdcli(cmd)
        elapsed = time.time() - t0

        if rc != 0:
            if rc == -9 and _user_killed:
                raise gr.Error(f"Cancelled (user disconnected after {elapsed:.0f}s)")
            err = stderr.decode(errors="replace")[-500:] if stderr else "Unknown error"
            if rc == -9:
                raise gr.Error(f"OOM at {w}x{h}, {steps} steps, seed={seed} after {elapsed:.0f}s. Try 512x512.")
            raise gr.Error(f"sd-cli failed (code {rc}) after {elapsed:.0f}s: {err}")

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

        from PIL import Image

        with Image.open(output_path) as opened_img:
            opened_img.load()
            img = opened_img.copy()
        status = f"Generated in {elapsed:.1f}s ({w}x{h}, {steps} steps, CFG {CFG}, {SAMPLER}/{SCHEDULER})"
        print(f"[gen] {status}")
        return img, status

    except subprocess.TimeoutExpired:
        elapsed = time.time() - t0
        raise gr.Error(f"Generation timed out after {elapsed:.0f}s ({TIMEOUT//60} min limit)")
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Error: {e}")
    finally:
        try:
            os.remove(output_path)
        except FileNotFoundError:
            pass


def generate_cli(prompt, negative_prompt, resolution, steps, seed, output):
    w, h = (int(x) for x in resolution.split("x"))
    seed = int(seed)
    steps = int(steps)

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

    cmd = _build_cmd(prompt, negative_prompt, w, h, steps, seed, tmp_path)
    print(f"[cli] {w}x{h} steps={steps} seed={seed} CFG={CFG} {SAMPLER}/{SCHEDULER}")
    print(f"[cli] prompt: {prompt[:120]}")
    t0 = time.time()

    rc, stdout, stderr = _run_sdcli(cmd)
    elapsed = time.time() - t0

    if stdout:
        print(stdout.decode(errors="replace"))
    if stderr:
        print(stderr.decode(errors="replace"), file=sys.stderr)

    if rc != 0:
        print(f"[cli] sd-cli exited with code {rc}", file=sys.stderr)
        sys.exit(rc)

    if not os.path.exists(tmp_path) or os.path.getsize(tmp_path) == 0:
        print("[cli] No output image generated", file=sys.stderr)
        sys.exit(1)

    final_path = output or f"anima_output_{int(time.time())}.png"
    shutil.move(tmp_path, final_path)
    print(f"[cli] Saved to {final_path} in {elapsed:.1f}s")


def _init_models():
    global SDCLI_PATH, aio_path
    if SDCLI_PATH and aio_path:
        return
    print("[init] Ensuring sd-cli binary...")
    SDCLI_PATH = _ensure_sdcli()
    print("[init] Ensuring AIO GGUF model file...")
    t0 = time.time()
    aio_path = _ensure_file(
        "n-Arno/Anima-P3-Turbo-AIO-Q4_K", "Anima-V1-Turbo-AIO-Q4_K.gguf",
        MODELS_DIR,
    )
    print(f"[init] Model ready in {time.time()-t0:.1f}s")


def main():
    parser = argparse.ArgumentParser(description="Anima Base v1.0 + Turbo (CPU)")
    sub = parser.add_subparsers(dest="command")

    infer = sub.add_parser("infer", help="Generate an image (CLI mode)")
    infer.add_argument("-p", "--prompt", required=True, help="Text prompt")
    infer.add_argument("-n", "--negative", default=DEFAULT_NEGATIVE,
                       help="Negative prompt")
    infer.add_argument("-r", "--resolution", default=DEFAULT_RESOLUTION,
                       help="WxH resolution (default: 512x512)")
    infer.add_argument("-s", "--steps", type=int, default=DEFAULT_STEPS,
                       help=f"Sampling steps (default: {DEFAULT_STEPS})")
    infer.add_argument("--seed", type=int, default=-1, help="Seed (-1=random)")
    infer.add_argument("-o", "--output", default=None, help="Output path")

    args = parser.parse_args()

    _init_models()
    _trim_process_memory()

    if args.command == "infer":
        generate_cli(args.prompt, args.negative, args.resolution, args.steps, args.seed, args.output)
        return

    import gradio as gr

    with gr.Blocks(
        title="Anima Base v1.0 (CPU)",
        analytics_enabled=False,
    ) as demo:
        gr.Markdown(
            "**[n-Arno Anima-V1-Turbo AIO Q4_K](https://huggingface.co/n-Arno/Anima-P3-Turbo-AIO-Q4_K)** GGUF "
            "(VAE + LLM + [DiT anima-base-v1.0](https://huggingface.co/circlestone-labs/Anima) "
            "+ [Turbo LoRA](https://civitai.com/models/2560840/anima-turbo-lora) merged) "
            "via [sd.cpp](https://github.com/leejet/stable-diffusion.cpp) | "
            f"CFG {CFG}, {SAMPLER} | CPU ~15m @512x512"
        )
        with gr.Row():
            with gr.Column():
                prompt_input = gr.Textbox(label="Prompt", lines=3,
                    placeholder="masterpiece, best quality, score_9, score_8_up, 1girl, solo, cyberpunk, neon lights")
                neg_input = gr.Textbox(label="Negative Prompt", lines=2,
                    value=DEFAULT_NEGATIVE)
                with gr.Row():
                    res_input = gr.Dropdown(choices=RESOLUTIONS, value=DEFAULT_RESOLUTION, label="Resolution")
                    steps_input = gr.Slider(8, 16, value=DEFAULT_STEPS, step=1, label="Steps")
                    seed_input = gr.Number(value=-1, label="Seed (-1=random)", 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, seed_input],
            outputs=[output_img, status_box],
            concurrency_limit=1,
            api_name="generate")

        # No _on_unload: can't reliably distinguish browser vs API disconnect.
    demo.queue(default_concurrency_limit=1)
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True,
                theme="NoCrypt/miku", mcp_server=True, ssr_mode=False,
                pwa=False, enable_monitoring=False)


if __name__ == "__main__":
    main()