# -*- coding: utf-8 -*-
"""POCKET-Zimage-CPU — FastAPI demo (CPU only, no GPU).
Serves index.html and runs stable-diffusion.cpp on CPU.
"""
import os, re, json, time, base64, shutil, subprocess, threading, uuid, urllib.request
from pathlib import Path
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

ROOT = Path(__file__).parent
MODELS = ROOT / "models"
OUT = ROOT / "outputs"
MODELS.mkdir(exist_ok=True)
OUT.mkdir(exist_ok=True)

SD_BIN = os.environ.get("SD_BIN", str(ROOT / "bin" / "sd-cli"))
THREADS = int(os.environ.get("SD_THREADS", str(min(8, os.cpu_count() or 4))))

FILES = {
    "diffusion": ("z_image_turbo-Q4_0-pocket.gguf",
                  "https://huggingface.co/FINAL-Bench/POCKET-Zimage-CPU/resolve/main/z_image_turbo-Q4_0-pocket.gguf"),
    "vae": ("ae.safetensors",
            "https://huggingface.co/Tongyi-MAI/Z-Image-Turbo/resolve/main/vae/diffusion_pytorch_model.safetensors"),
    "llm": ("Qwen3-4B-Instruct-2507-Q4_K_M.gguf",
            "https://huggingface.co/lmstudio-community/Qwen3-4B-Instruct-2507-GGUF/resolve/main/Qwen3-4B-Instruct-2507-Q4_K_M.gguf"),
}

_state = {"ready": False, "msg": "idle", "pct": 0}
_lock = threading.Lock()


HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


def _fetch(name: str, url: str):
    dst = MODELS / name
    if dst.exists() and dst.stat().st_size > 10_000_000:
        return
    tmp = dst.with_suffix(dst.suffix + ".part")
    hdr = {"User-Agent": "pocket-zimage/1.0"}
    if HF_TOKEN and "huggingface.co" in url:
        hdr["Authorization"] = "Bearer " + HF_TOKEN
    req = urllib.request.Request(url, headers=hdr)
    with urllib.request.urlopen(req) as r, open(tmp, "wb") as f:
        total = int(r.headers.get("Content-Length") or 0)
        done = 0
        while True:
            chunk = r.read(1 << 20)
            if not chunk:
                break
            f.write(chunk)
            done += len(chunk)
            if total:
                _state["pct"] = int(done * 100 / total)
                _state["msg"] = f"downloading {name} {_state['pct']}%"
    tmp.rename(dst)


def _prepare():
    try:
        for key, (name, url) in FILES.items():
            _state["msg"] = f"fetching {name}"
            _fetch(name, url)
        _state.update(ready=True, msg="ready", pct=100)
    except Exception as e:  # noqa: BLE001
        _state.update(ready=False, msg=f"error: {e}")


threading.Thread(target=_prepare, daemon=True).start()

app = FastAPI(title="POCKET-Zimage-CPU")


class GenReq(BaseModel):
    prompt: str = Field(min_length=1, max_length=600)
    steps: int = Field(default=3, ge=1, le=8)
    size: int = Field(default=512)
    seed: int = Field(default=-1)


@app.get("/api/status")
def status():
    return {**_state, "threads": THREADS, "gpu": False}


@app.post("/api/generate")
def generate(req: GenReq):
    if not _state["ready"]:
        raise HTTPException(503, _state["msg"])
    if req.size not in (384, 512, 768):
        raise HTTPException(400, "size must be 384, 512 or 768")
    if not _lock.acquire(blocking=False):
        raise HTTPException(429, "busy - one image at a time on CPU")
    try:
        tag = uuid.uuid4().hex[:10]
        dst = OUT / f"{tag}.png"
        seed = req.seed if req.seed >= 0 else int(time.time()) % 2_000_000_000
        cmd = [
            SD_BIN,
            "--diffusion-model", str(MODELS / FILES["diffusion"][0]),
            "--vae", str(MODELS / FILES["vae"][0]),
            "--llm", str(MODELS / FILES["llm"][0]),
            "-p", req.prompt,
            "--cfg-scale", "1.0",
            "--steps", str(req.steps),
            "--fa", "--vae-tiling",
            "-t", str(THREADS),
            "-H", str(req.size), "-W", str(req.size),
            "--seed", str(seed),
            "-o", str(dst),
        ]
        t0 = time.time()
        p = subprocess.run(cmd, capture_output=True, text=True, timeout=900)
        took = round(time.time() - t0, 1)
        if not dst.exists():
            tail = (p.stderr or p.stdout or "")[-400:]
            raise HTTPException(500, f"generation failed: {tail}")
        b64 = base64.b64encode(dst.read_bytes()).decode()
        dst.unlink(missing_ok=True)
        return JSONResponse({"image": f"data:image/png;base64,{b64}",
                             "seconds": took, "seed": seed,
                             "steps": req.steps, "size": req.size})
    finally:
        _lock.release()


@app.get("/")
def index():
    return FileResponse(ROOT / "index.html")


app.mount("/static", StaticFiles(directory=ROOT / "static"), name="static")