KawAI/launcher.py

"""First-run launcher.

Bootstraps a clean Python 3.11 environment via `uv`, regardless of the system
Python the user invoked us with. Keeps the user on a single supported runtime
while the AI ecosystem stabilizes around newer Python versions.

uv install strategy (in order):
1. Direct binary download from GitHub releases (no shell, no admin).
2. PowerShell / shell installer.
3. pip fallback, invoked as `python -m uv`.
"""
from __future__ import annotations

import argparse
import io
import json
import os
import platform
import shutil
import subprocess
import sys
import urllib.request
import zipfile
from pathlib import Path

ROOT = Path(__file__).parent.resolve()
VENV_DIR = ROOT / "venv"
MARKER = VENV_DIR / ".kawai_ready"
HARDWARE_CACHE = ROOT / "config.local.json"
UV_CACHE_DIR = ROOT / ".tools"
PYTHON_TARGET = "3.11"


def venv_python() -> Path:
    if os.name == "nt":
        return VENV_DIR / "Scripts" / "python.exe"
    return VENV_DIR / "bin" / "python"


# --- uv discovery / install ------------------------------------------------

def _local_uv_path() -> Path:
    return UV_CACHE_DIR / ("uv.exe" if os.name == "nt" else "uv")


def _uv_argv() -> list[str] | None:
    """Return argv prefix to invoke uv. None if not available."""
    local = _local_uv_path()
    if local.exists():
        return [str(local)]
    found = shutil.which("uv")
    if found:
        return [found]
    # Installed via pip into current Python (works as module).
    try:
        subprocess.check_output(
            [sys.executable, "-m", "uv", "--version"],
            stderr=subprocess.STDOUT,
            timeout=10,
        )
        return [sys.executable, "-m", "uv"]
    except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
        return None


def _uv_release_asset() -> str | None:
    """Pick the GitHub release asset name for this OS+arch."""
    machine = platform.machine().lower()
    arch_win = "x86_64" if machine in ("amd64", "x86_64") else "aarch64" if "arm" in machine else None
    if os.name == "nt" and arch_win:
        return f"uv-{arch_win}-pc-windows-msvc.zip"
    if sys.platform == "darwin":
        arch = "aarch64" if "arm" in machine else "x86_64"
        return f"uv-{arch}-apple-darwin.tar.gz"
    if sys.platform.startswith("linux"):
        arch = "aarch64" if "aarch64" in machine or "arm64" in machine else "x86_64"
        return f"uv-{arch}-unknown-linux-gnu.tar.gz"
    return None


def _download_uv() -> bool:
    """Download uv binary directly from GitHub releases. Most reliable path."""
    asset = _uv_release_asset()
    if asset is None:
        return False
    url = f"https://github.com/astral-sh/uv/releases/latest/download/{asset}"
    UV_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    print(f"[kawai] Downloading uv from {url}")
    try:
        with urllib.request.urlopen(url, timeout=60) as resp:
            data = resp.read()
    except Exception as e:
        print(f"[kawai]   download failed: {e}")
        return False

    try:
        if asset.endswith(".zip"):
            with zipfile.ZipFile(io.BytesIO(data)) as z:
                for name in z.namelist():
                    if name.endswith("uv.exe") or name.endswith("/uv"):
                        target = _local_uv_path()
                        target.write_bytes(z.read(name))
                        if os.name != "nt":
                            target.chmod(0o755)
                        return target.exists()
        else:
            import tarfile
            with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as t:
                for member in t.getmembers():
                    if member.name.endswith("/uv") or member.name == "uv":
                        f = t.extractfile(member)
                        if f is None:
                            continue
                        target = _local_uv_path()
                        target.write_bytes(f.read())
                        target.chmod(0o755)
                        return target.exists()
    except Exception as e:
        print(f"[kawai]   extract failed: {e}")
        return False
    return False


def _install_uv_via_shell() -> bool:
    """Use astral.sh installer scripts. Often blocked on locked-down systems."""
    UV_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    env = os.environ.copy()
    env["UV_INSTALL_DIR"] = str(UV_CACHE_DIR)
    env["UV_NO_MODIFY_PATH"] = "1"
    try:
        if os.name == "nt":
            subprocess.check_call(
                [
                    "powershell",
                    "-NoProfile",
                    "-ExecutionPolicy", "Bypass",
                    "-Command",
                    "irm https://astral.sh/uv/install.ps1 | iex",
                ],
                env=env,
            )
        else:
            subprocess.check_call(
                ["bash", "-c", "curl -LsSf https://astral.sh/uv/install.sh | sh"],
                env=env,
            )
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False
    return _local_uv_path().exists()


def _install_uv_via_pip() -> bool:
    print("[kawai] Installing uv via pip...")
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "--user", "--upgrade", "uv"])
    except subprocess.CalledProcessError:
        return False
    # Verify it's invokable as a module.
    try:
        subprocess.check_output(
            [sys.executable, "-m", "uv", "--version"],
            stderr=subprocess.STDOUT,
            timeout=10,
        )
        return True
    except Exception:
        return False


def _ensure_uv() -> list[str]:
    argv = _uv_argv()
    if argv:
        return argv

    print("[kawai] Installing uv...")
    if _download_uv():
        return _uv_argv() or []

    if _install_uv_via_shell():
        argv = _uv_argv()
        if argv:
            return argv

    if _install_uv_via_pip():
        argv = _uv_argv()
        if argv:
            return argv

    raise RuntimeError(
        "Failed to install uv. Install manually from https://astral.sh/uv "
        "and rerun this launcher."
    )


# --- venv + deps -----------------------------------------------------------

def _create_venv(uv: list[str]) -> None:
    if venv_python().exists():
        return
    print(f"[kawai] Creating venv with Python {PYTHON_TARGET} (uv will download it if needed)...")
    subprocess.check_call([*uv, "venv", str(VENV_DIR), "--python", PYTHON_TARGET])


def _uv_pip(uv: list[str], args: list[str]) -> None:
    cmd = [*uv, "pip", "install", "--python", str(venv_python()), *args]
    print(f"[kawai] uv pip install {' '.join(args)}")
    subprocess.check_call(cmd)


def detect_and_install(
    uv: list[str],
    force_backend: str | None = None,
    force_vendor: str | None = None,
) -> dict:
    sys.path.insert(0, str(ROOT))
    from backends import hardware

    info = hardware.detect(force_backend=force_backend, force_vendor=force_vendor)
    forced_note = " (forced)" if force_backend and force_backend != "auto" else ""
    print(
        f"[kawai] Backend: {info.backend}{forced_note} | "
        f"{info.vendor} / {info.device_name} / {info.vram_gb:.1f} GB / tier={info.tier}"
    )

    _uv_pip(uv, hardware.torch_install_args(info))
    _uv_pip(uv, ["-r", str(ROOT / "requirements.txt")])

    payload = {
        "vendor": info.vendor,
        "backend": info.backend,
        "device_name": info.device_name,
        "vram_gb": info.vram_gb,
        "tier": info.tier,
        "forced": bool(force_backend and force_backend != "auto"),
    }
    HARDWARE_CACHE.write_text(json.dumps(payload, indent=2))
    MARKER.write_text("ok")
    return payload


def already_in_venv() -> bool:
    try:
        return Path(sys.executable).resolve() == venv_python().resolve()
    except OSError:
        return False


def relaunch_in_venv(forwarded_args: list[str]) -> None:
    """Re-exec the launcher inside the venv. Use subprocess on Windows because
    os.execv mangles argv with spaces in paths."""
    print("[kawai] Relaunching inside venv...")
    py = str(venv_python())
    script = str(ROOT / "launcher.py")
    argv = [py, script, *forwarded_args]
    if os.name == "nt":
        result = subprocess.run(argv)
        sys.exit(result.returncode)
    else:
        os.execv(py, argv)


def _build_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="kawai",
        description="Local AI image/video generator. Auto-detects GPU; pass --backend to override.",
    )
    p.add_argument(
        "--backend",
        choices=["auto", "cuda", "rocm", "directml", "mps", "cpu"],
        default="auto",
        help=(
            "Force torch backend. cuda=NVIDIA, rocm=AMD on Linux, directml=AMD/Intel on Windows, "
            "mps=Apple Silicon, cpu=fallback. Default: auto-detect."
        ),
    )
    p.add_argument(
        "--vendor",
        choices=["nvidia", "amd", "intel", "apple", "cpu"],
        default=None,
        help="Override detected vendor (rarely needed; useful when pairing --backend directml with intel).",
    )
    p.add_argument(
        "--reinstall",
        action="store_true",
        help="Force re-detect and reinstall torch (clears the install marker).",
    )
    return p


def main() -> None:
    args = _build_arg_parser().parse_args()
    forced = args.backend if args.backend != "auto" else None

    if args.reinstall and MARKER.exists():
        print("[kawai] --reinstall: clearing install marker")
        MARKER.unlink()

    if already_in_venv():
        if not MARKER.exists():
            uv = _ensure_uv()
            detect_and_install(uv, force_backend=forced, force_vendor=args.vendor)
        from app import run
        run()
        return

    uv = _ensure_uv()
    _create_venv(uv)
    if not MARKER.exists():
        detect_and_install(uv, force_backend=forced, force_vendor=args.vendor)
    relaunch_in_venv(sys.argv[1:])


if __name__ == "__main__":
    main()