import os
import subprocess
import time
import urllib.request
import zipfile

import gradio as gr
import uvicorn
from fastapi import FastAPI
from huggingface_hub import hf_hub_download

# --- Config -----------------------------------------------------------------
SD_TAG = "master-709-92a3b73"
SD_ZIP = "sd-master-92a3b73-bin-Linux-Ubuntu-24.04-x86_64.zip"
SD_URL = f"https://github.com/leejet/stable-diffusion.cpp/releases/download/{SD_TAG}/{SD_ZIP}"

MODEL_REPO = "Green-Sky/SD-Turbo-GGUF"
MODEL_FILE = "sd_turbo-f16-q8_0.gguf"

N_THREADS = int(os.environ.get("N_THREADS", "2"))
WORK = os.path.abspath("runtime")
os.makedirs(WORK, exist_ok=True)


def log(*a):
    print("[startup]", *a, flush=True)


# --- Fetch sd.cpp binary ----------------------------------------------------
def fetch_sd():
    sd_dir = os.path.join(WORK, "sdcpp")
    if not os.path.isdir(sd_dir):
        zip_path = os.path.join(WORK, "sd.zip")
        log("downloading stable-diffusion.cpp ...")
        urllib.request.urlretrieve(SD_URL, zip_path)
        with zipfile.ZipFile(zip_path) as z:
            z.extractall(sd_dir)
    sd_bin = os.path.join(sd_dir, "sd-cli")
    os.chmod(sd_bin, 0o755)
    return sd_dir, sd_bin


log("fetching binary + model ...")
SD_DIR, SD_BIN = fetch_sd()
MODEL_PATH = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
log("ready.")


def generate(prompt, negative, steps, width, height, seed):
    if not prompt or not prompt.strip():
        raise gr.Error("Введите prompt")
    out = os.path.join(WORK, "out.png")
    if os.path.exists(out):
        os.remove(out)
    env = dict(os.environ)
    env["LD_LIBRARY_PATH"] = SD_DIR + ":" + env.get("LD_LIBRARY_PATH", "")
    cmd = [
        SD_BIN,
        "-m", MODEL_PATH,
        "-p", prompt,
        "-n", negative or "",
        "--cfg-scale", "1.0",          # turbo models: cfg ~1
        "--steps", str(int(steps)),
        "--sampling-method", "euler",
        "-W", str(int(width)),
        "-H", str(int(height)),
        "-s", str(int(seed)),
        "-t", str(N_THREADS),
        "-o", out,
    ]
    t0 = time.time()
    r = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=1200)
    dt = time.time() - t0
    if not os.path.exists(out):
        tail = (r.stderr or r.stdout or "no output")[-800:]
        raise gr.Error(f"Генерация не удалась:\n{tail}")
    return out, f"{dt:.1f}s · steps={int(steps)} · {int(width)}x{int(height)} · seed={int(seed)}"


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", value="a cute cat astronaut floating in space, digital art, highly detailed"),
        gr.Textbox(label="Negative prompt", value="blurry, low quality, deformed"),
        gr.Slider(1, 8, value=3, step=1, label="Steps (turbo: 1-4 достаточно)"),
        gr.Slider(256, 768, value=512, step=64, label="Width"),
        gr.Slider(256, 768, value=512, step=64, label="Height"),
        gr.Number(value=42, label="Seed", precision=0),
    ],
    outputs=[
        gr.Image(label="Результат", type="filepath"),
        gr.Textbox(label="Инфо / тайминг"),
    ],
    title="SD-Turbo — CPU text→image (stable-diffusion.cpp)",
    description=(
        "SD-Turbo на CPU Basic (2 vCPU). 512px, few-step (1-4 шага). "
        "Одна картинка ~30-90 с — наберитесь терпения, это CPU."
    ),
    allow_flagging="never",
)

# Use mount_gradio_app + uvicorn (avoids demo.launch() self-check which crashes
# on gradio_client 1.3.0 get_api_info bug under Python 3.12).
demo.queue(max_size=4)
fastapi_app = FastAPI()
app = gr.mount_gradio_app(fastapi_app, demo, path="/")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)