OBLITERATUS/obliteratus/local_ui.py

"""Local UI launcher for OBLITERATUS — beautiful terminal experience + Gradio web UI.

Usage:
    obliteratus ui                          # launch with defaults
    obliteratus ui --port 8080 --share      # public link on custom port
    obliteratus ui --no-browser             # don't auto-open browser
    obliteratus ui --auth user:pass         # basic auth
"""

from __future__ import annotations

import os
import pathlib
import platform
import shutil
import sys
import tempfile
import time

from rich.console import Console
from rich.panel import Panel
from rich.table import Table

from obliteratus.credential_sources import secret_available

console = Console()

# ── ASCII banner ────────────────────────────────────────────────────────────

_BANNER = r"""
[bold green]
  ░▒█▀▀▀█ ░▒█▀▀▄ ░▒█░░░ ▀█▀ ▀▀█▀▀ ░▒█▀▀▀ ░▒█▀▀█ ▒█▀▀█ ▀▀█▀▀ ░▒█░░▒█ ░▒█▀▀▀█
  ░▒█░░▒█ ░▒█▀▀▄ ░▒█░░░ ░█░ ░░█░░ ░▒█▀▀▀ ░▒█▄▄▀ ▒█▄▄█ ░░█░░ ░▒█░░▒█ ░░▀▀▀▄▄
  ░▒█▄▄▄█ ░▒█▄▄▀ ░▒█▄▄█ ▄█▄ ░░▀░░ ░▒█▄▄▄ ░▒█░▒█ ▒█░▒█ ░░▀░░ ░░▒█▄▄█ ░▒█▄▄▄█
[/bold green]
[dim]  ════════════════════════════════════════════════════════════════════[/dim]
[bold white]   MASTER ABLATION SUITE — LOCAL EDITION[/bold white] [dim]//[/dim] [bold green]Free the mind.[/bold green]
[dim]  ════════════════════════════════════════════════════════════════════[/dim]
"""

_ASCII_BANNER = """
[bold green]OBLITERATUS[/bold green]
[dim]====================================================================[/dim]
[bold white]MASTER ABLATION SUITE - LOCAL EDITION[/bold white] [dim]//[/dim] [bold green]Free the mind.[/bold green]
[dim]====================================================================[/dim]
"""


def _banner_for_console() -> str:
    """Return a startup banner supported by the active stdout encoding."""
    encoding = getattr(sys.stdout, "encoding", None) or "utf-8"
    try:
        _BANNER.encode(encoding)
    except (LookupError, UnicodeEncodeError):
        return _ASCII_BANNER
    return _BANNER


def _detect_gpu() -> list[dict]:
    """Detect available GPUs and return info dicts."""
    gpus = []
    try:
        import torch

        if torch.cuda.is_available():
            for i in range(torch.cuda.device_count()):
                props = torch.cuda.get_device_properties(i)
                gpus.append(
                    {
                        "index": i,
                        "name": props.name,
                        "vram_gb": round(props.total_memory / 1024**3, 1),
                        "compute": f"{props.major}.{props.minor}",
                    }
                )
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            gpus.append(
                {
                    "index": 0,
                    "name": "Apple Silicon (MPS)",
                    "vram_gb": 0,  # shared memory
                    "compute": "mps",
                }
            )
    except (AttributeError, ImportError, OSError, RuntimeError, ValueError):
        pass
    return gpus


def _get_disk_free_gb(path: str = "/") -> float:
    """Return free disk space in GB."""
    try:
        usage = shutil.disk_usage(path)
        return round(usage.free / 1024**3, 1)
    except Exception:
        return 0.0


def _get_ram_gb() -> float:
    """Return total system RAM in GB."""
    try:
        import psutil

        return round(psutil.virtual_memory().total / 1024**3, 1)
    except (AttributeError, ImportError, OSError, RuntimeError, ValueError):
        pass

    # Fallback: read from /proc/meminfo on Linux.
    try:
        with open("/proc/meminfo") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    kb = int(line.split()[1])
                    return round(kb / 1024**2, 1)
    except (OSError, ValueError):
        pass

    # macOS and other POSIX systems generally expose physical RAM via sysconf.
    try:
        pages = os.sysconf("SC_PHYS_PAGES")
        page_size = os.sysconf("SC_PAGE_SIZE")
        if pages > 0 and page_size > 0:
            return round(pages * page_size / 1024**3, 1)
    except (AttributeError, OSError, ValueError):
        pass
    return 0.0


def _print_system_info(gpus: list[dict]) -> None:
    """Print a beautiful system info panel."""
    table = Table(show_header=False, border_style="green", expand=True, padding=(0, 2))
    table.add_column("Key", style="dim", width=18)
    table.add_column("Value", style="cyan")

    table.add_row("Platform", f"{platform.system()} {platform.release()}")
    table.add_row("Python", f"{platform.python_version()}")

    ram = _get_ram_gb()
    if ram > 0:
        table.add_row("System RAM", f"{ram} GB")

    temp_dir = tempfile.gettempdir()
    disk = _get_disk_free_gb(temp_dir)
    if disk > 0:
        table.add_row(f"Disk Free ({temp_dir})", f"{disk} GB")

    # PyTorch version
    try:
        import torch

        table.add_row("PyTorch", torch.__version__)
    except ImportError:
        table.add_row("PyTorch", "[red]NOT INSTALLED[/red]")

    # Transformers version
    try:
        import transformers

        table.add_row("Transformers", transformers.__version__)
    except ImportError:
        table.add_row("Transformers", "[red]NOT INSTALLED[/red]")

    # Gradio version
    try:
        import gradio

        table.add_row("Gradio", gradio.__version__)
    except ImportError:
        table.add_row("Gradio", "[red]NOT INSTALLED[/red]")

    # GPU info
    if gpus:
        for gpu in gpus:
            vram = f"{gpu['vram_gb']} GB" if gpu["vram_gb"] > 0 else "shared"
            table.add_row(
                f"GPU {gpu['index']}",
                f"[bold green]{gpu['name']}[/bold green]  ({vram}, compute {gpu['compute']})",
            )
    else:
        table.add_row("GPU", "[yellow]None detected — CPU mode[/yellow]")

    # HF Token
    if secret_available("HF_TOKEN"):
        table.add_row("HF Token", "[green]set[/green]")
    else:
        table.add_row("HF Token", "[dim]not set (gated models won't work)[/dim]")

    console.print(Panel(table, title="[bold green]System Info[/bold green]", border_style="green"))


def _compute_tier(gpus: list[dict]) -> str:
    """Determine the compute tier based on available hardware."""
    if not gpus:
        return "cpu"
    max_vram = max(g["vram_gb"] for g in gpus)
    if max_vram >= 80:
        return "frontier"
    elif max_vram >= 24:
        return "large"
    elif max_vram >= 8:
        return "medium"
    elif max_vram > 0:
        return "small"
    else:
        return "mps"  # Apple Silicon


def _print_recommendations(tier: str) -> None:
    """Print model recommendations based on detected hardware."""
    recs = {
        "cpu": [
            ("distilbert/distilgpt2", "Tiny — 82M params, runs on anything"),
            ("TinyLlama/TinyLlama-1.1B-Chat-v1.0", "Small — 1.1B params, workable on CPU"),
        ],
        "small": [
            ("TinyLlama/TinyLlama-1.1B-Chat-v1.0", "1.1B — fits in 4GB"),
            ("Qwen/Qwen2.5-0.5B-Instruct", "0.5B — very fast on small GPUs"),
            ("Qwen/Qwen2.5-3B-Instruct", "3B — good quality, ~6GB VRAM"),
        ],
        "mps": [
            ("TinyLlama/TinyLlama-1.1B-Chat-v1.0", "1.1B — fast on Apple Silicon"),
            ("Qwen/Qwen2.5-3B-Instruct", "3B — good with shared memory"),
            ("meta-llama/Llama-3.2-3B-Instruct", "3B — great quality on MPS"),
        ],
        "medium": [
            ("Qwen/Qwen2.5-7B-Instruct", "7B — excellent quality at 16GB"),
            ("meta-llama/Llama-3.1-8B-Instruct", "8B — the standard benchmark"),
            ("mistralai/Mistral-7B-Instruct-v0.3", "7B — fast and capable"),
        ],
        "large": [
            ("openai/gpt-oss-20b", "20B MoE — flagship, try surgical/nuclear"),
            ("meta-llama/Llama-3.1-70B-Instruct", "70B — with 4bit quant"),
            ("Qwen/Qwen3-30B-A3B", "30B MoE — only 3B active"),
        ],
        "frontier": [
            ("deepseek-ai/DeepSeek-V3", "671B MoE — the frontier"),
            ("meta-llama/Llama-3.1-70B-Instruct", "70B — full precision"),
            ("Qwen/Qwen3.5-397B-A17B", "397B MoE — massive"),
        ],
    }

    tier_recs = recs.get(tier, recs["cpu"])

    table = Table(
        title=f"Recommended Models for Your Hardware ({tier.upper()})",
        border_style="green",
    )
    table.add_column("Model", style="cyan", min_width=40)
    table.add_column("Notes", style="dim")

    for model, notes in tier_recs:
        table.add_row(model, notes)

    console.print(table)


def _print_launch_info(
    host: str,
    port: int,
    share: bool,
    auth: tuple[str, str] | None,
) -> None:
    """Print launch configuration."""
    parts = []
    url = f"http://{'localhost' if host == '0.0.0.0' else host}:{port}"
    parts.append(f"[bold green]Local URL:[/bold green]     {url}")
    if host == "0.0.0.0":
        parts.append(f"[bold green]Network URL:[/bold green]   http://<your-ip>:{port}")
    if share:
        parts.append("[bold green]Share link:[/bold green]    [yellow]generating...[/yellow]")
    if auth:
        parts.append(f"[bold green]Auth:[/bold green]          {auth[0]}:{'*' * len(auth[1])}")
    parts.append("")
    parts.append("[dim]Press Ctrl+C to stop the server[/dim]")

    console.print(
        Panel(
            "\n".join(parts),
            title="[bold green]Launching OBLITERATUS UI[/bold green]",
            border_style="green",
        )
    )


def launch_local_ui(
    host: str = "0.0.0.0",
    port: int = 7860,
    share: bool = False,
    open_browser: bool = True,
    auth: tuple[str, str] | None = None,
    quiet: bool = False,
) -> None:
    """Launch the OBLITERATUS Gradio UI with a beautiful local experience.

    This is the main entry point for ``obliteratus ui``.
    """
    # ── Beautiful startup ──────────────────────────────────────────────
    if not quiet:
        console.print(_banner_for_console())

        gpus = _detect_gpu()
        _print_system_info(gpus)
        console.print()

        tier = _compute_tier(gpus)
        _print_recommendations(tier)
        console.print()

        _print_launch_info(host, port, share, auth)
        console.print()

    # ── Check dependencies ─────────────────────────────────────────────
    try:
        import gradio  # noqa: F401
    except ImportError:
        console.print(
            "[bold red]Error:[/bold red] Gradio is not installed.\n"
            "Install it with: [cyan]pip install -e '.[spaces]'[/cyan]\n"
            "Or: [cyan]pip install gradio>=5.0[/cyan]"
        )
        sys.exit(1)

    # ── Import and launch the app ──────────────────────────────────────
    console.print("[dim]Loading OBLITERATUS UI (this may take a moment on first run)...[/dim]")
    start = time.time()

    # app.py lives at the project root, one level above this package.
    # When installed via pip the root isn't on sys.path, so add it.
    _project_root = str(pathlib.Path(__file__).resolve().parent.parent)
    if _project_root not in sys.path:
        sys.path.insert(0, _project_root)

    from app import launch as app_launch

    elapsed = time.time() - start
    if not quiet:
        console.print(f"[green]UI loaded in {elapsed:.1f}s[/green]")
        console.print()

    app_launch(
        server_name=host,
        server_port=port,
        share=share,
        inbrowser=open_browser,
        auth=auth,
        quiet=quiet,
    )