cli.py

#!/usr/bin/env python3
"""
ezlocalai CLI - Run local AI inference with ease.

This CLI manages local LLM, TTS, STT, and image generation.
Supports Docker mode (default on x86_64) and native mode (ARM64/Jetson or --native flag).
Automatically detects GPU availability and architecture to pick the best mode.

Usage:
    ezlocalai start [--model MODEL] [--uri URI] [--api-key KEY] [--ngrok TOKEN] [--native]
    ezlocalai stop
    ezlocalai restart [--model MODEL] [--uri URI] [--api-key KEY] [--ngrok TOKEN] [--native]
    ezlocalai status
    ezlocalai logs [-f]
    ezlocalai prompt "your prompt" [-m MODEL] [-temp TEMPERATURE] [-tp TOP_P] [-image PATH]
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import platform
import re
import shutil
import signal
import socket
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Optional

# Version
__version__ = "1.0.5"

# Configuration
DOCKER_IMAGE = "joshxt/ezlocalai:latest"
DOCKER_IMAGE_CUDA = "ezlocalai:cuda"  # Built locally, not from DockerHub
DOCKER_IMAGE_ROCM = "ezlocalai:rocm"  # Built locally for AMD GPUs
DOCKER_IMAGE_JETSON = "ezlocalai:jetson"  # Built locally on Jetson ARM64+CUDA
CONTAINER_NAME = "ezlocalai"
DEFAULT_PORT = 8091
STATE_DIR = Path.home() / ".ezlocalai"
STATE_DIR.mkdir(parents=True, exist_ok=True)
ENV_FILE = STATE_DIR / ".env"
LOG_FILE = STATE_DIR / "ezlocalai.log"
PID_FILE = STATE_DIR / "ezlocalai.pid"
SOURCE_DIR_FILE = STATE_DIR / "source_dir"
REPO_URL = "https://github.com/DevXT-LLC/ezlocalai.git"
REPO_DIR = STATE_DIR / "repo"

# Cache for uv availability (checked once per process)
_uv_available: Optional[bool] = None


def get_local_ip() -> str:
    """Detect the local network IP address.

    Uses a UDP socket trick to find the IP of the interface that would
    route to the internet, without actually sending any data.
    Falls back to 'localhost' if detection fails.
    """
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
        finally:
            s.close()
        return ip
    except Exception:
        return "localhost"


def _ensure_uv_installed() -> bool:
    """Ensure uv is installed. Installs it if missing.

    uv is a fast Python package manager (10-100x faster than pip).
    Used as a drop-in replacement for 'pip install' throughout the CLI.
    Falls back to pip if uv cannot be installed.

    Returns True if uv is available.
    """
    global _uv_available
    if _uv_available is not None:
        return _uv_available

    if shutil.which("uv"):
        _uv_available = True
        return True

    # Try installing uv via pip (fastest, doesn't need curl/sudo)
    python = sys.executable
    result = subprocess.run(
        [python, "-m", "pip", "install", "uv", "-q"],
        capture_output=True,
        check=False,
    )
    if result.returncode == 0 and shutil.which("uv"):
        _uv_available = True
        print("⚡ Installed uv for faster package management")
        return True

    # Try the official installer as fallback
    try:
        result = subprocess.run(
            ["sh", "-c", "curl -LsSf https://astral.sh/uv/install.sh | sh"],
            capture_output=True,
            check=False,
        )
        # uv installs to ~/.local/bin or ~/.cargo/bin
        for uv_dir in [
            Path.home() / ".local" / "bin",
            Path.home() / ".cargo" / "bin",
        ]:
            uv_path = uv_dir / "uv"
            if uv_path.exists():
                if str(uv_dir) not in os.environ.get("PATH", ""):
                    os.environ["PATH"] = f"{uv_dir}:{os.environ.get('PATH', '')}"
                _uv_available = True
                print("⚡ Installed uv for faster package management")
                return True
    except Exception:
        pass

    _uv_available = False
    return False


def _get_pip_cmd(python: str = None, subcommand: str = None) -> list[str]:
    """Get the pip command prefix with optional subcommand.

    Returns ['uv', 'pip'] or [python, '-m', 'pip'] as the base.
    If subcommand is given (e.g. 'install', 'uninstall'), it is appended
    along with --python for uv (which requires it after the subcommand).

    Examples:
        _get_pip_cmd(python)              -> ['uv', 'pip'] or [python, '-m', 'pip']
        _get_pip_cmd(python, 'install')   -> ['uv', 'pip', 'install', '--python', python]
                                          or [python, '-m', 'pip', 'install']
    """
    if python is None:
        python = sys.executable
    use_uv = _uv_available or _ensure_uv_installed()
    if use_uv:
        cmd = ["uv", "pip"]
        if subcommand:
            cmd.append(subcommand)
            cmd.extend(["--python", python])
        return cmd
    cmd = [python, "-m", "pip"]
    if subcommand:
        cmd.append(subcommand)
    return cmd


def _pip_install(
    packages: list[str],
    python: str = None,
    extra_args: list[str] = None,
    **kwargs,
) -> subprocess.CompletedProcess:
    """Install packages using uv (fast) or pip (fallback).

    Args:
        packages: List of package specs (e.g. ['torch', 'numpy>=1.20'])
        python: Python executable path. Defaults to sys.executable.
        extra_args: Additional args like ['--no-deps', '--index-url', '...']
        **kwargs: Passed to subprocess.run (e.g. capture_output, cwd, env)
    """
    if python is None:
        python = sys.executable
    cmd = _get_pip_cmd(python, "install") + packages
    if extra_args:
        cmd.extend(extra_args)
    kwargs.setdefault("check", False)
    return subprocess.run(cmd, **kwargs)


def _pip_uninstall(
    packages: list[str],
    python: str = None,
    **kwargs,
) -> subprocess.CompletedProcess:
    """Uninstall packages using uv (fast) or pip (fallback)."""
    if python is None:
        python = sys.executable
    cmd = _get_pip_cmd(python, "uninstall") + packages + ["-y"]
    kwargs.setdefault("check", False)
    return subprocess.run(cmd, **kwargs)


def is_arm64() -> bool:
    """Check if running on ARM64/aarch64 architecture."""
    machine = platform.machine().lower()
    return machine in ("aarch64", "arm64")


def is_jetson() -> bool:
    """Check if running on an NVIDIA Jetson device."""
    # Check for Jetson-specific files
    if Path("/etc/nv_tegra_release").exists():
        return True
    try:
        model_path = Path("/proc/device-tree/model")
        if model_path.exists():
            model = model_path.read_text(errors="ignore").lower()
            if "jetson" in model or "tegra" in model:
                return True
    except Exception:
        pass
    return False


def has_jetson_cuda() -> bool:
    """Check if CUDA is available on Jetson (via JetPack/tegrastats)."""
    # Jetson doesn't have nvidia-smi, check for CUDA libraries
    cuda_paths = [
        Path("/usr/local/cuda"),
        Path("/usr/lib/aarch64-linux-gnu/libcuda.so"),
    ]
    if any(p.exists() for p in cuda_paths):
        return True
    # Check tegrastats
    if shutil.which("tegrastats"):
        return True
    return False


def _detect_jetpack_version() -> Optional[str]:
    """Detect the JetPack version on a Jetson device.

    Returns a version string like '60' (JetPack 6.0) or '61' (JetPack 6.1)
    suitable for constructing NVIDIA wheel download URLs. Returns None if
    detection fails.
    """
    # Method 1: dpkg query for nvidia-jetpack meta-package
    try:
        result = subprocess.run(
            ["dpkg-query", "--showformat=${Version}", "-W", "nvidia-jetpack"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0 and result.stdout.strip():
            ver = result.stdout.strip()
            # Parse "6.0-b123" or "6.1.1-b456" -> major.minor
            match = re.match(r"(\d+)\.(\d+)", ver)
            if match:
                return f"{match.group(1)}{match.group(2)}"
    except FileNotFoundError:
        pass

    # Method 2: dpkg query for nvidia-l4t-core (L4T version -> JetPack mapping)
    try:
        result = subprocess.run(
            ["dpkg-query", "--showformat=${Version}", "-W", "nvidia-l4t-core"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0 and result.stdout.strip():
            ver = result.stdout.strip()
            match = re.match(r"(\d+)\.(\d+)", ver)
            if match:
                l4t_major, l4t_minor = int(match.group(1)), int(match.group(2))
                # L4T R36.x = JetPack 6.x
                if l4t_major == 36:
                    if l4t_minor >= 4:
                        return "61"
                    return "60"
                # L4T R35.x = JetPack 5.1
                if l4t_major == 35:
                    return "51"
                # L4T R34.x = JetPack 5.0
                if l4t_major == 34:
                    return "50"
    except FileNotFoundError:
        pass

    # Method 3: Use CUDA version as a rough proxy
    cuda_ver = get_cuda_version()
    if cuda_ver:
        cuda_major_minor = tuple(int(x) for x in cuda_ver.split(".")[:2])
        if cuda_major_minor >= (12, 6):
            return "61"
        if cuda_major_minor >= (12, 2):
            return "60"
        if cuda_major_minor >= (11, 4):
            return "51"

    return None


def _install_jetson_torch(python: str) -> bool:
    """Install CUDA-enabled PyTorch on Jetson.

    PyPI's torch package is CPU-only on aarch64. This function tries multiple
    sources to find a CUDA-enabled wheel that matches the device's JetPack
    version and Python version:

    1. NVIDIA's official redistribution index (cp38/cp310 only typically)
    2. PyTorch's own CUDA index (https://download.pytorch.org/whl/cuXXX)
    3. Jetson AI Lab community index (https://pypi.jetson-ai-lab.dev/simple/)

    Must be called BEFORE installing other requirements so that transitive
    torch dependencies don't pull in the CPU-only version from PyPI.

    Returns True if a CUDA-enabled torch was installed, False otherwise.
    """
    from html.parser import HTMLParser

    jp_ver = _detect_jetpack_version()
    if not jp_ver:
        print(
            "⚠️  Could not detect JetPack version — "
            "skipping NVIDIA PyTorch install (torch will be CPU-only)"
        )
        return False

    py_tag = f"cp{sys.version_info.major}{sys.version_info.minor}"
    print(f"🔍 Detected JetPack {jp_ver[0]}.{jp_ver[1:]}, Python {py_tag}")

    # --- Source 1: NVIDIA official wheel index ---
    base_url = (
        f"https://developer.download.nvidia.com/compute/redist/"
        f"jp/v{jp_ver}/pytorch/"
    )
    print(f"📦 Checking NVIDIA PyTorch wheels at {base_url} ...")

    class _WheelLinkParser(HTMLParser):
        def __init__(self):
            super().__init__()
            self.wheels: list[str] = []

        def handle_starttag(self, tag, attrs):
            if tag == "a":
                for attr_name, attr_val in attrs:
                    if attr_name == "href" and attr_val and attr_val.endswith(".whl"):
                        self.wheels.append(attr_val)

    nvidia_wheel_url = None
    try:
        req = urllib.request.Request(base_url, headers={"User-Agent": "ezlocalai"})
        with urllib.request.urlopen(req, timeout=30) as resp:
            html = resp.read().decode("utf-8", errors="replace")
        parser = _WheelLinkParser()
        parser.feed(html)
        matching = [
            w for w in parser.wheels if f"-{py_tag}-" in w and "linux_aarch64" in w
        ]
        if matching:
            matching.sort()
            wheel_name = matching[-1]
            nvidia_wheel_url = (
                wheel_name if wheel_name.startswith("http") else base_url + wheel_name
            )
            print(f"   Found NVIDIA wheel: {wheel_name}")
        else:
            available_tags = set()
            for w in parser.wheels:
                m = re.search(r"-(cp\d+)-", w)
                if m:
                    available_tags.add(m.group(1))
            print(
                f"   No {py_tag} wheel on NVIDIA index "
                f"(available: {', '.join(sorted(available_tags)) or 'none'})"
            )
    except Exception as exc:
        print(f"   NVIDIA index unavailable: {exc}")

    if nvidia_wheel_url:
        print(f"⬇️  Installing NVIDIA PyTorch: {nvidia_wheel_url.split('/')[-1]}")
        _pip_uninstall(["torch"], python=python, capture_output=True)
        result = _pip_install(
            [nvidia_wheel_url], python=python, extra_args=["--no-cache-dir"]
        )
        if result.returncode == 0 and _verify_torch_cuda(python):
            return True
        print("   NVIDIA wheel installation failed, trying alternatives...")

    # --- Source 2: PyTorch's official CUDA index ---
    # Map JetPack CUDA version to PyTorch's index URL
    cuda_ver = get_cuda_version()
    cuda_index_tag = None
    if cuda_ver:
        cv = tuple(int(x) for x in cuda_ver.split(".")[:2])
        # Match to closest PyTorch CUDA index
        if cv >= (12, 6):
            cuda_index_tag = "cu126"
        elif cv >= (12, 4):
            cuda_index_tag = "cu124"
        elif cv >= (12, 1):
            cuda_index_tag = "cu121"
        elif cv >= (11, 8):
            cuda_index_tag = "cu118"
        elif cv >= (11, 4):
            cuda_index_tag = "cu118"  # closest available

    if cuda_index_tag:
        pytorch_index = f"https://download.pytorch.org/whl/{cuda_index_tag}"
        print(f"📦 Trying PyTorch CUDA index ({cuda_index_tag})...")
        _pip_uninstall(["torch"], python=python, capture_output=True)
        result = _pip_install(
            ["torch"],
            python=python,
            extra_args=["--index-url", pytorch_index, "--no-cache-dir"],
            capture_output=True,
        )
        if result.returncode == 0 and _verify_torch_cuda(python):
            return True
        print(f"   No matching aarch64+CUDA wheel on PyTorch {cuda_index_tag} index")

    # --- Source 3: Jetson AI Lab community index ---
    jetson_index = "https://pypi.jetson-ai-lab.dev/simple/"
    print(f"📦 Trying Jetson AI Lab community index...")
    _pip_uninstall(["torch"], python=python, capture_output=True)
    result = _pip_install(
        ["torch"],
        python=python,
        extra_args=["--extra-index-url", jetson_index, "--no-cache-dir"],
        capture_output=True,
    )
    if result.returncode == 0 and _verify_torch_cuda(python):
        return True
    print("   No matching wheel on Jetson AI Lab index")

    # --- Source 4: PyPI default (CPU-only, last resort) ---
    print("⚠️  No CUDA-enabled PyTorch wheel found for this platform.")
    print("   Installing CPU-only torch from PyPI as fallback...")
    _pip_uninstall(["torch"], python=python, capture_output=True)
    result = _pip_install(
        ["torch"], python=python, extra_args=["--no-cache-dir"], capture_output=True
    )
    if result.returncode == 0:
        print("   torch installed (CPU-only). GPU inference will be slower.")
        return True

    print("   ❌ Failed to install torch from any source")
    return False


def _verify_torch_cuda(python: str) -> bool:
    """Verify that torch is installed and CUDA-enabled."""
    result = subprocess.run(
        [
            python,
            "-c",
            (
                "import torch; "
                "cuda = torch.cuda.is_available(); "
                "print('CUDA available:', cuda); "
                "print('GPU:', torch.cuda.get_device_name(0) if cuda else 'none')"
            ),
        ],
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode == 0 and "CUDA available: True" in result.stdout:
        print(f"   ✅ {result.stdout.strip()}")
        return True
    if result.returncode == 0:
        # torch imported but CUDA not available
        print(f"   ⚠️  {result.stdout.strip()}")
    else:
        stderr = result.stderr.strip().split("\n")[-1] if result.stderr else "unknown"
        print(f"   ⚠️  torch import failed: {stderr}")
    return False


def should_use_native_mode(force_native: bool = False) -> bool:
    """Determine if native mode should be used instead of Docker.

    Native mode is used when:
    - --native flag is passed
    - Docker is not available or not running
    - On Jetson without Docker + nvidia-container-runtime

    On Jetson with Docker available, we prefer Docker mode using the
    locally-built Jetson image (avoids ARM64 dependency hell).
    """
    if force_native:
        return True
    if not is_tool_installed("docker") or not is_docker_running():
        return True
    # Jetson with Docker: use Docker mode if nvidia-container-runtime is available
    if is_arm64() and is_jetson():
        # Check for NVIDIA container runtime (required for GPU access in Docker)
        result = subprocess.run(
            ["docker", "info"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0 and "nvidia" in result.stdout.lower():
            return False  # Docker mode with Jetson image
        # No NVIDIA runtime — fall back to native
        print("⚠️  Docker found but no NVIDIA container runtime.")
        print("   Install with: sudo apt-get install nvidia-container")
        print("   Falling back to native mode...")
        return True
    # Non-Jetson ARM64 without a Docker image
    if is_arm64():
        return True
    return False


def is_ezlocalai_folder(folder: Path) -> bool:
    """Check if the given folder is the ezlocalai source folder.

    Detects the ezlocalai folder by checking for key files that exist
    in the source repository but not in typical installation locations.
    """
    key_files = [
        "docker-compose.yml",
        "docker-compose-cuda.yml",
        "Dockerfile",
        "cuda.Dockerfile",
        "app.py",  # Main app file
    ]
    key_dirs = [
        "ezlocalai",  # Python module folder
    ]
    files_exist = all((folder / f).exists() for f in key_files)
    dirs_exist = all((folder / d).is_dir() for d in key_dirs)
    return files_exist and dirs_exist


def get_ezlocalai_source_dir() -> Optional[Path]:
    """Get the ezlocalai source directory.

    Checks (in order):
    1. Current working directory (if it's the ezlocalai folder)
    2. Persisted source directory from ~/.ezlocalai/source_dir

    When found via cwd, the path is persisted so future invocations
    from other directories still find the source .env and compose files.
    """
    cwd = Path.cwd()
    if is_ezlocalai_folder(cwd):
        # Persist for future use from other directories
        try:
            SOURCE_DIR_FILE.write_text(str(cwd), encoding="utf-8")
        except OSError:
            pass
        return cwd

    # Check persisted path
    if SOURCE_DIR_FILE.exists():
        try:
            saved = Path(SOURCE_DIR_FILE.read_text(encoding="utf-8").strip())
            if saved.exists() and is_ezlocalai_folder(saved):
                return saved
        except (OSError, ValueError):
            pass

    return None


class CLIError(RuntimeError):
    """Raised for recoverable CLI errors."""


def is_tool_installed(tool: str) -> bool:
    """Check if a command-line tool is installed."""
    try:
        result = subprocess.run(
            [tool, "--version"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=False,
        )
        return result.returncode == 0
    except FileNotFoundError:
        return False


def is_docker_running() -> bool:
    """Check if Docker daemon is running."""
    try:
        result = subprocess.run(
            ["docker", "info"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=False,
        )
        return result.returncode == 0
    except FileNotFoundError:
        return False


def has_nvidia_gpu() -> bool:
    """Check if NVIDIA GPU is available."""
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0 and result.stdout.strip():
            return True
    except FileNotFoundError:
        pass
    return False


def has_nvidia_container_toolkit() -> bool:
    """Check if NVIDIA Container Toolkit is installed."""
    try:
        result = subprocess.run(
            [
                "docker",
                "run",
                "--rm",
                "--gpus",
                "all",
                "nvidia/cuda:12.8.1-base-ubuntu24.04",
                "nvidia-smi",
            ],
            capture_output=True,
            text=True,
            check=False,
            timeout=60,
        )
        return result.returncode == 0
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return False


def get_nvidia_gpu_info() -> Optional[str]:
    """Get NVIDIA GPU information."""
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0:
            return result.stdout.strip()
    except FileNotFoundError:
        pass
    return None


def has_amd_gpu() -> bool:
    """Check if AMD GPU is available via ROCm."""
    try:
        result = subprocess.run(
            ["rocm-smi", "--showproductname"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0 and result.stdout.strip():
            return True
    except FileNotFoundError:
        pass
    # Also check for /dev/kfd which indicates ROCm-capable hardware
    if Path("/dev/kfd").exists() and Path("/dev/dri").exists():
        return True
    return False


def has_rocm_support() -> bool:
    """Check if ROCm is properly installed and functional."""
    try:
        result = subprocess.run(
            ["rocminfo"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0 and "Agent" in result.stdout:
            return True
    except FileNotFoundError:
        pass
    return False


def get_amd_gpu_info() -> Optional[str]:
    """Get AMD GPU information."""
    try:
        result = subprocess.run(
            ["rocm-smi", "--showproductname"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0:
            # Parse the output to get GPU name
            for line in result.stdout.splitlines():
                if "GPU" in line or "Card" in line:
                    return line.strip()
            return result.stdout.strip()
    except FileNotFoundError:
        pass
    # Fallback: try lspci
    try:
        result = subprocess.run(
            ["lspci"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0:
            for line in result.stdout.splitlines():
                if "VGA" in line and ("AMD" in line or "Radeon" in line):
                    return line.split(":")[-1].strip()
    except FileNotFoundError:
        pass
    return None


def prompt_user(prompt: str, default: str = "") -> str:
    """Prompt user for input with optional default."""
    if default:
        user_input = input(f"{prompt} (default: {default}): ").strip()
    else:
        user_input = input(f"{prompt}: ").strip()
    return user_input if user_input else default


def _ensure_git_access(repo_dir: Path) -> None:
    """Ensure git operations will work in the given directory.

    Fixes two common issues:
    1. 'dubious ownership' — when repo was cloned with sudo but run as a normal user.
       Adds the directory to git's global safe.directory list.
    2. File ownership — chowns the directory to the current user if not writable.
    """
    if not (repo_dir / ".git").exists():
        return

    # Add to safe.directory if not already there
    result = subprocess.run(
        ["git", "config", "--global", "--get-all", "safe.directory"],
        capture_output=True,
        text=True,
        check=False,
    )
    safe_dirs = result.stdout.strip().splitlines() if result.returncode == 0 else []
    repo_str = str(repo_dir)
    if repo_str not in safe_dirs:
        subprocess.run(
            ["git", "config", "--global", "--add", "safe.directory", repo_str],
            capture_output=True,
            check=False,
        )

    # Fix ownership if directory is not writable by current user
    if not os.access(repo_dir, os.W_OK):
        import getpass

        user = getpass.getuser()
        subprocess.run(
            ["sudo", "chown", "-R", f"{user}:{user}", repo_str],
            capture_output=True,
            check=False,
        )


def clone_or_update_repo() -> Path:
    """Clone or update the ezlocalai repository for building CUDA image.

    Returns the path to the ezlocalai source directory (either local or cloned).
    """
    # First check if we're running from within the ezlocalai folder
    local_source = get_ezlocalai_source_dir()
    if local_source:
        print("📦 Using local ezlocalai source folder...")
        print(f"   Path: {local_source}")
        _ensure_git_access(local_source)
        return local_source

    # Fall back to cloning/updating the repo
    if REPO_DIR.exists():
        _ensure_git_access(REPO_DIR)
        print("📦 Updating ezlocalai repository...")
        result = subprocess.run(
            ["git", "pull"],
            cwd=REPO_DIR,
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode != 0:
            print(f"⚠️  Failed to update repo: {result.stderr}")
            # Try to continue with existing repo
            return REPO_DIR
        print("✅ Repository updated")
        return REPO_DIR
    else:
        print("📦 Cloning ezlocalai repository...")
        result = subprocess.run(
            ["git", "clone", REPO_URL, str(REPO_DIR)],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode != 0:
            print(f"❌ Failed to clone repo: {result.stderr}")
            return None
        print("✅ Repository cloned")
        return REPO_DIR


def build_cuda_image() -> bool:
    """Build the CUDA Docker image from source using docker-compose."""
    source_dir = clone_or_update_repo()
    if not source_dir:
        return False

    print("\n🔨 Building CUDA image (this may take 10-20 minutes)...")
    print("   Building from: docker-compose-cuda.yml")
    print(f"   Source directory: {source_dir}")

    # Build using docker-compose (handles complex builds better)
    result = subprocess.run(
        ["docker", "compose", "-f", "docker-compose-cuda.yml", "build"],
        cwd=source_dir,
        check=False,
    )

    if result.returncode != 0:
        print("❌ Failed to build CUDA image")
        return False

    # Tag the image with our expected name
    # docker-compose names it based on folder name
    print("   Tagging image as ezlocalai:cuda...")

    # Determine the expected image name based on folder
    folder_name = source_dir.name
    expected_names = [
        f"{folder_name}-ezlocalai:latest",
        "repo-ezlocalai:latest",
        f"{folder_name}_ezlocalai:latest",
    ]

    tagged = False
    for expected_name in expected_names:
        tag_result = subprocess.run(
            ["docker", "tag", expected_name, DOCKER_IMAGE_CUDA],
            capture_output=True,
            check=False,
        )
        if tag_result.returncode == 0:
            tagged = True
            break

    if not tagged:
        print("⚠️  Could not tag image, trying to find it...")
        # List images and try to find one that matches
        list_result = subprocess.run(
            ["docker", "images", "--format", "{{.Repository}}:{{.Tag}}"],
            capture_output=True,
            text=True,
            check=False,
        )
        if list_result.returncode == 0:
            for line in list_result.stdout.splitlines():
                if "ezlocalai" in line.lower() and "cuda" not in line:
                    subprocess.run(
                        ["docker", "tag", line.strip(), DOCKER_IMAGE_CUDA],
                        check=False,
                    )
                    tagged = True
                    break

    print("✅ CUDA image built successfully")
    return True


def cuda_image_exists() -> bool:
    """Check if the CUDA image exists locally."""
    result = subprocess.run(
        ["docker", "images", "-q", DOCKER_IMAGE_CUDA],
        capture_output=True,
        text=True,
        check=False,
    )
    return bool(result.stdout.strip())


def rocm_image_exists() -> bool:
    """Check if the ROCm image exists locally."""
    result = subprocess.run(
        ["docker", "images", "-q", DOCKER_IMAGE_ROCM],
        capture_output=True,
        text=True,
        check=False,
    )
    return bool(result.stdout.strip())


def jetson_image_exists() -> bool:
    """Check if the Jetson image exists locally."""
    result = subprocess.run(
        ["docker", "images", "-q", DOCKER_IMAGE_JETSON],
        capture_output=True,
        text=True,
        check=False,
    )
    return bool(result.stdout.strip())


def build_jetson_image() -> bool:
    """Build the Jetson Docker image from source using docker-compose.

    Must be run ON the Jetson itself (ARM64 native build).
    Detects JetPack version and CUDA architecture automatically.
    """
    source_dir = clone_or_update_repo()
    if not source_dir:
        return False

    # Auto-detect L4T tag and CUDA arch
    l4t_tag = "r36.4.0"  # Default: JetPack 6
    cuda_arch = "87"  # Default: Orin

    jp_ver = _detect_jetpack_version()
    if jp_ver:
        jp_major = int(jp_ver[0])
        if jp_major == 5:
            l4t_tag = "r35.4.1"  # JetPack 5.1
        elif jp_major == 6:
            l4t_tag = "r36.4.0"  # JetPack 6

    detected_arch = detect_jetson_cuda_arch()
    if detected_arch:
        cuda_arch = detected_arch

    print(f"\n🔨 Building Jetson image (this may take 15-30 minutes)...")
    print(f"   Building from: docker-compose-jetson.yml")
    print(f"   L4T base: {l4t_tag}, CUDA arch: {cuda_arch}")
    print(f"   Source directory: {source_dir}")

    build_env = os.environ.copy()
    build_env["L4T_TAG"] = l4t_tag
    build_env["CUDA_ARCH"] = cuda_arch

    result = subprocess.run(
        ["docker", "compose", "-f", "docker-compose-jetson.yml", "build"],
        cwd=source_dir,
        env=build_env,
        check=False,
    )

    if result.returncode != 0:
        print("❌ Failed to build Jetson image")
        return False

    # Tag the image with our expected name
    print("   Tagging image as ezlocalai:jetson...")
    folder_name = source_dir.name
    expected_names = [
        f"{folder_name}-ezlocalai:latest",
        "repo-ezlocalai:latest",
        f"{folder_name}_ezlocalai:latest",
    ]

    tagged = False
    for expected_name in expected_names:
        tag_result = subprocess.run(
            ["docker", "tag", expected_name, DOCKER_IMAGE_JETSON],
            capture_output=True,
            check=False,
        )
        if tag_result.returncode == 0:
            tagged = True
            break

    if not tagged:
        list_result = subprocess.run(
            ["docker", "images", "--format", "{{.Repository}}:{{.Tag}}"],
            capture_output=True,
            text=True,
            check=False,
        )
        if list_result.returncode == 0:
            for line in list_result.stdout.splitlines():
                if "ezlocalai" in line.lower():
                    subprocess.run(
                        ["docker", "tag", line.strip(), DOCKER_IMAGE_JETSON],
                        check=False,
                    )
                    tagged = True
                    break

    print("✅ Jetson image built successfully")
    return True


def build_rocm_image() -> bool:
    """Build the ROCm Docker image from source using docker-compose."""
    source_dir = clone_or_update_repo()
    if not source_dir:
        return False

    print("\n🔨 Building ROCm image (this may take 10-20 minutes)...")
    print("   Building from: docker-compose-rocm.yml")
    print(f"   Source directory: {source_dir}")

    # Build using docker-compose (handles complex builds better)
    result = subprocess.run(
        ["docker", "compose", "-f", "docker-compose-rocm.yml", "build"],
        cwd=source_dir,
        check=False,
    )

    if result.returncode != 0:
        print("❌ Failed to build ROCm image")
        return False

    # Tag the image with our expected name
    # docker-compose names it based on folder name
    print("   Tagging image as ezlocalai:rocm...")

    # Determine the expected image name based on folder
    folder_name = source_dir.name
    expected_names = [
        f"{folder_name}-ezlocalai:latest",
        "repo-ezlocalai:latest",
        f"{folder_name}_ezlocalai:latest",
    ]

    tagged = False
    for expected_name in expected_names:
        tag_result = subprocess.run(
            ["docker", "tag", expected_name, DOCKER_IMAGE_ROCM],
            capture_output=True,
            check=False,
        )
        if tag_result.returncode == 0:
            tagged = True
            break

    if not tagged:
        print("⚠️  Could not tag image, trying to find it...")
        # List images and try to find one that matches
        list_result = subprocess.run(
            ["docker", "images", "--format", "{{.Repository}}:{{.Tag}}"],
            capture_output=True,
            text=True,
            check=False,
        )
        if list_result.returncode == 0:
            for line in list_result.stdout.splitlines():
                if "ezlocalai" in line.lower() and "rocm" not in line:
                    subprocess.run(
                        ["docker", "tag", line.strip(), DOCKER_IMAGE_ROCM],
                        check=False,
                    )
                    tagged = True
                    break

    print("✅ ROCm image built successfully")
    return True


def install_docker_linux() -> bool:
    """Attempt to install Docker on Linux."""
    system = platform.system().lower()
    if system != "linux":
        return False

    install = prompt_user(
        "Docker is not installed. Would you like to install it? (y/n)", "y"
    )
    if install.lower() != "y":
        return False

    print("\n🔧 Installing Docker...")

    # Detect package manager and install
    if is_tool_installed("apt-get"):
        commands = [
            "sudo apt-get update",
            "sudo apt-get install -y docker.io",
            "sudo systemctl start docker",
            "sudo systemctl enable docker",
            "sudo usermod -aG docker $USER",
        ]
    elif is_tool_installed("dnf"):
        commands = [
            "sudo dnf install -y docker",
            "sudo systemctl start docker",
            "sudo systemctl enable docker",
            "sudo usermod -aG docker $USER",
        ]
    elif is_tool_installed("yum"):
        commands = [
            "sudo yum install -y docker",
            "sudo systemctl start docker",
            "sudo systemctl enable docker",
            "sudo usermod -aG docker $USER",
        ]
    elif is_tool_installed("pacman"):
        commands = [
            "sudo pacman -Sy --noconfirm docker",
            "sudo systemctl start docker",
            "sudo systemctl enable docker",
            "sudo usermod -aG docker $USER",
        ]
    else:
        print("❌ Unsupported package manager. Please install Docker manually.")
        return False

    try:
        for cmd in commands:
            print(f"  Running: {cmd}")
            subprocess.run(cmd, shell=True, check=True)

        print("\n✅ Docker installed successfully!")
        print(
            "⚠️  You may need to log out and back in for group changes to take effect."
        )
        print("   Or run: newgrp docker")
        return True
    except subprocess.CalledProcessError as e:
        print(f"\n❌ Failed to install Docker: {e}")
        return False


def install_nvidia_container_toolkit() -> bool:
    """Attempt to install NVIDIA Container Toolkit on Linux."""
    system = platform.system().lower()
    if system != "linux":
        return False

    install = prompt_user(
        "\n🎮 NVIDIA GPU detected but Container Toolkit not found.\n"
        "Would you like to install NVIDIA Container Toolkit for GPU acceleration? (y/n)",
        "y",
    )
    if install.lower() != "y":
        return False

    print("\n🔧 Installing NVIDIA Container Toolkit...")

    try:
        # Add NVIDIA repository and install
        if is_tool_installed("apt-get"):
            commands = [
                "curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg",
                "curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list",
                "sudo apt-get update",
                "sudo apt-get install -y nvidia-container-toolkit",
                "sudo nvidia-ctk runtime configure --runtime=docker",
                "sudo systemctl restart docker",
            ]
        elif is_tool_installed("dnf") or is_tool_installed("yum"):
            pkg_mgr = "dnf" if is_tool_installed("dnf") else "yum"
            commands = [
                "curl -s -L https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo | sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo",
                f"sudo {pkg_mgr} install -y nvidia-container-toolkit",
                "sudo nvidia-ctk runtime configure --runtime=docker",
                "sudo systemctl restart docker",
            ]
        else:
            print(
                "❌ Unsupported package manager. Please install NVIDIA Container Toolkit manually:"
            )
            print(
                "   https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html"
            )
            return False

        for cmd in commands:
            print(f"  Running: {cmd}")
            subprocess.run(cmd, shell=True, check=True)

        print("\n✅ NVIDIA Container Toolkit installed successfully!")
        return True
    except subprocess.CalledProcessError as e:
        print(f"\n❌ Failed to install NVIDIA Container Toolkit: {e}")
        print("   You can still run ezlocalai on CPU.")
        return False


def check_prerequisites(native_mode: bool = False) -> tuple[bool, str]:
    """
    Check and install prerequisites.

    Args:
        native_mode: If True, skip Docker checks (running start.py directly).

    Returns:
        Tuple of (docker_available, gpu_type) where gpu_type is 'nvidia', 'amd', or 'cpu'
    """
    system = platform.system().lower()

    if native_mode:
        # Native mode: skip Docker, just detect GPU
        print("ℹ️  Running in native mode (no Docker)")

        gpu_type = "cpu"
        # Check for NVIDIA GPU (nvidia-smi or Jetson)
        if has_nvidia_gpu():
            gpu_info = get_nvidia_gpu_info()
            print(f"NVIDIA GPU detected: {gpu_info}")
            gpu_type = "nvidia"
        elif is_jetson() and has_jetson_cuda():
            print("NVIDIA Jetson detected with CUDA support")
            gpu_type = "nvidia"
        elif has_amd_gpu() and has_rocm_support():
            gpu_info = get_amd_gpu_info()
            print(f"AMD GPU detected: {gpu_info}")
            gpu_type = "amd"
        else:
            print("ℹ️  No GPU detected, running on CPU")

        return True, gpu_type

    # Docker mode: check Docker
    if not is_tool_installed("docker"):
        print("❌ Docker is not installed.")
        if system == "linux":
            if not install_docker_linux():
                print("\n📦 Please install Docker manually:")
                print("   https://docs.docker.com/engine/install/")
                sys.exit(1)
        else:
            print("\n📦 Please install Docker Desktop:")
            if system == "darwin":
                print("   https://docs.docker.com/desktop/install/mac-install/")
            elif system == "windows":
                print("   https://docs.docker.com/desktop/install/windows-install/")
            else:
                print("   https://docs.docker.com/engine/install/")
            sys.exit(1)

    # Check if Docker daemon is running
    if not is_docker_running():
        print("❌ Docker daemon is not running.")
        print("   Please start Docker and try again.")
        if system == "linux":
            print("   Try: sudo systemctl start docker")
        sys.exit(1)

    # print("✅ Docker is installed and running")

    # Check for NVIDIA GPU first
    gpu_type = "cpu"
    if is_jetson() and has_jetson_cuda():
        # Jetson detected with CUDA — use the Jetson Docker image
        print("NVIDIA Jetson detected with CUDA support")
        gpu_type = "jetson"
    elif has_nvidia_gpu():
        gpu_info = get_nvidia_gpu_info()
        print(f"NVIDIA GPU detected: {gpu_info}")

        if has_nvidia_container_toolkit():
            # print("✅ NVIDIA Container Toolkit is installed")
            gpu_type = "nvidia"
        else:
            if system == "linux":
                if install_nvidia_container_toolkit():
                    if has_nvidia_container_toolkit():
                        gpu_type = "nvidia"
                else:
                    print("⚠️  Running on CPU (GPU acceleration not available)")
            else:
                print("⚠️  NVIDIA Container Toolkit not detected.")
                print("   For GPU acceleration, install it from:")
                print(
                    "   https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html"
                )
                print("   Continuing with CPU mode...")
    # Check for AMD GPU if no NVIDIA GPU
    elif has_amd_gpu():
        gpu_info = get_amd_gpu_info()
        print(f"✅ AMD GPU detected: {gpu_info}")

        if has_rocm_support():
            print("✅ ROCm is installed and functional")
            gpu_type = "amd"
        else:
            print("⚠️  ROCm not detected or not functional.")
            if system == "linux":
                print("   For AMD GPU acceleration, install ROCm:")
                print(
                    "   https://rocm.docs.amd.com/projects/install-on-linux/en/latest/"
                )
            print("   Continuing with CPU mode...")
    else:
        print("ℹ️  No GPU detected, running on CPU")

    return True, gpu_type


def is_container_running() -> bool:
    """Check if any ezlocalai container is running.

    Uses docker ps name filter (substring match) so it catches both
    'docker run --name ezlocalai' and compose-created containers like
    'ezlocalai-ezlocalai-1'.
    """
    try:
        result = subprocess.run(
            ["docker", "ps", "-q", "-f", f"name={CONTAINER_NAME}"],
            capture_output=True,
            text=True,
            check=False,
        )
        return bool(result.stdout.strip())
    except FileNotFoundError:
        return False


def _wait_for_container_stop(timeout: int = 30) -> bool:
    """Wait until no ezlocalai container is running.

    Returns True if the container stopped within the timeout, False otherwise.
    """
    for _ in range(timeout):
        if not is_container_running():
            return True
        time.sleep(1)
    return False


def get_container_status() -> Optional[str]:
    """Get container status (running, exited, etc.)."""
    try:
        result = subprocess.run(
            [
                "docker",
                "ps",
                "-a",
                "--format",
                "{{.Status}}",
                "-f",
                f"name={CONTAINER_NAME}",
            ],
            capture_output=True,
            text=True,
            check=False,
        )
        return result.stdout.strip() if result.stdout.strip() else None
    except FileNotFoundError:
        return None


def _parse_env_file(path: Path) -> dict:
    """Parse a .env file and return key-value pairs."""
    env_vars = {}
    if path.exists():
        for line in path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" in line:
                key, value = line.split("=", 1)
                # Remove quotes
                value = value.strip().strip('"').strip("'")
                env_vars[key.strip()] = value
    return env_vars


def load_env_file() -> dict:
    """Load environment variables, preferring the source directory .env.

    Priority (highest wins):
    1. .env in the ezlocalai source directory (if running from within it)
    2. ~/.ezlocalai/.env (CLI state file)
    """
    env_vars = {}
    # Start with CLI state file
    env_vars.update(_parse_env_file(ENV_FILE))
    # Override with source directory .env if available
    source_dir = get_ezlocalai_source_dir()
    if source_dir:
        local_env = source_dir / ".env"
        env_vars.update(_parse_env_file(local_env))
    return env_vars


def save_env_file(env_vars: dict) -> None:
    """Save environment variables to state file."""
    lines = []
    for key, value in sorted(env_vars.items()):
        lines.append(f'{key}="{value}"')
    ENV_FILE.write_text("\n".join(lines) + "\n", encoding="utf-8")


def _get_compose_file(gpu_type: str) -> str:
    """Return the docker-compose filename for the given GPU type."""
    return {
        "nvidia": "docker-compose-cuda.yml",
        "jetson": "docker-compose-jetson.yml",
        "amd": "docker-compose-rocm.yml",
        "cpu": "docker-compose.yml",
    }.get(gpu_type, "docker-compose.yml")


def _get_compose_source_dir() -> Optional[Path]:
    """Get the source directory where docker-compose files live.

    Returns the local source dir if running from within it,
    otherwise the cloned repo dir if it exists, or the editable
    install location (pip install -e .), or None.
    """
    local = get_ezlocalai_source_dir()
    if local:
        return local
    if REPO_DIR.exists() and is_ezlocalai_folder(REPO_DIR):
        return REPO_DIR
    # Check editable install location (pip install -e . stores __file__ path)
    try:
        import ezlocalai as _pkg

        pkg_dir = Path(_pkg.__file__).resolve().parent.parent
        if is_ezlocalai_folder(pkg_dir):
            return pkg_dir
    except Exception:
        pass
    return None


def _write_compose_env(source_dir: Path, env_vars: dict) -> None:
    """Write env vars to the source directory .env for Docker Compose.

    Docker Compose reads .env from the project directory (where the
    compose file lives). This ensures the compose services pick up
    CLI-managed configuration.
    """
    env_file = source_dir / ".env"
    # Sensitive keys are passed via process environment to docker compose,
    # so exclude them from the on-disk .env file to avoid clear-text storage.
    _sensitive_suffixes = ("_KEY", "_TOKEN", "_SECRET", "_PASSWORD")
    lines = ["# Managed by ezlocalai CLI — do not edit manually"]
    for key, value in sorted(env_vars.items()):
        if not key.startswith("_") and not key.upper().endswith(_sensitive_suffixes):
            lines.append(f"{key}={value}")
    content = "\n".join(lines) + "\n"
    try:
        fd = os.open(str(env_file), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        try:
            os.write(fd, content.encode("utf-8"))
        finally:
            os.close(fd)
    except PermissionError:
        # Try with sudo-like workaround: write to temp then move
        import tempfile

        try:
            fd, tmp_path = tempfile.mkstemp(suffix=".env", dir=str(source_dir))
            os.fchmod(fd, 0o600)
            os.write(fd, content.encode("utf-8"))
            os.close(fd)
            os.replace(tmp_path, str(env_file))
        except (PermissionError, OSError):
            # Last resort: write to STATE_DIR and pass via env vars instead
            print(
                f"⚠️  Cannot write {env_file} (permission denied). "
                f"Using environment variables directly."
            )


def get_default_env() -> dict:
    """Get default environment variables.

    These should match the defaults in Globals.py so Docker and native
    modes behave identically when no .env is present.
    """
    return {
        "EZLOCALAI_URL": f"http://{get_local_ip()}:{DEFAULT_PORT}",
        "EZLOCALAI_API_KEY": "",
        "DEFAULT_MODEL": "unsloth/Qwen3.5-4B-GGUF",
        "WHISPER_MODEL": "large-v3",
        "IMG_MODEL": "none",
        "VIDEO_MODEL": "none",
        "NGROK_TOKEN": "",
        "MAX_QUEUE_SIZE": "100",
        "REQUEST_TIMEOUT": "300",
    }


def _get_effective_url(env_vars: dict = None) -> str:
    """Get the effective ezlocalai URL from env vars or defaults.

    Loads env vars if not provided, falls back to default.
    """
    if env_vars is None:
        env_vars = load_env_file()
    url = env_vars.get("EZLOCALAI_URL", "")
    if not url:
        url = f"http://{get_local_ip()}:{DEFAULT_PORT}"
    return url.rstrip("/")


# ── Native mode process management ──────────────────────────────────────────


def get_native_pid() -> Optional[int]:
    """Get the PID of the running native ezlocalai process."""
    if not PID_FILE.exists():
        return None
    try:
        pid = int(PID_FILE.read_text().strip())
        # Check if process is actually running
        os.kill(pid, 0)
        return pid
    except (ValueError, ProcessLookupError, PermissionError):
        # PID file is stale
        PID_FILE.unlink(missing_ok=True)
        return None


def is_native_running() -> bool:
    """Check if ezlocalai is running in native mode."""
    return get_native_pid() is not None


def get_source_dir_for_native() -> Path:
    """Get the ezlocalai source directory for native mode.

    Checks (in order):
    1. Current working directory (if it's the ezlocalai folder)
    2. Cloned repo in ~/.ezlocalai/repo
    """
    local_source = get_ezlocalai_source_dir()
    if local_source:
        return local_source
    if REPO_DIR.exists() and is_ezlocalai_folder(REPO_DIR):
        return REPO_DIR
    return None


def get_cuda_version() -> Optional[str]:
    """Detect the installed CUDA toolkit version.

    Returns a string like '12.8' or '12.4', or None if not found.
    """
    # Try nvcc first
    try:
        result = subprocess.run(
            ["nvcc", "--version"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0:
            # Parse "release 12.8" or "V12.8.93"
            match = re.search(r"(?:release |V)(\d+\.\d+)", result.stdout)
            if match:
                return match.group(1)
    except FileNotFoundError:
        pass

    # Try reading from CUDA path
    for cuda_dir in [
        Path("/usr/local/cuda"),
        Path("/usr/local/cuda-12.8"),
        Path("/usr/local/cuda-12.4"),
    ]:
        version_file = cuda_dir / "version.json"
        if version_file.exists():
            try:
                data = json.loads(version_file.read_text())
                ver = data.get("cuda", {}).get("version", "")
                if ver:
                    parts = ver.split(".")
                    return f"{parts[0]}.{parts[1]}" if len(parts) >= 2 else ver
            except Exception:
                pass
        version_txt = cuda_dir / "version.txt"
        if version_txt.exists():
            try:
                match = re.search(r"(\d+\.\d+)", version_txt.read_text())
                if match:
                    return match.group(1)
            except Exception:
                pass

    return None


def get_rocm_version() -> Optional[str]:
    """Detect the installed ROCm version.

    Returns a string like '6.4.1' or '6.3.4', or None if not found.
    """
    # Check /opt/rocm/.info/version
    version_file = Path("/opt/rocm/.info/version")
    if version_file.exists():
        try:
            ver = version_file.read_text().strip()
            if ver:
                return ver
        except Exception:
            pass

    # Try rocm-smi
    try:
        result = subprocess.run(
            ["rocm-smi", "--showdriverversion"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0:
            match = re.search(r"(\d+\.\d+\.\d+)", result.stdout)
            if match:
                return match.group(1)
    except FileNotFoundError:
        pass

    return None


XLLAMACPP_REPO = "https://github.com/xorbitsai/xllamacpp.git"
XLLAMACPP_BUILD_DIR = STATE_DIR / "xllamacpp-build"


def build_xllamacpp_from_source(gpu_type: str = "nvidia") -> list[str]:
    """Build xllamacpp from source for platforms without prebuilt GPU wheels (e.g. Jetson ARM64).

    Clones the repo, sets the appropriate build env vars, and does 'pip install .'.
    Returns the pip install command list (for compatibility with get_xllamacpp_install_cmd),
    but actually performs the build inline and returns an empty install command.
    """
    python = sys.executable

    # Ensure ~/.cargo/bin is in PATH (Rust may already be installed)
    cargo_bin = Path.home() / ".cargo" / "bin"
    if cargo_bin.exists() and str(cargo_bin) not in os.environ.get("PATH", ""):
        os.environ["PATH"] = f"{cargo_bin}:{os.environ.get('PATH', '')}"

    # Check for build prerequisites
    if not shutil.which("cmake"):
        print("   ⚠️  cmake not found, attempting to install...")
        subprocess.run(
            ["sudo", "apt-get", "install", "-y", "cmake", "build-essential"],
            check=False,
        )

    # Check for Rust (required by xllamacpp build)
    if not shutil.which("rustc"):
        print("   ⚠️  Rust toolchain not found, installing...")
        result = subprocess.run(
            [
                "sh",
                "-c",
                "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y",
            ],
            check=False,
        )
        if result.returncode != 0:
            print("   ❌ Failed to install Rust. xllamacpp will fall back to CPU.")
            return _get_pip_cmd(python, "install") + ["xllamacpp", "-q"]
        # Add cargo to PATH for this session
        if cargo_bin.exists() and str(cargo_bin) not in os.environ.get("PATH", ""):
            os.environ["PATH"] = f"{cargo_bin}:{os.environ.get('PATH', '')}"

    # Install build dependencies (cython, setuptools needed for xllamacpp)
    print("   Installing build dependencies...")
    _pip_install(
        ["cython", "setuptools", "wheel"],
        python=python,
        extra_args=["-q"],
        capture_output=True,
    )

    # Clone or update xllamacpp source
    if XLLAMACPP_BUILD_DIR.exists():
        print("   Updating xllamacpp source...")
        subprocess.run(
            ["git", "pull"],
            cwd=XLLAMACPP_BUILD_DIR,
            capture_output=True,
            check=False,
        )
    else:
        print("   Cloning xllamacpp source...")
        result = subprocess.run(
            ["git", "clone", "--recursive", XLLAMACPP_REPO, str(XLLAMACPP_BUILD_DIR)],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode != 0:
            print(f"   ❌ Failed to clone xllamacpp: {result.stderr.strip()}")
            print("   Falling back to CPU-only aarch64 wheel from PyPI...")
            return _get_pip_cmd(python, "install") + ["xllamacpp", "-q"]

    # Ensure submodules are initialized
    subprocess.run(
        ["git", "submodule", "update", "--init", "--recursive"],
        cwd=XLLAMACPP_BUILD_DIR,
        capture_output=True,
        check=False,
    )

    # Install xllamacpp's own build requirements if present
    xllamacpp_reqs = XLLAMACPP_BUILD_DIR / "requirements.txt"
    if xllamacpp_reqs.exists():
        _pip_install(
            ["-r", str(xllamacpp_reqs)],
            python=python,
            extra_args=["-q"],
            capture_output=True,
        )

    # Set build environment for CUDA
    build_env = os.environ.copy()
    if gpu_type == "nvidia":
        build_env["XLLAMACPP_BUILD_CUDA"] = "1"
        # Detect CUDA arch for Jetson (e.g. sm_72 for Xavier, sm_87 for Orin)
        cuda_arch = detect_jetson_cuda_arch()
        if cuda_arch:
            build_env["CMAKE_CUDA_ARCHITECTURES"] = cuda_arch
            print(f"   Building with CUDA arch: {cuda_arch}")

        # Ensure nvcc is available
        nvcc = shutil.which("nvcc")
        if not nvcc:
            cuda_bin = Path("/usr/local/cuda/bin")
            if cuda_bin.exists():
                build_env["PATH"] = f"{cuda_bin}:{build_env.get('PATH', '')}"
                print(f"   Added {cuda_bin} to PATH for nvcc")
            else:
                print("   ⚠️  nvcc not found — CUDA build may fail")
                print("   Try: sudo apt-get install nvidia-cuda-toolkit")

    # On memory-constrained devices (e.g. Jetson), limit parallel compile jobs
    # to avoid OOM kills. Jetson Orin NX 16GB shares RAM with GPU.
    try:
        import multiprocessing

        total_ram_gb = (
            os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / (1024**3)
        )
        cpus = multiprocessing.cpu_count()
        if total_ram_gb < 24:
            # ~2GB per compile job is safe for CUDA builds
            max_jobs = max(1, min(cpus, int(total_ram_gb / 2.5)))
            build_env["CMAKE_BUILD_PARALLEL_LEVEL"] = str(max_jobs)
            build_env["MAKEFLAGS"] = f"-j{max_jobs}"
            print(
                f"   Limiting to {max_jobs} parallel compile jobs "
                f"({total_ram_gb:.0f}GB RAM, {cpus} CPUs)"
            )
    except Exception:
        pass

    # Build and install (use pip for source builds — uv doesn't support build envs well)
    print("   Building xllamacpp from source (this may take several minutes)...")
    result = subprocess.run(
        [
            python,
            "-m",
            "pip",
            "install",
            ".",
            "--force-reinstall",
            "--no-build-isolation",
            "-v",
        ],
        cwd=XLLAMACPP_BUILD_DIR,
        env=build_env,
        capture_output=True,
        text=True,
        check=False,
    )

    if result.returncode != 0:
        # Capture the actual compilation error — look for g++/nvcc errors
        all_output = (result.stderr or "") + "\n" + (result.stdout or "")
        error_lines = []
        for line in all_output.splitlines():
            low = line.lower()
            if any(
                kw in low
                for kw in [
                    "error:",
                    "fatal error",
                    "cannot find",
                    "no such file",
                    "killed",
                    "oom",
                    "out of memory",
                    "failed",
                ]
            ):
                error_lines.append(line.strip())

        print("   ❌ Source build failed:")
        if error_lines:
            # Show unique, deduplicated error lines (up to 20)
            seen = set()
            for line in error_lines:
                if line not in seen:
                    seen.add(line)
                    print(f"      {line}")
                    if len(seen) >= 20:
                        break
        else:
            # Fall back to last N lines of stderr
            for line in (result.stderr or "").splitlines()[-15:]:
                print(f"      {line}")

        # Check for common fixable issues
        if "cuda_runtime" in all_output or "cuda.h" in all_output:
            print(
                "   💡 Missing CUDA headers. Try: sudo apt-get install nvidia-cuda-toolkit nvidia-cuda-dev"
            )
        if "killed" in all_output.lower() or "cannot allocate" in all_output.lower():
            print(
                "   💡 Build killed (likely OOM). Try closing other apps or reducing GPU memory usage."
            )

        print("   Falling back to CPU-only aarch64 wheel from PyPI...")
        return _get_pip_cmd(python, "install") + ["xllamacpp", "-q"]

    print("   ✅ xllamacpp built from source with CUDA support")
    # Return a no-op command since we already installed
    return [python, "-c", "pass"]


def detect_jetson_cuda_arch() -> Optional[str]:
    """Detect the CUDA compute capability for Jetson devices.

    Returns the SM architecture string (e.g. '72' for Xavier, '87' for Orin).
    """
    # Known Jetson models and their CUDA architectures
    jetson_archs = {
        "nano": "53",
        "tx1": "53",
        "tx2": "62",
        "xavier": "72",
        "agx xavier": "72",
        "xavier nx": "72",
        "orin": "87",
        "agx orin": "87",
        "orin nx": "87",
        "orin nano": "87",
    }

    # Try reading device model
    try:
        model_path = Path("/proc/device-tree/model")
        if model_path.exists():
            model = model_path.read_text(errors="ignore").lower().strip("\x00").strip()
            for name, arch in jetson_archs.items():
                if name in model:
                    return arch
    except Exception:
        pass

    # Try checking /etc/nv_tegra_release for SoC info
    try:
        tegra_path = Path("/etc/nv_tegra_release")
        if tegra_path.exists():
            content = tegra_path.read_text().lower()
            # R36 = Orin, R35 = Orin, R32 = Xavier/Nano
            if "r36" in content or "r35" in content:
                return "87"
            elif "r32" in content:
                return "72"
    except Exception:
        pass

    return None


def get_xllamacpp_install_cmd(gpu_type: str = "cpu") -> list[str]:
    """Build the pip install command for xllamacpp based on GPU type.

    Returns the full command list for subprocess.run().
    Uses GPU-specific index URLs for CUDA/ROCm wheels, or plain PyPI for CPU/ARM64.
    """
    python = sys.executable
    base_cmd = _get_pip_cmd(python, "install") + ["xllamacpp"]

    if gpu_type == "nvidia" and not is_arm64():
        # x86_64 NVIDIA: use CUDA index URL
        cuda_ver = get_cuda_version()
        if cuda_ver:
            major_minor = cuda_ver  # e.g. "12.8"
            major, minor = major_minor.split(".")
            if int(major) >= 12 and int(minor) >= 8:
                index_url = "https://xorbitsai.github.io/xllamacpp/whl/cu128"
            elif int(major) >= 12 and int(minor) >= 4:
                index_url = "https://xorbitsai.github.io/xllamacpp/whl/cu124"
            else:
                # Older CUDA, try cu124 as closest
                print(f"   ⚠️  CUDA {cuda_ver} detected, trying cu124 wheel...")
                index_url = "https://xorbitsai.github.io/xllamacpp/whl/cu124"
            base_cmd.extend(["--force-reinstall", "--index-url", index_url])
            print(f"   Using CUDA {cuda_ver} wheel from {index_url}")
        else:
            print("   ⚠️  Could not detect CUDA version, installing CPU xllamacpp")

    elif gpu_type == "nvidia" and is_arm64():
        # ARM64 + NVIDIA (Jetson): no prebuilt CUDA wheels for aarch64
        # Must build from source with CUDA enabled
        print("   ARM64 + CUDA detected (Jetson) — building xllamacpp from source...")
        return build_xllamacpp_from_source(gpu_type="nvidia")

    elif gpu_type == "amd":
        # AMD ROCm: use ROCm index URL
        rocm_ver = get_rocm_version()
        if rocm_ver:
            # Map to nearest supported version
            major_minor_patch = rocm_ver.split(".")
            ver_tuple = tuple(int(x) for x in major_minor_patch[:3])
            if ver_tuple >= (6, 4, 0):
                index_url = "https://xorbitsai.github.io/xllamacpp/whl/rocm-6.4.1"
            else:
                index_url = "https://xorbitsai.github.io/xllamacpp/whl/rocm-6.3.4"
            base_cmd.extend(["--force-reinstall", "--index-url", index_url])
            print(f"   Using ROCm {rocm_ver} wheel from {index_url}")
        else:
            print("   ⚠️  Could not detect ROCm version, installing CPU xllamacpp")

    # CPU or ARM64: plain PyPI install (aarch64 CPU wheels available on PyPI)
    base_cmd.append("-q")
    return base_cmd


def install_native_dependencies(source_dir: Path, gpu_type: str = "cpu") -> bool:
    """Install Python dependencies for native mode.

    Uses the source directory's requirements files. On Jetson/ARM64 with CUDA,
    installs cuda-requirements.txt; otherwise uses requirements.txt.
    Installs xllamacpp with the correct GPU-specific wheel (CUDA/ROCm index URLs).

    On ARM64, installs packages one-by-one to avoid a single missing wheel
    (e.g. torchcodec) from blocking all other dependencies.
    """
    python = sys.executable

    # Ensure uv is available for fast package installation
    _ensure_uv_installed()

    # On Jetson, install NVIDIA's CUDA-enabled PyTorch BEFORE other packages
    # PyPI's torch is CPU-only; Jetson needs the JetPack-compatible wheel
    if is_jetson() and has_jetson_cuda():
        _install_jetson_torch(python)

    # Determine which requirements file to use
    if gpu_type == "nvidia" or (is_jetson() and has_jetson_cuda()):
        req_file = source_dir / "cuda-requirements.txt"
        if not req_file.exists():
            req_file = source_dir / "requirements.txt"
        print("📦 Installing CUDA dependencies (this may take several minutes)...")
    else:
        req_file = source_dir / "requirements.txt"
        print("📦 Installing dependencies...")

    # Install xllamacpp with correct GPU backend
    print("   Installing xllamacpp...")
    xllamacpp_cmd = get_xllamacpp_install_cmd(gpu_type)
    result = subprocess.run(
        xllamacpp_cmd,
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        print(f"⚠️  xllamacpp install warning: {result.stderr.strip()}")
        if gpu_type in ("nvidia", "amd"):
            print("   Falling back to CPU-only xllamacpp...")
            result = _pip_install(
                ["xllamacpp"],
                python=python,
                extra_args=["-q"],
                capture_output=True,
                text=True,
            )
            if result.returncode != 0:
                print(f"⚠️  xllamacpp CPU fallback also failed: {result.stderr.strip()}")

    # Install chatterbox-tts with --no-deps (same as Dockerfile)
    print("   Installing chatterbox-tts...")
    result = _pip_install(
        ["chatterbox-tts"],
        python=python,
        extra_args=["--no-deps", "-q"],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        print(f"⚠️  chatterbox-tts install warning: {result.stderr.strip()}")

    # Install main requirements
    print(f"   Installing from {req_file.name}...")

    if is_arm64():
        # On ARM64, install packages one-by-one so a missing wheel for one package
        # (e.g. torchcodec) doesn't prevent all other deps from being installed.
        _install_requirements_individually(python, req_file)
    else:
        # On x86_64, install all at once (faster, all wheels available)
        result = _pip_install(
            ["-r", str(req_file)],
            python=python,
            extra_args=["-q"],
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            print(f"⚠️  Some dependencies failed to install:")
            for line in result.stderr.splitlines():
                if "error" in line.lower():
                    print(f"      {line.strip()}")
            print("   Continuing anyway — some features may be unavailable.")

    print("✅ Dependencies installed")
    return True


def _install_requirements_individually(python: str, req_file: Path) -> None:
    """Install requirements with ARM64 resilience.

    Strategy: try batch install first (preserves pip's version resolver for
    compatible huggingface_hub/transformers/diffusers versions). If batch fails,
    identify failing packages and install the rest individually.

    On Jetson with NVIDIA torch installed, packages that have strict torch
    version coupling (torchaudio, torchcodec) are installed with --no-deps
    to prevent pip from replacing the NVIDIA CUDA-enabled torch with PyPI's
    CPU-only version.
    """
    jetson_with_cuda = is_jetson() and has_jetson_cuda()

    # Packages that have exact torch version constraints (e.g. torch==2.4.0)
    # which would cause pip to replace our NVIDIA CUDA torch with CPU-only PyPI torch
    TORCH_COUPLED_PKGS = {"torch", "torchaudio", "torchcodec", "torchvision"}

    # First, try batch install — this keeps the dependency resolver intact
    result = _pip_install(
        ["-r", str(req_file)],
        python=python,
        extra_args=["-q"],
        capture_output=True,
        text=True,
    )
    if result.returncode == 0:
        if jetson_with_cuda:
            # Verify NVIDIA torch wasn't overwritten
            verify = subprocess.run(
                [
                    python,
                    "-c",
                    "import torch; assert torch.cuda.is_available(), 'CUDA lost'",
                ],
                capture_output=True,
                text=True,
                check=False,
            )
            if verify.returncode != 0:
                print(
                    "⚠️  Batch install replaced NVIDIA torch — reinstalling CUDA torch..."
                )
                _install_jetson_torch(python)
        print("   All packages installed successfully (batch).")
        return

    print("   Batch install failed, falling back to individual package install...")

    # Parse requirements
    lines = req_file.read_text(encoding="utf-8").splitlines()
    packages = []
    for line in lines:
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        # Strip inline comments (e.g. "esp-ppq>=0.2.0  # ESP32 quantization")
        if " #" in line:
            line = line[: line.index(" #")].strip()
        if not line:
            continue
        packages.append(line)

    # Identify which packages can't be installed individually
    failed = []
    succeeded = []
    skipped = []
    total = len(packages)
    for i, pkg in enumerate(packages, 1):
        display_name = re.split(r"[>=<\[@ ]", pkg)[0]

        # On Jetson with CUDA, skip packages that would replace NVIDIA torch
        if jetson_with_cuda and display_name.lower() in TORCH_COUPLED_PKGS:
            if display_name.lower() == "torch":
                skipped.append(display_name)
                continue  # Already installed from NVIDIA
            # For torchaudio/torchcodec/torchvision, try --no-deps to avoid
            # pip pulling in CPU-only torch from PyPI
            result = _pip_install(
                [pkg],
                python=python,
                extra_args=["--no-deps", "-q"],
                capture_output=True,
                text=True,
            )
            if result.returncode != 0:
                skipped.append(display_name)
            else:
                succeeded.append(display_name)
            continue

        result = _pip_install(
            [pkg],
            python=python,
            extra_args=["-q"],
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            failed.append(display_name)
        else:
            succeeded.append(display_name)

    if skipped:
        print(
            f"   ℹ️  {len(skipped)} torch-coupled package(s) handled separately on Jetson:"
        )
        for name in skipped:
            print(f"      - {name}")
    if failed:
        print(f"   ⚠️  {len(failed)} package(s) skipped (no ARM64 wheel):")
        for name in failed:
            print(f"      - {name}")
        print(f"   {len(succeeded)}/{total} packages installed successfully.")
        print("   Non-critical features may be unavailable.")
    else:
        print(f"   All {total - len(skipped)} packages installed successfully.")


def start_native(
    model: Optional[str] = None,
    uri: Optional[str] = None,
    api_key: Optional[str] = None,
    ngrok: Optional[str] = None,
    gpu_type: str = "cpu",
) -> None:
    """Start ezlocalai in native mode (directly running start.py)."""

    # Check if already running
    if is_native_running():
        print(f"✅ ezlocalai is already running (native mode, PID {get_native_pid()})!")
        print(f"   ezLocalai is online at {_get_effective_url()}")
        return

    # Also check if Docker container is running
    try:
        if is_container_running():
            print("⚠️  ezlocalai Docker container is already running.")
            print("   Stop it first with 'ezlocalai stop' before using native mode.")
            return
    except Exception:
        pass  # Docker not available, which is expected in native mode

    # Find source directory
    source_dir = get_source_dir_for_native()
    if not source_dir:
        # Need to clone
        print("📦 ezlocalai source not found, cloning...")
        source_dir = clone_or_update_repo()
        if not source_dir:
            print("❌ Failed to get ezlocalai source code.")
            sys.exit(1)

    # Auto-update: check for new commits and pull if available
    auto_update_native(source_dir, gpu_type)

    # Load existing env or defaults
    env_vars = get_default_env()
    saved_env = load_env_file()
    env_vars.update(saved_env)

    # Apply command line overrides
    if model:
        env_vars["DEFAULT_MODEL"] = model
    if uri:
        env_vars["EZLOCALAI_URL"] = uri
    if api_key:
        env_vars["EZLOCALAI_API_KEY"] = api_key
    if ngrok:
        env_vars["NGROK_TOKEN"] = ngrok

    # Save updated env
    save_env_file(env_vars)

    # Install dependencies if needed — verify critical imports actually work
    deps_marker = STATE_DIR / ".deps_installed"
    needs_deps = not deps_marker.exists()
    if not needs_deps:
        # Verify critical modules are actually importable
        python = sys.executable
        check = subprocess.run(
            [python, "-c", "import uvicorn; import dotenv; import fastapi"],
            capture_output=True,
            check=False,
        )
        if check.returncode != 0:
            print("⚠️  Missing critical dependencies, reinstalling...")
            needs_deps = True

    if needs_deps:
        install_native_dependencies(source_dir, gpu_type)
        deps_marker.write_text("1")
    else:
        print(
            "✅ Dependencies verified (delete ~/.ezlocalai/.deps_installed to force reinstall)"
        )

    # Prepare data directories (same as Docker mode)
    data_dir = STATE_DIR / "data"
    for d in ["models", "outputs", "voices", "hf"]:
        (data_dir / d).mkdir(parents=True, exist_ok=True)

    # Ensure outputs dir exists in source_dir (app.py mounts it as static files)
    (source_dir / "outputs").mkdir(parents=True, exist_ok=True)

    # Build environment for the subprocess
    proc_env = os.environ.copy()
    proc_env.update(env_vars)
    proc_env["HOST"] = "0.0.0.0"
    proc_env["PORT"] = str(DEFAULT_PORT)
    # Point model/data paths to our state directory
    proc_env["MODELS_PATH"] = str(data_dir / "models")
    proc_env["HF_HOME"] = str(data_dir / "hf")

    mode = (
        "Jetson CUDA"
        if is_jetson()
        else ("NVIDIA GPU" if gpu_type == "nvidia" else "CPU")
    )
    print(f"\n🚀 Starting ezlocalai in native mode ({mode})...")
    print(f"   Model: {env_vars.get('DEFAULT_MODEL', 'default')}")
    print(f"   Source: {source_dir}")

    # Start start.py as a background process
    log_file = open(LOG_FILE, "a")
    proc = subprocess.Popen(
        [sys.executable, "start.py"],
        cwd=str(source_dir),
        env=proc_env,
        stdout=log_file,
        stderr=subprocess.STDOUT,
        start_new_session=True,  # Detach from terminal
    )

    # Save PID
    PID_FILE.write_text(str(proc.pid))
    print(f"   PID: {proc.pid}")

    # Wait for server to be ready
    print("\n⏳ Waiting for server to be ready...")
    max_wait = 300
    start_time = time.time()

    while time.time() - start_time < max_wait:
        # Check if process died
        if proc.poll() is not None:
            print(f"\n❌ ezlocalai process exited with code {proc.returncode}")
            print(f"   Check logs: cat {LOG_FILE}")
            PID_FILE.unlink(missing_ok=True)
            sys.exit(1)

        try:
            ezlocalai_url = _get_effective_url(env_vars)
            req = urllib.request.Request(f"{ezlocalai_url}/v1/models")
            with urllib.request.urlopen(req, timeout=5) as response:
                if response.status == 200:
                    print("\n✅ ezlocalai is ready! (native mode)")
                    print(f"\n   ezLocalai is online at {ezlocalai_url}")
                    print(f"   📖 API Docs: {ezlocalai_url}/docs")
                    print(f"   📋 Logs: {LOG_FILE}")
                    return
        except Exception:
            pass

        elapsed = int(time.time() - start_time)
        if elapsed % 10 == 0 and elapsed > 0:
            print(f"   Still loading... ({elapsed}s)")
        time.sleep(2)

    print("\n⚠️  Server started but not responding yet.")
    print("   This is normal for first-time model downloads.")
    print(f"   Check logs: tail -f {LOG_FILE}")
    print(f"\n   ezLocalai is online at {_get_effective_url(env_vars)}")


def _pid_alive(pid: int) -> bool:
    """Check if a process is still running."""
    try:
        os.kill(pid, 0)
        return True
    except (ProcessLookupError, PermissionError):
        return False


def _get_process_tree(root_pid: int) -> list[int]:
    """Get a list of all PIDs in a process tree (parent + all descendants).

    Uses /proc on Linux to walk the process tree without external tools.
    Returns pids in child-first order so children are killed before parents.
    """
    all_pids = {root_pid}

    try:
        # Read all /proc/[pid]/stat to build parent->children map
        children_map: dict[int, list[int]] = {}
        proc = Path("/proc")
        for entry in proc.iterdir():
            if not entry.name.isdigit():
                continue
            try:
                stat = (entry / "stat").read_text()
                # Format: pid (comm) state ppid ...
                # Need to handle comm with spaces/parens: find last ')' then split
                last_paren = stat.rfind(")")
                fields_after = stat[last_paren + 2 :].split()
                ppid = int(
                    fields_after[1]
                )  # ppid is field index 3, but index 1 after ')'
                pid_val = int(entry.name)
                children_map.setdefault(ppid, []).append(pid_val)
            except (OSError, ValueError, IndexError):
                continue

        # BFS from root_pid to find all descendants
        queue = [root_pid]
        while queue:
            current = queue.pop(0)
            for child in children_map.get(current, []):
                all_pids.add(child)
                queue.append(child)
    except OSError:
        pass

    # Return children first, parent last
    result = sorted(all_pids)
    result.reverse()
    return result


def stop_native() -> None:
    """Stop the native ezlocalai process and all its children.

    Since start_native() uses start_new_session=True, the process runs in its
    own process group. We kill the entire process group to ensure uvicorn,
    xllamacpp, and any other child processes are terminated.
    """
    pid = get_native_pid()
    if not pid:
        # PID file missing or stale — try to find and kill any orphaned processes
        _kill_orphaned_ezlocalai()
        return

    print(f"🛑 Stopping ezlocalai (PID {pid})...")

    # Collect all PIDs to kill: the tracked PID + all its descendants
    pids_to_kill = _get_process_tree(pid)

    # First try graceful SIGTERM on all of them
    for p in pids_to_kill:
        try:
            os.kill(p, signal.SIGTERM)
        except (ProcessLookupError, PermissionError):
            pass

    # Wait up to 3 seconds for graceful shutdown
    for _ in range(6):
        time.sleep(0.5)
        alive = [p for p in pids_to_kill if _pid_alive(p)]
        if not alive:
            break
    else:
        # Force kill anything still alive
        alive = [p for p in pids_to_kill if _pid_alive(p)]
        if alive:
            print(f"   Force killing {len(alive)} remaining process(es)...")
            for p in alive:
                try:
                    os.kill(p, signal.SIGKILL)
                except (ProcessLookupError, PermissionError):
                    pass
            time.sleep(1)

    # Also try process group kill as a belt-and-suspenders measure
    try:
        pgid = os.getpgid(pid)
        os.killpg(pgid, signal.SIGKILL)
    except (ProcessLookupError, PermissionError, OSError):
        pass

    # Final cleanup: kill anything still holding the port
    _kill_orphaned_ezlocalai()

    PID_FILE.unlink(missing_ok=True)
    print("✅ ezlocalai stopped")


def _kill_orphaned_ezlocalai() -> None:
    """Find and kill any orphaned ezlocalai-related processes.

    Catches processes that survived process group kill or were started
    outside the tracked PID (e.g. stale uvicorn or xllamacpp-server).
    Uses multiple strategies to ensure everything is stopped.
    """
    killed_any = False

    # Strategy 1: pgrep for known process patterns
    patterns = [
        "ezlocalai.*start\\.py",
        "uvicorn.*app:app",
        "xllamacpp",
        "python.*start\\.py",
    ]
    my_pid = str(os.getpid())

    for pattern in patterns:
        try:
            result = subprocess.run(
                ["pgrep", "-f", pattern],
                capture_output=True,
                text=True,
                check=False,
            )
            if result.returncode == 0 and result.stdout.strip():
                for p in result.stdout.strip().split("\n"):
                    p = p.strip()
                    if p and p != my_pid:
                        try:
                            os.kill(int(p), signal.SIGKILL)
                            killed_any = True
                        except (ProcessLookupError, PermissionError, ValueError):
                            pass
        except FileNotFoundError:
            break  # pgrep not available

    # Strategy 2: kill whatever is holding our port
    try:
        result = subprocess.run(
            ["fuser", f"{DEFAULT_PORT}/tcp"],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0 and result.stdout.strip():
            pids = result.stdout.strip().split()
            for p in pids:
                p = p.strip()
                if p and p != my_pid:
                    try:
                        os.kill(int(p), signal.SIGKILL)
                        killed_any = True
                    except (ProcessLookupError, PermissionError, ValueError):
                        pass
    except FileNotFoundError:
        # fuser not available, try lsof
        try:
            result = subprocess.run(
                ["lsof", "-ti", f":{DEFAULT_PORT}"],
                capture_output=True,
                text=True,
                check=False,
            )
            if result.returncode == 0 and result.stdout.strip():
                for p in result.stdout.strip().split("\n"):
                    p = p.strip()
                    if p and p != my_pid:
                        try:
                            os.kill(int(p), signal.SIGKILL)
                            killed_any = True
                        except (ProcessLookupError, PermissionError, ValueError):
                            pass
        except FileNotFoundError:
            pass  # Neither fuser nor lsof available

    if killed_any:
        time.sleep(1)  # Give OS time to release the port


def show_native_status() -> None:
    """Show status of native ezlocalai process."""
    pid = get_native_pid()
    env_vars = load_env_file()
    if not env_vars:
        env_vars = get_default_env()

    ezlocalai_url = _get_effective_url(env_vars)

    if pid:
        print(f"✅ ezlocalai is running (native mode, PID {pid})")
        print(f"\n   ezLocalai is online at {ezlocalai_url}")

        # Show loaded model from API
        try:
            req = urllib.request.Request(f"{ezlocalai_url}/v1/models")
            with urllib.request.urlopen(req, timeout=5) as response:
                data = json.loads(response.read().decode())
                models = [m.get("id") for m in data.get("data", [])]
                if models:
                    print(f"\n   🧠 Active model: {models[0]}")
        except Exception:
            print("   ⚠️  Server not responding (may still be loading)")
    else:
        print("ℹ️  ezlocalai is not running in native mode")
        print("   Run 'ezlocalai start --native' to start in native mode")

    # Show config
    print(f"\n   ⚙️  Configuration:")
    models = env_vars.get("DEFAULT_MODEL", "unsloth/Qwen3.5-4B-GGUF")
    configured = [m.strip() for m in models.split(",")]
    print(f"      LLM models:")
    for m in configured:
        print(f"        - {m}")

    whisper = env_vars.get("WHISPER_MODEL", "large-v3")
    print(f"      Speech-to-text: {whisper if whisper else 'disabled'}")
    img_model = env_vars.get("IMG_MODEL", "")
    print(f"      Image generation: {img_model if img_model else 'disabled'}")
    video_model = env_vars.get("VIDEO_MODEL", "")
    print(f"      Video generation: {video_model if video_model else 'disabled'}")

    print(f"\n   📋 Logs: {LOG_FILE}")


def show_native_logs(follow: bool = False) -> None:
    """Show logs from native mode."""
    if not LOG_FILE.exists():
        print("ℹ️  No log file found")
        return

    if follow:
        try:
            subprocess.run(["tail", "-f", str(LOG_FILE)], check=False)
        except KeyboardInterrupt:
            pass
    else:
        # Show last 100 lines
        try:
            subprocess.run(["tail", "-n", "100", str(LOG_FILE)], check=False)
        except FileNotFoundError:
            # Windows fallback
            lines = LOG_FILE.read_text(encoding="utf-8", errors="replace").splitlines()
            for line in lines[-100:]:
                print(line)


# ── Docker mode functions ───────────────────────────────────────────────────


def start_container(
    model: Optional[str] = None,
    uri: Optional[str] = None,
    api_key: Optional[str] = None,
    ngrok: Optional[str] = None,
    gpu_type: str = "cpu",
) -> None:
    """Start the ezlocalai container via Docker Compose."""

    # Check if already running
    if is_container_running():
        print(f"ezLocalai is online at {_get_effective_url()}")
        return

    # Load existing env or defaults
    env_vars = get_default_env()
    saved_env = load_env_file()
    env_vars.update(saved_env)

    # Apply command line overrides
    if model:
        env_vars["DEFAULT_MODEL"] = model
    if uri:
        env_vars["EZLOCALAI_URL"] = uri
    if api_key:
        env_vars["EZLOCALAI_API_KEY"] = api_key
    if ngrok:
        env_vars["NGROK_TOKEN"] = ngrok

    # Save GPU type so stop/restart can find the right compose file
    env_vars["_GPU_TYPE"] = gpu_type

    # Save to CLI state
    save_env_file(env_vars)

    # Resolve source directory (contains compose files)
    source_dir = _get_compose_source_dir()
    if not source_dir:
        print("📦 ezlocalai source not found, cloning...")
        source_dir = clone_or_update_repo()
        if not source_dir:
            print("❌ Failed to get ezlocalai source code.")
            sys.exit(1)

    compose_file = _get_compose_file(gpu_type)
    compose_path = source_dir / compose_file
    if not compose_path.exists():
        print(f"❌ Compose file not found: {compose_path}")
        sys.exit(1)

    # Write env vars to source dir so Docker Compose picks them up
    _write_compose_env(source_dir, env_vars)

    # Auto-update: pull latest image or rebuild GPU image if source changed
    auto_update_docker(gpu_type)

    # Ensure image exists — build if needed for GPU types
    if gpu_type == "nvidia":
        if not cuda_image_exists():
            print("\n🔨 CUDA image not found, building from source...")
            if not build_cuda_image():
                print("❌ Failed to build CUDA image. Falling back to CPU mode.")
                gpu_type = "cpu"
                compose_file = _get_compose_file(gpu_type)
                env_vars["_GPU_TYPE"] = gpu_type
                save_env_file(env_vars)
                _write_compose_env(source_dir, env_vars)
    elif gpu_type == "jetson":
        if not jetson_image_exists():
            print("\n🔨 Jetson image not found, building from source...")
            print("   (This is a one-time build, may take 15-30 minutes)")
            if not build_jetson_image():
                print("❌ Failed to build Jetson image.")
                sys.exit(1)
    elif gpu_type == "amd":
        if not rocm_image_exists():
            print("\n🔨 ROCm image not found, building from source...")
            if not build_rocm_image():
                print("❌ Failed to build ROCm image. Falling back to CPU mode.")
                gpu_type = "cpu"
                compose_file = _get_compose_file(gpu_type)
                env_vars["_GPU_TYPE"] = gpu_type
                save_env_file(env_vars)
                _write_compose_env(source_dir, env_vars)
    else:
        # CPU image: pull from DockerHub
        print(f"\n📦 Pulling latest CPU image...")
        subprocess.run(
            ["docker", "pull", DOCKER_IMAGE],
            capture_output=True,
            text=True,
            check=False,
        )

    # Start via Docker Compose
    mode_names = {
        "nvidia": "NVIDIA GPU",
        "jetson": "Jetson GPU",
        "amd": "AMD GPU",
        "cpu": "CPU",
    }
    mode = mode_names.get(gpu_type, "CPU")
    print(f"\n🚀 Starting ezlocalai ({mode} mode)...")
    print(f"   Model: {env_vars.get('DEFAULT_MODEL', 'default')}")
    print(f"   Compose: {compose_file}")

    # Build compose env for Jetson build args
    compose_env = os.environ.copy()
    for key, value in env_vars.items():
        if value and not key.startswith("_"):
            compose_env[key] = value

    result = subprocess.run(
        ["docker", "compose", "-f", compose_file, "up", "-d"],
        cwd=source_dir,
        env=compose_env,
        capture_output=True,
        text=True,
        check=False,
    )

    if result.returncode != 0:
        print(f"❌ Failed to start container: {result.stderr}")
        sys.exit(1)

    # Wait for container to be healthy
    print("\n⏳ Waiting for server to be ready...")
    max_wait = 300  # 5 minutes for model download
    start_time = time.time()

    while time.time() - start_time < max_wait:
        try:
            import urllib.request

            ezlocalai_url = _get_effective_url(env_vars)
            req = urllib.request.Request(f"{ezlocalai_url}/v1/models")
            with urllib.request.urlopen(req, timeout=5) as response:
                if response.status == 200:
                    print("\n✅ ezlocalai is ready!")
                    print(f"\n   ezLocalai is online at {ezlocalai_url}")
                    print(f"   📖 API Docs: {ezlocalai_url}/docs")
                    return
        except Exception:
            pass

        # Show progress
        elapsed = int(time.time() - start_time)
        if elapsed % 10 == 0 and elapsed > 0:
            print(f"   Still loading... ({elapsed}s)")
            # After 60s, show a tail of the logs to help debug
            if elapsed == 60:
                print("\n   📋 Recent container logs:")
                log_result = subprocess.run(
                    ["docker", "logs", "--tail", "15", CONTAINER_NAME],
                    capture_output=True,
                    text=True,
                    check=False,
                )
                if log_result.stdout:
                    for line in log_result.stdout.strip().splitlines()[-10:]:
                        print(f"      {line}")
                elif log_result.stderr:
                    for line in log_result.stderr.strip().splitlines()[-10:]:
                        print(f"      {line}")
                print()
        time.sleep(2)

    print("\n⚠️  Container started but server not responding yet.")
    print("   This is normal for first-time model downloads.")
    print(f"   Check logs with: ezlocalai logs")
    print(f"\n   ezLocalai is online at {_get_effective_url(env_vars)}")


def _cleanup_all_ezlocalai_containers(
    source_dir: Optional[Path] = None, compose_file: str = ""
) -> None:
    """Remove all stopped ezlocalai containers (compose and docker-run)."""
    # Remove any stopped containers matching the name
    result = subprocess.run(
        [
            "docker",
            "ps",
            "-a",
            "-q",
            "-f",
            f"name={CONTAINER_NAME}",
            "-f",
            "status=exited",
        ],
        capture_output=True,
        text=True,
        check=False,
    )
    stopped_ids = [
        cid.strip() for cid in result.stdout.strip().splitlines() if cid.strip()
    ]
    if stopped_ids:
        for cid in stopped_ids:
            subprocess.run(
                ["docker", "rm", "-f", cid],
                capture_output=True,
                check=False,
            )


def stop_container() -> None:
    """Stop and remove the ezlocalai container via Docker Compose."""
    source_dir = _get_compose_source_dir()

    # Determine GPU type from saved env
    saved_env = load_env_file()
    gpu_type = saved_env.get("_GPU_TYPE", "cpu")
    compose_file = _get_compose_file(gpu_type)

    # Always clean up any old containers (from prior docker-run or compose)
    _cleanup_all_ezlocalai_containers(source_dir, compose_file)

    if not is_container_running():
        print("ℹ️  ezlocalai is not running")
        return

    print("🛑 Stopping ezlocalai...")

    # Try Docker Compose down if we have a source dir
    if source_dir and (source_dir / compose_file).exists():
        result = subprocess.run(
            ["docker", "compose", "-f", compose_file, "down", "--remove-orphans"],
            cwd=source_dir,
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode != 0:
            print(f"⚠️  Compose down failed, trying direct stop...")
            subprocess.run(
                ["docker", "stop", CONTAINER_NAME],
                capture_output=True,
                check=False,
            )
    else:
        # Fallback: direct docker stop/rm
        subprocess.run(
            ["docker", "stop", CONTAINER_NAME],
            capture_output=True,
            check=False,
        )
        subprocess.run(
            ["docker", "rm", "-f", CONTAINER_NAME],
            capture_output=True,
            check=False,
        )

    # Wait until the container is actually gone
    if _wait_for_container_stop(timeout=30):
        print("✅ ezlocalai stopped")
    else:
        # Force kill anything remaining
        subprocess.run(
            ["docker", "rm", "-f", CONTAINER_NAME],
            capture_output=True,
            check=False,
        )
        # Also kill any compose-named containers
        result = subprocess.run(
            ["docker", "ps", "-q", "-f", f"name={CONTAINER_NAME}"],
            capture_output=True,
            text=True,
            check=False,
        )
        for cid in result.stdout.strip().splitlines():
            if cid:
                subprocess.run(
                    ["docker", "rm", "-f", cid],
                    capture_output=True,
                    check=False,
                )
        print("✅ ezlocalai stopped (forced)")


def show_status() -> None:
    """Show ezlocalai status."""
    status = get_container_status()

    # Load config (saved or defaults)
    env_vars = load_env_file()
    if not env_vars:
        env_vars = get_default_env()

    ezlocalai_url = _get_effective_url(env_vars)

    if not status:
        print("ℹ️  ezlocalai container not found")
        print("   Run 'ezlocalai start' to start the server")
    elif is_container_running():
        print(f"✅ ezlocalai is running")
        print(f"   Status: {status}")
        print(f"\n   ezLocalai is online at {ezlocalai_url}")

        # Show loaded model from API
        try:
            import urllib.request
            import json

            req = urllib.request.Request(f"{ezlocalai_url}/v1/models")
            with urllib.request.urlopen(req, timeout=5) as response:
                data = json.loads(response.read().decode())
                models = [m.get("id") for m in data.get("data", [])]
                if models:
                    print(f"\n   🧠 Active model: {models[0]}")
        except Exception:
            pass
    else:
        print(f"❌ ezlocalai is not running")
        print(f"   Status: {status}")
        print("   Run 'ezlocalai start' to start the server")

    # Always show configuration
    print(f"\n   ⚙️  Configuration:")

    # LLM models
    models = env_vars.get("DEFAULT_MODEL", None)
    if not models:
        # Get models from the list
        env = get_default_env()
        models = env.get("DEFAULT_MODEL", "unsloth/Qwen3-4B-Instruct-2507-GGUF")
    configured = [m.strip() for m in models.split(",")]
    print(f"      LLM models:")
    for m in configured:
        print(f"        - {m}")

    # Whisper
    whisper = env_vars.get("WHISPER_MODEL", "large-v3")
    if whisper:
        print(f"      Speech-to-text: {whisper}")
    else:
        print(f"      Speech-to-text: disabled")

    # Image model
    img_model = env_vars.get("IMG_MODEL", "")
    if img_model:
        print(f"      Image generation: {img_model}")
    else:
        print(f"      Image generation: disabled")

    # Video model
    video_model = env_vars.get("VIDEO_MODEL", "")
    if video_model:
        print(f"      Video generation: {video_model}")
    else:
        print(f"      Video generation: disabled")

    print(f"\n   💡 To change settings:")
    print(f"      ezlocalai start --model <model>")
    print(f"      Or edit ~/.ezlocalai/.env")


def show_logs(follow: bool = False) -> None:
    """Show container logs."""
    if not get_container_status():
        print("ℹ️  ezlocalai container not found")
        return

    cmd = ["docker", "logs"]
    if follow:
        cmd.append("-f")
    cmd.extend(["--tail", "100", CONTAINER_NAME])

    try:
        subprocess.run(cmd, check=False)
    except KeyboardInterrupt:
        pass


def _check_for_updates(source_dir: Path) -> bool:
    """Check if the remote repo has new commits.

    Returns True if updates are available, False otherwise.
    """
    if not (source_dir / ".git").exists():
        return False

    _ensure_git_access(source_dir)

    # Fetch remote quietly
    result = subprocess.run(
        ["git", "fetch", "--quiet"],
        cwd=source_dir,
        capture_output=True,
        check=False,
    )
    if result.returncode != 0:
        return False

    # Check if local is behind remote
    result = subprocess.run(
        ["git", "rev-list", "--count", "HEAD..@{u}"],
        cwd=source_dir,
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        return False

    behind = int(result.stdout.strip() or "0")
    return behind > 0


def auto_update_native(source_dir: Path, gpu_type: str = "cpu") -> None:
    """Auto-update native mode: git pull + pip install -e . + reinstall deps if needed."""
    if not _check_for_updates(source_dir):
        return

    print("📦 New version available, auto-updating...")

    # Git pull
    result = subprocess.run(
        ["git", "pull", "--ff-only"],
        cwd=source_dir,
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        print(f"   ⚠️  git pull failed: {result.stderr.strip()}")
        print("   Continuing with current version...")
        return

    output = result.stdout.strip()
    print(f"   ✅ Updated: {output.splitlines()[-1] if output else 'done'}")

    # Check if requirements changed
    result = subprocess.run(
        ["git", "diff", "HEAD~1", "--name-only"],
        cwd=source_dir,
        capture_output=True,
        text=True,
        check=False,
    )
    changed_files = result.stdout.strip().splitlines() if result.returncode == 0 else []

    deps_changed = any(
        f in changed_files
        for f in [
            "requirements.txt",
            "cuda-requirements.txt",
            "pyproject.toml",
            "setup.py",
        ]
    )

    # Reinstall package
    python = sys.executable
    _pip_install(
        ["-e", "."],
        python=python,
        extra_args=["-q"],
        capture_output=True,
        cwd=source_dir,
    )

    # Reinstall deps if requirements changed
    if deps_changed:
        print("   Dependencies changed, reinstalling...")
        deps_marker = STATE_DIR / ".deps_installed"
        deps_marker.unlink(missing_ok=True)
        install_native_dependencies(source_dir, gpu_type)
        deps_marker.write_text("1")

    print("   ✅ Auto-update complete\n")


def auto_update_docker(gpu_type: str = "cpu") -> None:
    """Auto-update Docker mode: pull latest image, rebuild CUDA/ROCm if source changed."""
    # For CPU image: quick pull check
    if gpu_type == "cpu":
        print("🔄 Checking for image updates...")
        result = subprocess.run(
            ["docker", "pull", DOCKER_IMAGE],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode == 0 and "up to date" not in result.stdout.lower():
            print("   ✅ New image pulled")
        return

    # For GPU images: check if source repo has updates and rebuild if so
    source_dir = get_ezlocalai_source_dir()
    if not source_dir:
        if REPO_DIR.exists() and is_ezlocalai_folder(REPO_DIR):
            source_dir = REPO_DIR
        else:
            return  # No source dir, can't auto-update GPU images

    if not _check_for_updates(source_dir):
        return

    print("📦 New version available, updating source and rebuilding image...")
    result = subprocess.run(
        ["git", "pull", "--ff-only"],
        cwd=source_dir,
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        print(f"   ⚠️  git pull failed, continuing with current image...")
        return

    # Rebuild the appropriate GPU image
    if gpu_type == "nvidia":
        print("   Rebuilding CUDA image...")
        build_cuda_image()
    elif gpu_type == "jetson":
        print("   Rebuilding Jetson image...")
        build_jetson_image()
    elif gpu_type == "amd":
        print("   Rebuilding ROCm image...")
        build_rocm_image()

    print("   ✅ Auto-update complete\n")


def update_native() -> None:
    """Update native mode installation: git pull + pip install -e . + reinstall deps."""
    print("📦 Updating ezlocalai (native mode)...")

    source_dir = get_ezlocalai_source_dir()
    if not source_dir:
        # Check if REPO_DIR exists
        if REPO_DIR.exists() and is_ezlocalai_folder(REPO_DIR):
            source_dir = REPO_DIR
        else:
            print("❌ Cannot find ezlocalai source directory.")
            print("   Run from within the ezlocalai folder, or clone it first.")
            return

    print(f"   Source: {source_dir}")

    _ensure_git_access(source_dir)

    # Git pull
    print("\n📥 Pulling latest changes...")
    result = subprocess.run(
        ["git", "pull", "--ff-only"],
        cwd=source_dir,
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        print(f"   ⚠️  git pull failed: {result.stderr.strip()}")
        print("   Continuing with current code...")
    else:
        output = result.stdout.strip()
        if "Already up to date" in output:
            print("   Already up to date.")
        else:
            print(f"   ✅ Updated: {output.splitlines()[-1] if output else 'done'}")

    # Reinstall package in editable mode
    print("\n📦 Reinstalling ezlocalai...")
    python = sys.executable
    result = _pip_install(
        ["-e", "."],
        python=python,
        extra_args=["-q"],
        capture_output=True,
        text=True,
        cwd=source_dir,
    )
    if result.returncode != 0:
        print(f"   ❌ install -e . failed: {result.stderr.strip()}")
        return
    print("   ✅ ezlocalai reinstalled")

    # Reinstall dependencies (clear marker to force reinstall)
    deps_marker = STATE_DIR / ".deps_installed"
    deps_marker.unlink(missing_ok=True)

    # Detect GPU type for dependency install
    gpu_type = "cpu"
    if has_nvidia_gpu() or (is_jetson() and has_jetson_cuda()):
        gpu_type = "nvidia"
    elif has_amd_gpu() and has_rocm_support():
        gpu_type = "amd"

    install_native_dependencies(source_dir, gpu_type)
    deps_marker.write_text("1")

    print("\n✅ Update complete!")
    print("   Run 'ezlocalai restart' to use the new version.")


def update_images() -> None:
    """Pull latest CPU image and rebuild CUDA/ROCm/Jetson image (Docker mode)."""
    print("📦 Updating ezlocalai (Docker mode)...")

    # Check if we're in the ezlocalai source folder
    local_source = get_ezlocalai_source_dir()
    if local_source:
        print(f"ℹ️  Running from ezlocalai source folder: {local_source}")

    # Jetson: skip CPU pull (no ARM64 manifest on DockerHub), rebuild Jetson image
    if is_jetson():
        print(f"\n🔨 Jetson detected — rebuilding Jetson Docker image from source...")
        # Update source first
        source_dir = local_source or (REPO_DIR if REPO_DIR.exists() else None)
        if source_dir:
            print("📥 Pulling latest source...")
            subprocess.run(
                ["git", "pull", "--ff-only"],
                cwd=source_dir,
                capture_output=True,
                text=True,
                check=False,
            )
        if build_jetson_image():
            print("   ✅ Jetson image rebuilt")
        else:
            print("   ⚠️  Failed to build Jetson image")
        print("\n✅ Update complete!")
        print("   Run 'ezlocalai restart' to use the new version.")
        return

    # Pull CPU image from DockerHub
    print(f"\n📥 Pulling {DOCKER_IMAGE}...")
    result = subprocess.run(
        ["docker", "pull", DOCKER_IMAGE],
        check=False,
    )
    if result.returncode != 0:
        print(f"   ⚠️  Failed to pull CPU image")
    else:
        print(f"   ✅ CPU image updated")

    # Build CUDA image from source if NVIDIA GPU available
    if has_nvidia_gpu():
        print(f"\n🔨 Building CUDA image from source...")
        if build_cuda_image():
            print("   ✅ CUDA image rebuilt")
        else:
            print("   ⚠️  Failed to build CUDA image")
    else:
        print("\nℹ️  No NVIDIA GPU detected, skipping CUDA image build")

    # Build ROCm image from source if AMD GPU available
    if has_amd_gpu():
        print(f"\n🔨 Building ROCm image from source...")
        if build_rocm_image():
            print("   ✅ ROCm image rebuilt")
        else:
            print("   ⚠️  Failed to build ROCm image")
    else:
        print("\nℹ️  No AMD GPU detected, skipping ROCm image build")

    print("\n✅ Update complete!")
    print("   Run 'ezlocalai restart' to use the new version.")


def encode_image_to_base64(image_path: str) -> str:
    """Encode a local image file to base64."""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def get_image_mime_type(image_path: str) -> str:
    """Get the MIME type based on file extension."""
    ext = Path(image_path).suffix.lower()
    mime_types = {
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".png": "image/png",
        ".gif": "image/gif",
        ".webp": "image/webp",
        ".bmp": "image/bmp",
    }
    return mime_types.get(ext, "image/jpeg")


def get_video_mime_type(video_path: str) -> str:
    """Get the MIME type based on video file extension."""
    ext = Path(video_path).suffix.lower()
    mime_types = {
        ".mp4": "video/mp4",
        ".webm": "video/webm",
        ".avi": "video/x-msvideo",
        ".mov": "video/quicktime",
        ".mkv": "video/x-matroska",
        ".m4v": "video/mp4",
        ".wmv": "video/x-ms-wmv",
        ".flv": "video/x-flv",
    }
    return mime_types.get(ext, "video/mp4")


def is_video_file(path: str) -> bool:
    """Check if the path is a video file based on extension."""
    ext = Path(path).suffix.lower()
    video_extensions = {".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv", ".flv"}
    return ext in video_extensions


def is_url(path: str) -> bool:
    """Check if the path is a URL."""
    return path.startswith("http://") or path.startswith("https://")


def send_prompt(
    prompt_text: str,
    model: Optional[str] = None,
    temperature: Optional[float] = None,
    top_p: Optional[float] = None,
    image_path: Optional[str] = None,
    video_path: Optional[str] = None,
    show_stats: bool = False,
) -> None:
    """Send a prompt to the ezlocalai server and print the response."""
    # Check if server is running (Docker or native)
    running = is_native_running()
    if not running:
        try:
            running = is_container_running()
        except Exception:
            pass
    if not running:
        print("❌ ezlocalai is not running. Start it with: ezlocalai start")
        sys.exit(1)

    # Load env to get API key and default model if set
    env_vars = load_env_file()
    api_key = env_vars.get("EZLOCALAI_API_KEY", "")
    ezlocalai_url = _get_effective_url(env_vars)

    # Get default model from running server if not specified
    if not model:
        try:
            models_url = f"{ezlocalai_url}/v1/models"
            req = urllib.request.Request(models_url)
            with urllib.request.urlopen(req, timeout=10) as response:
                models_data = json.loads(response.read().decode("utf-8"))
                if models_data.get("data"):
                    model = models_data["data"][0].get("id")
        except Exception:
            # Fall back to configured default model
            default_models = env_vars.get("DEFAULT_MODEL", "")
            if default_models:
                model = default_models.split(",")[0].strip()

    # Build the messages array
    content = []

    # Handle image if provided
    if image_path:
        if is_url(image_path):
            # Image URL - use directly
            content.append(
                {
                    "type": "image_url",
                    "image_url": {"url": image_path},
                }
            )
        else:
            # Local file - encode to base64
            image_file = Path(image_path)
            if not image_file.exists():
                print(f"❌ Image file not found: {image_path}")
                sys.exit(1)

            mime_type = get_image_mime_type(image_path)
            base64_image = encode_image_to_base64(image_path)
            content.append(
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:{mime_type};base64,{base64_image}"},
                }
            )

    # Handle video if provided
    if video_path:
        if is_url(video_path):
            # Video URL - pass directly, server will handle frame extraction
            content.append(
                {
                    "type": "video_url",
                    "video_url": {"url": video_path},
                }
            )
            print(f"📹 Processing video from URL...")
        else:
            # Local file - check it exists and pass the path
            # The server will extract frames from the local path
            video_file = Path(video_path)
            if not video_file.exists():
                print(f"❌ Video file not found: {video_path}")
                sys.exit(1)

            # For local files, we encode to base64 data URL
            print(f"📹 Processing video: {video_path}")
            mime_type = get_video_mime_type(video_path)
            base64_video = encode_image_to_base64(
                video_path
            )  # Same function works for any file
            content.append(
                {
                    "type": "video_url",
                    "video_url": {"url": f"data:{mime_type};base64,{base64_video}"},
                }
            )

    # Add the text prompt
    content.append({"type": "text", "text": prompt_text})

    messages = [{"role": "user", "content": content}]

    # Build the request payload
    payload = {
        "messages": messages,
        "stream": True,
    }

    # Add optional parameters
    if model:
        payload["model"] = model
    if temperature is not None:
        payload["temperature"] = temperature
    if top_p is not None:
        payload["top_p"] = top_p

    # Build headers
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    # Make the request
    url = f"{ezlocalai_url}/v1/chat/completions"

    try:
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(url, data=data, headers=headers, method="POST")

        start_time = time.time()
        print()  # Start on a new line for streaming output

        # Track streaming statistics
        completion_tokens = 0
        prompt_tokens = 0
        actual_model = model or "unknown"
        full_content = ""

        with urllib.request.urlopen(req, timeout=300) as response:
            # Process Server-Sent Events (SSE) stream
            buffer = ""
            for chunk in iter(lambda: response.read(1024).decode("utf-8"), ""):
                if not chunk:
                    break
                buffer += chunk

                # Process complete SSE messages
                while "\n\n" in buffer or "\r\n\r\n" in buffer:
                    # Split on either \n\n or \r\n\r\n
                    if "\r\n\r\n" in buffer:
                        message, buffer = buffer.split("\r\n\r\n", 1)
                    else:
                        message, buffer = buffer.split("\n\n", 1)

                    # Process each line in the message
                    for line in message.split("\n"):
                        line = line.strip()
                        if line.startswith("data: "):
                            json_data = line[6:]  # Remove "data: " prefix

                            # Check for stream end
                            if json_data == "[DONE]":
                                continue

                            try:
                                chunk_data = json.loads(json_data)

                                # Get model from first chunk
                                if chunk_data.get("model"):
                                    actual_model = chunk_data["model"]

                                # Extract and print delta content
                                choices = chunk_data.get("choices", [])
                                if choices:
                                    delta = choices[0].get("delta", {})
                                    content = delta.get("content", "")
                                    if content:
                                        print(content, end="", flush=True)
                                        full_content += content
                                        completion_tokens += (
                                            1  # Approximate token count
                                        )

                                # Get usage from final chunk if available
                                usage = chunk_data.get("usage")
                                if usage:
                                    prompt_tokens = usage.get("prompt_tokens", 0)
                                    completion_tokens = usage.get(
                                        "completion_tokens", completion_tokens
                                    )

                            except json.JSONDecodeError:
                                pass  # Skip malformed JSON

        elapsed = time.time() - start_time
        print()  # End with newline

        if not full_content:
            print("⚠️  Empty response received")

        # Show stats if requested
        if show_stats:
            total_tokens = prompt_tokens + completion_tokens

            print(f"\n{'─' * 50}")
            print(f"📊 Statistics")
            print(f"{'─' * 50}")
            print(f"   Model: {actual_model}")
            print(f"   Prompt tokens: {prompt_tokens:,}")
            print(f"   Completion tokens: {completion_tokens:,}")
            print(f"   Total tokens: {total_tokens:,}")
            print(f"   Total time: {elapsed:.1f}s")

            if completion_tokens > 0:
                overall_speed = completion_tokens / elapsed
                print(f"   Overall speed: {overall_speed:.1f} tok/s")

    except urllib.error.HTTPError as e:
        error_body = e.read().decode("utf-8") if e.fp else ""
        print(f"❌ HTTP Error {e.code}: {e.reason}")
        if error_body:
            try:
                error_json = json.loads(error_body)
                print(f"   {error_json.get('detail', error_body)}")
            except json.JSONDecodeError:
                print(f"   {error_body}")
        sys.exit(1)
    except urllib.error.URLError as e:
        print(f"❌ Connection error: {e.reason}")
        print("   Is ezlocalai running? Check with: ezlocalai status")
        sys.exit(1)
    except TimeoutError:
        print("❌ Request timed out after 300 seconds")
        sys.exit(1)


def main():
    """Main entry point for the CLI."""
    parser = argparse.ArgumentParser(
        description="ezlocalai - Local AI inference made easy",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  ezlocalai start                           Start with default settings (Docker)
  ezlocalai start --native                  Start in native mode (no Docker)
  ezlocalai start --model unsloth/Qwen3-4B  Start with specific model
  ezlocalai restart                         Restart the server
  ezlocalai stop                            Stop the server
  ezlocalai status                          Show server status
  ezlocalai logs -f                         Follow logs
  ezlocalai prompt "Hello, world!"          Send a prompt to the AI
  ezlocalai prompt "What's in this image?" -image ./photo.jpg
  ezlocalai prompt "Describe this video" -video ./clip.mp4

Modes:
  Docker (default on x86_64): Runs in a Docker container
  Native (default on ARM64):  Runs start.py directly, auto-detected on ARM64/Jetson

Environment:
  Configuration is stored in ~/.ezlocalai/.env
  Data (models, outputs) is stored in ~/.ezlocalai/data/
""",
    )

    parser.add_argument(
        "-v",
        "--version",
        action="version",
        version=f"ezlocalai {__version__}",
    )

    subparsers = parser.add_subparsers(dest="command", help="Commands")

    # Start command
    start_parser = subparsers.add_parser("start", help="Start ezlocalai server")
    start_parser.add_argument(
        "--model",
        "-m",
        help="Model(s) to load (comma-separated for multiple)",
        default=None,
    )
    start_parser.add_argument(
        "--uri",
        help="Server URI (default: http://localhost:8091)",
        default=None,
    )
    start_parser.add_argument(
        "--api-key",
        help="API key for authentication",
        default=None,
    )
    start_parser.add_argument(
        "--ngrok",
        help="Ngrok token for public URL",
        default=None,
    )
    start_parser.add_argument(
        "--native",
        action="store_true",
        help="Run in native mode (directly, without Docker). Auto-enabled on ARM64/Jetson.",
    )

    # Stop command
    subparsers.add_parser("stop", help="Stop ezlocalai server")

    # Restart command
    restart_parser = subparsers.add_parser("restart", help="Restart ezlocalai server")
    restart_parser.add_argument(
        "--model",
        "-m",
        help="Model(s) to load (comma-separated for multiple)",
        default=None,
    )
    restart_parser.add_argument(
        "--uri",
        help="Server URI (default: http://localhost:8091)",
        default=None,
    )
    restart_parser.add_argument(
        "--api-key",
        help="API key for authentication",
        default=None,
    )
    restart_parser.add_argument(
        "--ngrok",
        help="Ngrok token for public URL",
        default=None,
    )
    restart_parser.add_argument(
        "--native",
        action="store_true",
        help="Run in native mode (directly, without Docker). Auto-enabled on ARM64/Jetson.",
    )

    # Status command
    subparsers.add_parser("status", help="Show server status")

    # Update command
    subparsers.add_parser("update", help="Pull latest Docker image")

    # Logs command
    logs_parser = subparsers.add_parser("logs", help="Show server logs")
    logs_parser.add_argument(
        "-f",
        "--follow",
        action="store_true",
        help="Follow log output",
    )

    # Prompt command
    prompt_parser = subparsers.add_parser(
        "prompt", help="Send a prompt to the AI and get a response"
    )
    prompt_parser.add_argument(
        "text",
        help="The prompt text to send to the AI",
    )
    prompt_parser.add_argument(
        "-m",
        "--model",
        help="Model to use for the prompt",
        default=None,
    )
    prompt_parser.add_argument(
        "-temp",
        "--temperature",
        type=float,
        help="Temperature for response generation (0.0-2.0)",
        default=None,
    )
    prompt_parser.add_argument(
        "-tp",
        "--top-p",
        type=float,
        help="Top-p (nucleus) sampling parameter (0.0-1.0)",
        default=None,
    )
    prompt_parser.add_argument(
        "-image",
        "--image",
        help="Path to an image file or URL to include with the prompt",
        default=None,
    )
    prompt_parser.add_argument(
        "-video",
        "--video",
        help="Path to a video file or URL to include with the prompt (frames will be extracted)",
        default=None,
    )
    prompt_parser.add_argument(
        "-stats",
        "--stats",
        action="store_true",
        help="Show statistics (tokens, speed, timing) after the response",
    )

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(0)

    # Determine if we should use native mode
    native = False
    if args.command in ("start", "restart"):
        native = should_use_native_mode(force_native=args.native)

    # For stop/status/logs/update, auto-detect which mode is active
    if args.command in ("stop", "status", "logs", "update"):
        native = is_native_running() or should_use_native_mode()

    # Check prerequisites for start/restart
    gpu_type = "cpu"
    if args.command in ("start", "restart"):
        _, gpu_type = check_prerequisites(native_mode=native)
        print()

    # Execute command
    if args.command == "start":
        if native:
            start_native(
                model=args.model,
                uri=args.uri,
                api_key=args.api_key,
                ngrok=args.ngrok,
                gpu_type=gpu_type,
            )
        else:
            start_container(
                model=args.model,
                uri=args.uri,
                api_key=args.api_key,
                ngrok=args.ngrok,
                gpu_type=gpu_type,
            )

    elif args.command == "stop":
        if native:
            stop_native()
        else:
            stop_container()

    elif args.command == "restart":
        if native:
            stop_native()
            time.sleep(2)
            start_native(
                model=args.model,
                uri=args.uri,
                api_key=args.api_key,
                ngrok=args.ngrok,
                gpu_type=gpu_type,
            )
        else:
            stop_container()
            start_container(
                model=args.model,
                uri=args.uri,
                api_key=args.api_key,
                ngrok=args.ngrok,
                gpu_type=gpu_type,
            )

    elif args.command == "status":
        if native:
            show_native_status()
        else:
            show_status()

    elif args.command == "update":
        if native:
            update_native()
        else:
            update_images()

    elif args.command == "logs":
        if native:
            show_native_logs(follow=args.follow)
        else:
            show_logs(follow=args.follow)

    elif args.command == "prompt":
        send_prompt(
            prompt_text=args.text,
            model=args.model,
            temperature=args.temperature,
            top_p=args.top_p,
            image_path=args.image,
            video_path=args.video,
            show_stats=args.stats,
        )


if __name__ == "__main__":
    main()