# app.py – Fooocus + Pony Diffusion V6 XL (Hugging Face Spaces ready)
# Komplett, standalone, mit moderner Torch-Version und Fehlerbehandlung

import os
import subprocess
import time
import threading
import sys
from pathlib import Path
from huggingface_hub import hf_hub_download

# ─── KONFIGURATION ──────────────────────────────────────────────────────────────

MODEL_NAME = "Pony Diffusion V6 XL"
MODEL_REPO = "Lykon/Pony_Diffusion_V6_XL"
MODEL_FILE = "ponyDiffusionV6XL_v6StartWithThisOne.safetensors"

# VAE (sehr empfohlen für bessere Qualität)
VAE_REPO = "stabilityai/sd-vae-ft-mse-original"
VAE_FILE = "vae-ft-mse-840000-ema-pruned.safetensors"

# Fooocus-Repository (aktuellster Branch)
FOOOCUS_REPO = "https://github.com/lllyasviel/Fooocus.git"

# ─── HELPER-FUNKTIONEN ──────────────────────────────────────────────────────────

def run_cmd(cmd, error_msg="Fehler beim Ausführen"):
    """Wrapper für subprocess mit besserer Fehlermeldung"""
    try:
        print(f"[CMD] {' '.join(cmd)}")
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        print(f"{error_msg}: {e}")
        sys.exit(1)


def install_fooocus():
    """Fooocus klonen + Abhängigkeiten installieren"""
    if Path("Fooocus").exists():
        print("Fooocus Ordner existiert bereits → überspringe Klonen")
    else:
        print("→ Klone Fooocus Repository ...")
        run_cmd(["git", "clone", FOOOCUS_REPO])

    os.chdir("Fooocus")

    print("→ Aktualisiere pip & wheel ...")
    run_cmd(["pip", "install", "--upgrade", "pip", "wheel", "setuptools"])

    print("→ Installiere aktuelle Torch-Version (CPU) ...")
    run_cmd([
        "pip", "install",
        "torch==2.5.0", "torchvision==0.20.0",
        "--index-url", "https://download.pytorch.org/whl/cpu",
        "--no-cache-dir"
    ])

    print("→ Installiere Fooocus Abhängigkeiten ...")
    run_cmd(["pip", "install", "-r", "requirements_versions.txt", "--no-cache-dir"])


def download_model():
    """Model und VAE laden – nur wenn nicht vorhanden"""
    checkpoints = Path("models/checkpoints")
    checkpoints.mkdir(parents=True, exist_ok=True)

    model_path = checkpoints / MODEL_FILE
    if not model_path.exists():
        print(f"→ Lade {MODEL_NAME} herunter ...")
        hf_hub_download(
            repo_id=MODEL_REPO,
            filename=MODEL_FILE,
            local_dir=checkpoints,
            local_dir_use_symlinks=False
        )
    else:
        print(f"{MODEL_FILE} bereits vorhanden → überspringe Download")

    vae_dir = Path("models/vae")
    vae_dir.mkdir(parents=True, exist_ok=True)
    vae_path = vae_dir / VAE_FILE
    if not vae_path.exists():
        print("→ Lade empfohlene VAE herunter ...")
        hf_hub_download(
            repo_id=VAE_REPO,
            filename=VAE_FILE,
            local_dir=vae_dir,
            local_dir_use_symlinks=False
        )
    else:
        print("VAE bereits vorhanden")


def start_fooocus():
    """Fooocus mit sinnvollen Parametern starten"""
    cmd = [
        "python", "entry_with_update.py",
        "--share",                        # öffentlicher Link
        "--preset", "anime",              # Pony V6 ist anime-optimiert
        "--disable-offload-from-vram",
        "--always-high-vram",
        "--theme", "dark",
        "--output-dir", "/tmp/outputs",
        "--disable-image-logger",
        "--listen"
    ]

    print("→ Starte Fooocus ... (kann 60–300 Sekunden dauern)")
    print("Der öffentliche Link erscheint im Build-Log oben[](https://*.gradio.live)")
    subprocess.Popen(cmd)


# ─── HAUPTPROGRAMM ──────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print(f"=== Fooocus Bildgenerator – {MODEL_NAME} ===")
    print("Hugging Face Spaces – CPU-kompatibel – zensurarm")
    print("─────────────────────────────────────────────")

    try:
        install_fooocus()
        download_model()

        # Fooocus im Hintergrund starten
        threading.Thread(target=start_fooocus, daemon=True).start()

        print("\n" + "═"*70)
        print("Fooocus läuft jetzt.")
        print("Warte bitte – der öffentliche Gradio-Link erscheint im Build-Log oben")
        print("Format: https://xxxxxx.gradio.live")
        print("Öffne diesen Link im Browser – das ist deine App!")
        print("═"*70 + "\n")

        # Container am Leben halten
        while True:
            time.sleep(300)

    except Exception as e:
        print(f"KRITISCHER FEHLER: {e}")
        sys.exit(1)