import os
import shutil
import subprocess
import threading
import time

FOOOCUS_REPO = "https://github.com/lllyasviel/Fooocus.git"
FOOOCUS_DIR = "Fooocus"

# Your direct download URL to the .safetensors file
MODEL_URL = "PASTE_YOUR_DIRECT_DOWNLOAD_URL_HERE"
MODEL_FILENAME = "realism_by_stable_yogi_pony.safetensors"

# Save outputs to /tmp so nothing persists
TMP_OUTPUTS = "/tmp/fooocus_outputs"
CLEAN_EVERY_SECONDS = 300


def sh(cmd, cwd=None, check=True):
    print(f"\n[CMD] {cmd}\n", flush=True)
    return subprocess.run(cmd, shell=True, cwd=cwd, check=check)


def ensure_dir(path):
    os.makedirs(path, exist_ok=True)


def reset_tmp_outputs():
    shutil.rmtree(TMP_OUTPUTS, ignore_errors=True)
    ensure_dir(TMP_OUTPUTS)


def link_outputs_to_tmp():
    # Fooocus writes to Fooocus/outputs — redirect to /tmp
    outputs_path = os.path.join(FOOOCUS_DIR, "outputs")

    # remove existing outputs folder/symlink
    if os.path.islink(outputs_path):
        os.unlink(outputs_path)
    elif os.path.exists(outputs_path):
        shutil.rmtree(outputs_path, ignore_errors=True)

    os.symlink(TMP_OUTPUTS, outputs_path)
    print(f"[OK] outputs -> {TMP_OUTPUTS}", flush=True)


def download_checkpoint():
    ckpt_dir = os.path.join(FOOOCUS_DIR, "models", "checkpoints")
    ensure_dir(ckpt_dir)

    dst = os.path.join(ckpt_dir, MODEL_FILENAME)

    if os.path.exists(dst) and os.path.getsize(dst) > 50_000_000:
        print(f"[OK] Model already exists: {dst}", flush=True)
        return

    # Download (fails loudly if URL is wrong/403)
    print(f"[INFO] Downloading model to: {dst}", flush=True)
    sh(f"wget -O '{dst}' '{MODEL_URL}'", check=True)

    # Sanity check: make sure it’s not HTML
    size = os.path.getsize(dst)
    print(f"[INFO] Downloaded size: {size} bytes", flush=True)
    if size < 50_000_000:
        print("[WARN] File looks too small for a checkpoint. Likely not a direct .safetensors download link or got blocked (403/HTML).", flush=True)


def install_fooocus_once():
    marker = os.path.join(FOOOCUS_DIR, ".installed")
    if os.path.exists(marker):
        print("[OK] Fooocus deps already installed.", flush=True)
        return

    sh("pip install -r requirements_versions.txt", cwd=FOOOCUS_DIR, check=True)
    with open(marker, "w") as f:
        f.write("ok\n")
    print("[OK] Fooocus deps installed.", flush=True)


def cleaner_loop():
    while True:
        try:
            # wipe outputs every minute
            for name in os.listdir(TMP_OUTPUTS):
                p = os.path.join(TMP_OUTPUTS, name)
                if os.path.isdir(p):
                    shutil.rmtree(p, ignore_errors=True)
                else:
                    try:
                        os.remove(p)
                    except FileNotFoundError:
                        pass
            print("[CLEAN] Cleared /tmp outputs", flush=True)
        except Exception as e:
            print(f"[CLEAN-WARN] {e}", flush=True)
        time.sleep(CLEAN_EVERY_SECONDS)


def main():
    print("[BOOT] Starting Space…", flush=True)

    # Clone Fooocus if missing
    if not os.path.exists(FOOOCUS_DIR):
        sh(f"git clone {FOOOCUS_REPO}", check=True)

    # outputs -> /tmp (non-persistent)
    reset_tmp_outputs()
    link_outputs_to_tmp()

    # install once
    install_fooocus_once()

    # download your model (only works if URL is direct and not blocked)
    if MODEL_URL and "PASTE_" not in MODEL_URL:
        download_checkpoint()
    else:
        print("[INFO] MODEL_URL not set. Skipping model download.", flush=True)

    # start cleaner
    threading.Thread(target=cleaner_loop, daemon=True).start()

    # launch Fooocus
    print("[BOOT] Launching Fooocus…", flush=True)
    sh("python entry_with_update.py --listen 0.0.0.0 --port 7860", cwd=FOOOCUS_DIR, check=True)


if __name__ == "__main__":
    main()