setup-wizard.py

#!/usr/bin/env python3
"""FreeAI Setup Wizard — interactive configuration assistant.

Guides users through initial setup:
  - GPU detection and CUDA verification
  - Provider API key configuration
  - Service port selection
  - Profile selection (dev/staging/prod)
  - .env file generation from .env.example

Usage:
    python setup-wizard.py              # interactive wizard
    python setup-wizard.py --auto       # auto-detect and apply defaults
    python setup-wizard.py --check      # system check only (no changes)
"""
import os
import sys
import shutil
import subprocess
from pathlib import Path

ROOT = Path(__file__).parent.resolve()
ENV_EXAMPLE = ROOT / ".env.example"
ENV_FILE = ROOT / ".env"


def detect_gpu():
    """Detect available GPUs and CUDA."""
    result = {"cuda": False, "gpu_count": 0, "gpus": []}
    try:
        nvcc = subprocess.run(
            ["nvcc", "--version"], capture_output=True, text=True, timeout=5
        )
        if nvcc.returncode == 0:
            result["cuda"] = True
            result["cuda_version"] = nvcc.stdout.strip()
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    try:
        nvidia_smi = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total,driver_version", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=5,
        )
        if nvidia_smi.returncode == 0:
            for line in nvidia_smi.stdout.strip().split("\n"):
                if line.strip():
                    result["gpus"].append(line.strip())
                    result["gpu_count"] += 1
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    return result


def check_port_available(port):
    """Check if a port is available."""
    import socket
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.bind(("0.0.0.0", port))
        s.close()
        return True
    except OSError:
        return False
    finally:
        s.close()


def read_env_example():
    """Read .env.example and return dict of defaults."""
    if not ENV_EXAMPLE.exists():
        return {}
    env = {}
    for line in ENV_EXAMPLE.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" in line:
            key, _, value = line.partition("=")
            env[key.strip()] = value.strip()
    return env


def write_env(env_dict):
    """Write .env file from dict."""
    with ENV_FILE.open("w", encoding="utf-8") as f:
        f.write("# FreeAI Environment Configuration\n")
        f.write("# Generated by setup-wizard.py\n")
        f.write("# Copy from .env.example and adjust as needed.\n\n")
        for key, value in env_dict.items():
            f.write(f"{key}={value}\n")
    print(f"[wizard] Wrote {ENV_FILE}")


def interactive_wizard():
    """Run interactive setup wizard."""
    print("=" * 60)
    print("  FreeAI Setup Wizard")
    print("=" * 60)
    print()

    # GPU Detection
    gpu_info = detect_gpu()
    print(f"[info] CUDA available: {gpu_info['cuda']}")
    print(f"[info] GPUs detected: {gpu_info['gpu_count']}")
    for gpu in gpu_info["gpus"]:
        print(f"  - {gpu}")
    if not gpu_info["cuda"]:
        print("[warn] No CUDA detected — MOCK_LLM will be enabled by default")
    print()

    # Environment profile
    profile = input("Environment profile [dev/staging/prod] (default: dev): ").strip() or "dev"
    print()

    # Read defaults from .env.example
    defaults = read_env_example()

    # Allow overriding defaults
    overrides = {}
    print("Override defaults? (leave blank to use defaults)")
    print("(Enter 'done' when finished)")
    while True:
        key = input("  Key (e.g. OPENAI_API_KEY) or 'done': ").strip()
        if key.lower() == "done" or not key:
            break
        value = input(f"  Value for {key}: ").strip()
        overrides[key] = value

    # Build final env
    env = dict(defaults)
    env["STACK_PROFILE"] = profile
    env["MOCK_LLM"] = "1" if not gpu_info["cuda"] else env.get("MOCK_LLM", "0")
    env.update(overrides)

    # Write .env
    write_env(env)
    print()
    print("[wizard] Setup complete!")
    print(f"[wizard] Profile: {profile}")
    print(f"[wizard] GPU: {'Yes' if gpu_info['cuda'] else 'No (mock mode)'}")
    print(f"[wizard] Run 'python launch.py' to start all services.")
    print(f"[wizard] Open http://localhost:8030 for the dashboard.")


def auto_mode():
    """Auto-detect and apply defaults."""
    print("[wizard] Running in auto mode...")
    gpu_info = detect_gpu()
    defaults = read_env_example()
    defaults["STACK_PROFILE"] = "dev"
    defaults["MOCK_LLM"] = "1" if not gpu_info["cuda"] else "0"
    write_env(defaults)
    print("[wizard] Auto setup complete!")
    print(f"[wizard] Profile: dev | GPU: {'Yes' if gpu_info['cuda'] else 'No'}")


def check_mode():
    """System check only."""
    print("[wizard] System check...")
    checks = {
        "Python 3.10+": sys.version_info >= (3, 10),
        "venv exists": (ROOT / "venv").exists() or (ROOT / ".venv").exists(),
        ".env.example exists": ENV_EXAMPLE.exists(),
        "requirements.txt exists": (ROOT / "requirements.txt").exists(),
        "launch.py exists": (ROOT / "launch.py").exists(),
        "dashboard/backend.py exists": (ROOT / "dashboard" / "backend.py").exists(),
    }
    all_ok = True
    for check, result in checks.items():
        status = "OK" if result else "FAIL"
        print(f"  [{status}] {check}")
        if not result:
            all_ok = False

    gpu_info = detect_gpu()
    print(f"  [{'OK' if gpu_info['cuda'] else 'WARN'}] CUDA: {gpu_info['cuda']} ({gpu_info['gpu_count']} GPU(s))")

    # Check key ports
    for port in [8010, 8050, 8080, 8090, 8100, 8110, 8120, 8130, 8140, 8150, 8160, 8170, 8180, 8192, 8196]:
        available = check_port_available(port)
        print(f"  [{'OK' if available else 'BUSY'}] Port {port}")

    print()
    if all_ok:
        print("[wizard] All checks passed!")
    else:
        print("[wizard] Some checks failed — review above.")


def main():
    import argparse
    parser = argparse.ArgumentParser(description="FreeAI Setup Wizard")
    parser.add_argument("--auto", action="store_true", help="Auto-detect and apply defaults")
    parser.add_argument("--check", action="store_true", help="System check only (no changes)")
    args = parser.parse_args()

    if args.check:
        check_mode()
    elif args.auto:
        auto_mode()
    else:
        interactive_wizard()


if __name__ == "__main__":
    main()