#!/usr/bin/env python3
"""
RunPod Fresh Setup Script for Qwen Image Edit
==============================================
Downloads everything from YOUR OWN HuggingFace repos:
  - App code: sdfafdfsdf/Qwen-Image-Edit-App
  - Models:   sdfafdfsdf/Phr00t-Qwen-Rapid-AIO (20GB)
              sdfafdfsdf/Qwen-Image-Edit-2511   (16GB)
              sdfafdfsdf/BFS-Best-Face-Swap      (1.3GB)
              sdfafdfsdf/Qwen-Image-Edit-2511-Multiple-Angles-LoRA (282MB)
              sdfafdfsdf/vllm-flash-attn3        (807MB)

Usage:
  1. Upload this script to your RunPod
  2. Run: python3 setup_runpod.py
  3. The app will auto-start at the end

Set your HF token as an environment variable before running:
  export HF_TOKEN="hf_your_token_here"

Or pass it as a command line argument:
  python3 setup_runpod.py --token hf_your_token_here
"""

import os
import sys
import subprocess
import argparse
import time

# ═══════════════════════════════════════════════════════
# CONFIGURATION — Change these if you rename your repos
# ═══════════════════════════════════════════════════════

HF_USERNAME = "sdfafdfsdf"

# App code repo (small files: app.py, qwenimage/, camera_control_ui, etc.)
APP_REPO = f"{HF_USERNAME}/Qwen-Image-Edit-App"

# Model repos (large files: transformer weights, VAE, text encoder, etc.)
MODEL_REPOS = [
    f"{HF_USERNAME}/Phr00t-Qwen-Rapid-AIO",           # 20GB - AIO transformer
    f"{HF_USERNAME}/Qwen-Image-Edit-2511",              # 16GB - Pipeline (VAE, text encoder)
    f"{HF_USERNAME}/BFS-Best-Face-Swap",                 # 1.3GB - Face swap LoRA
    f"{HF_USERNAME}/Qwen-Image-Edit-2511-Multiple-Angles-LoRA",  # 282MB - Angles LoRA
    f"{HF_USERNAME}/vllm-flash-attn3",                   # 807MB - Flash attention kernel
    f"{HF_USERNAME}/QIE_2511_Consistency_Lora",          # Consistency LoRA
]

# Directories
WORKSPACE_DIR = "/workspace"
APP_DIR = os.path.join(WORKSPACE_DIR, "Qwen-Image-Edit")
VENV_DIR = os.path.join(WORKSPACE_DIR, "venv")
CACHE_DIR = os.path.join(WORKSPACE_DIR, "cache")
HF_HOME = os.path.join(CACHE_DIR, "huggingface")
PIP_CACHE = os.path.join(CACHE_DIR, "pip")
TMP_DIR = os.path.join(WORKSPACE_DIR, "tmp")


# ═══════════════════════════════════════════════════════
# HELPERS
# ═══════════════════════════════════════════════════════

def log(msg, level="INFO"):
    """Print a formatted log message."""
    icons = {"INFO": "ℹ️ ", "OK": "✅", "WARN": "⚠️ ", "ERR": "❌", "DL": "📦", "RUN": "🚀"}
    icon = icons.get(level, "  ")
    print(f"{icon} {msg}")


def run(command, cwd=None, env_extra=None, check=True):
    """Run a shell command, stream output, return exit code."""
    log(f"Running: {command}", "RUN")
    env = os.environ.copy()
    env["TMPDIR"] = TMP_DIR
    env["PIP_CACHE_DIR"] = PIP_CACHE
    env["HF_HOME"] = HF_HOME
    if env_extra:
        env.update(env_extra)

    proc = subprocess.Popen(
        command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
        text=True, cwd=cwd, env=env
    )
    for line in proc.stdout:
        print(line, end="")
    proc.wait()

    if proc.returncode != 0 and check:
        log(f"Command failed (exit {proc.returncode}): {command}", "ERR")
    return proc.returncode


def ensure_dirs():
    """Create all necessary directories."""
    for d in [APP_DIR, VENV_DIR, CACHE_DIR, HF_HOME, PIP_CACHE, TMP_DIR]:
        os.makedirs(d, exist_ok=True)


# ═══════════════════════════════════════════════════════
# STEP 1: Download app code
# ═══════════════════════════════════════════════════════

def download_app_code(token):
    """Download app code from your HF repo."""
    log(f"Downloading app code from {APP_REPO}...", "DL")

    pip = os.path.join(VENV_DIR, "bin", "pip")
    python = os.path.join(VENV_DIR, "bin", "python")

    # Use huggingface_hub to download
    run(f"""{python} -c "
from huggingface_hub import snapshot_download
snapshot_download(
    repo_id='{APP_REPO}',
    repo_type='model',
    local_dir='{APP_DIR}',
    token='{token}'
)
print('App code downloaded successfully!')
" """)

    if os.path.exists(os.path.join(APP_DIR, "app.py")):
        log("App code downloaded successfully!", "OK")
    else:
        log("App code download may have failed — app.py not found!", "ERR")
        sys.exit(1)


# ═══════════════════════════════════════════════════════
# STEP 2: Download heavy models
# ═══════════════════════════════════════════════════════

def download_models(token):
    """Download all model repos (cached in HF_HOME)."""
    python = os.path.join(VENV_DIR, "bin", "python")

    for repo in MODEL_REPOS:
        log(f"Downloading model: {repo}...", "DL")
        run(f"""{python} -c "
from huggingface_hub import snapshot_download
snapshot_download(
    repo_id='{repo}',
    repo_type='model',
    token='{token}'
)
print('Downloaded: {repo}')
" """)

    log("All models downloaded!", "OK")


# ═══════════════════════════════════════════════════════
# STEP 3: Set up Python venv
# ═══════════════════════════════════════════════════════

def setup_venv():
    """Create virtual environment if it doesn't exist."""
    if not os.path.exists(os.path.join(VENV_DIR, "bin", "python")):
        log("Creating virtual environment...", "INFO")
        run(f"python3 -m venv {VENV_DIR}")
    else:
        log("Virtual environment already exists.", "OK")

    # Upgrade pip
    pip = os.path.join(VENV_DIR, "bin", "pip")
    run(f"{pip} install --upgrade pip")


# ═══════════════════════════════════════════════════════
# STEP 4: Install dependencies
# ═══════════════════════════════════════════════════════

def install_dependencies():
    """Install Python packages from requirements.txt."""
    pip = os.path.join(VENV_DIR, "bin", "pip")
    req_file = os.path.join(APP_DIR, "requirements.txt")

    if os.path.exists(req_file):
        log("Installing dependencies from requirements.txt...", "DL")
        run(f"{pip} install -r {req_file}")
        log("Dependencies installed!", "OK")
    else:
        log("No requirements.txt found!", "WARN")

    # Ensure huggingface_hub is available (needed for model downloads)
    run(f"{pip} install huggingface_hub")


# ═══════════════════════════════════════════════════════
# STEP 5: Start the app
# ═══════════════════════════════════════════════════════

def start_app():
    """Start the Gradio app."""
    python = os.path.join(VENV_DIR, "bin", "python")
    app_py = os.path.join(APP_DIR, "app.py")

    if not os.path.exists(app_py):
        log("app.py not found, cannot start!", "ERR")
        return

    log("Starting the Gradio app...", "RUN")

    env_extra = {
        "PYTHONPATH": APP_DIR,
        "PYTHONUNBUFFERED": "1",
    }

    # Run in foreground so user can see output
    run(f"{python} {app_py}", cwd=APP_DIR, env_extra=env_extra, check=False)


# ═══════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="Setup Qwen Image Edit on a fresh RunPod")
    parser.add_argument("--token", type=str, default=None,
                        help="HuggingFace token (or set HF_TOKEN env var)")
    parser.add_argument("--skip-models", action="store_true",
                        help="Skip model downloads (if already cached)")
    parser.add_argument("--skip-app", action="store_true",
                        help="Skip app code download (if already present)")
    parser.add_argument("--skip-deps", action="store_true",
                        help="Skip dependency installation")
    parser.add_argument("--no-start", action="store_true",
                        help="Don't auto-start the app after setup")
    args = parser.parse_args()

    # Get token
    token = args.token or os.environ.get("HF_TOKEN", "")
    if not token:
        log("No HuggingFace token provided!", "ERR")
        log("Set HF_TOKEN environment variable or pass --token", "INFO")
        sys.exit(1)

    print("=" * 60)
    print("  Qwen Image Edit — Fresh RunPod Setup")
    print("=" * 60)
    start_time = time.time()

    # Step 0: Create directories
    log("Creating directories...")
    ensure_dirs()

    # Step 1: Set up venv (needed before downloading, since we use huggingface_hub)
    log("Setting up Python virtual environment...")
    setup_venv()

    # Step 2: Install huggingface_hub first (needed for downloads)
    pip = os.path.join(VENV_DIR, "bin", "pip")
    run(f"{pip} install huggingface_hub")

    # Step 3: Download app code
    if not args.skip_app:
        download_app_code(token)
    else:
        log("Skipping app code download (--skip-app)", "INFO")

    # Step 4: Install full dependencies
    if not args.skip_deps:
        install_dependencies()
    else:
        log("Skipping dependency installation (--skip-deps)", "INFO")

    # Step 5: Download models
    if not args.skip_models:
        download_models(token)
    else:
        log("Skipping model downloads (--skip-models)", "INFO")

    elapsed = time.time() - start_time
    minutes = int(elapsed // 60)
    seconds = int(elapsed % 60)

    print("")
    print("=" * 60)
    log(f"Setup completed in {minutes}m {seconds}s!", "OK")
    print("=" * 60)

    # Step 6: Start the app
    if not args.no_start:
        start_app()
    else:
        log("Skipping app start (--no-start). To start manually:", "INFO")
        log(f"  bash /workspace/start_app.sh", "INFO")


if __name__ == "__main__":
    main()