setup/check_env.py

"""Environment verification script for the Multilingual Code-Switching Chatbot.

Checks Python version, CUDA availability, GPU info, available disk space,
and prints a clean summary table. Run this before installing dependencies.
"""

from __future__ import annotations

import os
import platform
import shutil
import subprocess
import sys
from typing import Any


# ---------------------------------------------------------------------------
# Individual checks
# ---------------------------------------------------------------------------

def check_python_version() -> dict[str, Any]:
    """Check that Python version is 3.9 or newer."""
    major, minor = sys.version_info.major, sys.version_info.minor
    ok = (major, minor) >= (3, 9)
    return {
        "name": "Python version",
        "value": f"{major}.{minor}.{sys.version_info.micro}",
        "ok": ok,
        "detail": "needs >= 3.9" if not ok else "",
    }


def check_platform() -> dict[str, Any]:
    """Report basic OS/platform info (always considered OK)."""
    return {
        "name": "Platform",
        "value": f"{platform.system()} {platform.release()} ({platform.machine()})",
        "ok": True,
        "detail": "",
    }


def check_pip() -> dict[str, Any]:
    """Check that pip is available."""
    try:
        out = subprocess.run(
            [sys.executable, "-m", "pip", "--version"],
            capture_output=True, text=True, check=True,
        )
        return {
            "name": "pip",
            "value": out.stdout.strip().split(" (")[0],
            "ok": True,
            "detail": "",
        }
    except Exception as exc:  # noqa: BLE001
        return {
            "name": "pip",
            "value": "NOT FOUND",
            "ok": False,
            "detail": str(exc),
        }


def check_torch_and_cuda() -> list[dict[str, Any]]:
    """Check whether torch is installed and whether CUDA is visible.

    Returns multiple rows: torch presence, CUDA availability, GPU details.
    """
    rows: list[dict[str, Any]] = []
    try:
        import torch  # type: ignore

        rows.append({
            "name": "PyTorch",
            "value": torch.__version__,
            "ok": True,
            "detail": "",
        })

        cuda_ok = torch.cuda.is_available()
        rows.append({
            "name": "CUDA available",
            "value": "yes" if cuda_ok else "no (CPU only)",
            "ok": True,  # CPU-only is acceptable, just informational
            "detail": "",
        })

        if cuda_ok:
            rows.append({
                "name": "CUDA version (torch)",
                "value": torch.version.cuda or "unknown",
                "ok": True,
                "detail": "",
            })
            n = torch.cuda.device_count()
            for i in range(n):
                props = torch.cuda.get_device_properties(i)
                vram_gb = props.total_memory / (1024 ** 3)
                rows.append({
                    "name": f"GPU {i}",
                    "value": f"{props.name} | {vram_gb:.1f} GB VRAM | CC {props.major}.{props.minor}",
                    "ok": True,
                    "detail": "",
                })
    except ImportError:
        rows.append({
            "name": "PyTorch",
            "value": "NOT INSTALLED",
            "ok": False,
            "detail": "will be installed in setup/install.py",
        })
        # Fall back to nvidia-smi to still report the GPU if present
        rows.extend(_nvidia_smi_fallback())
    except Exception as exc:  # noqa: BLE001
        rows.append({
            "name": "PyTorch",
            "value": "ERROR",
            "ok": False,
            "detail": str(exc),
        })
    return rows


def _nvidia_smi_fallback() -> list[dict[str, Any]]:
    """If torch is unavailable, try nvidia-smi to report GPU presence."""
    rows: list[dict[str, Any]] = []
    if shutil.which("nvidia-smi") is None:
        rows.append({
            "name": "nvidia-smi",
            "value": "not found (no NVIDIA GPU detected, or driver not installed)",
            "ok": True,
            "detail": "",
        })
        return rows
    try:
        out = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total,driver_version",
             "--format=csv,noheader"],
            capture_output=True, text=True, check=True,
        )
        for idx, line in enumerate(out.stdout.strip().splitlines()):
            parts = [p.strip() for p in line.split(",")]
            if len(parts) >= 2:
                name = parts[0]
                mem = parts[1]
                drv = parts[2] if len(parts) > 2 else "?"
                rows.append({
                    "name": f"GPU {idx} (nvidia-smi)",
                    "value": f"{name} | {mem} | driver {drv}",
                    "ok": True,
                    "detail": "",
                })
    except Exception as exc:  # noqa: BLE001
        rows.append({
            "name": "nvidia-smi",
            "value": "ERROR",
            "ok": False,
            "detail": str(exc),
        })
    return rows


def check_disk_space(path: str = ".") -> dict[str, Any]:
    """Report free / total disk space at the given path."""
    try:
        usage = shutil.disk_usage(os.path.abspath(path))
        free_gb = usage.free / (1024 ** 3)
        total_gb = usage.total / (1024 ** 3)
        # 10 GB recommended minimum for models + datasets
        ok = free_gb >= 10
        return {
            "name": "Disk space (cwd)",
            "value": f"{free_gb:.1f} GB free / {total_gb:.1f} GB total",
            "ok": ok,
            "detail": "" if ok else "recommend >= 10 GB free",
        }
    except Exception as exc:  # noqa: BLE001
        return {
            "name": "Disk space (cwd)",
            "value": "ERROR",
            "ok": False,
            "detail": str(exc),
        }


def check_ram() -> dict[str, Any]:
    """Best-effort RAM check via /proc/meminfo (Linux) or psutil if available."""
    # Try /proc/meminfo first
    meminfo = "/proc/meminfo"
    if os.path.exists(meminfo):
        try:
            with open(meminfo) as f:
                lines = f.readlines()
            total_kb = 0
            avail_kb = 0
            for ln in lines:
                if ln.startswith("MemTotal:"):
                    total_kb = int(ln.split()[1])
                elif ln.startswith("MemAvailable:"):
                    avail_kb = int(ln.split()[1])
            total_gb = total_kb / (1024 ** 2)
            avail_gb = avail_kb / (1024 ** 2)
            return {
                "name": "RAM",
                "value": f"{avail_gb:.1f} GB available / {total_gb:.1f} GB total",
                "ok": total_gb >= 4,
                "detail": "" if total_gb >= 4 else "recommend >= 4 GB",
            }
        except Exception:  # noqa: BLE001
            pass
    return {
        "name": "RAM",
        "value": "unknown",
        "ok": True,
        "detail": "could not read /proc/meminfo",
    }


# ---------------------------------------------------------------------------
# Pretty-printer
# ---------------------------------------------------------------------------

def print_table(rows: list[dict[str, Any]]) -> None:
    """Print a clean fixed-width summary table."""
    name_w = max(len(r["name"]) for r in rows) + 2
    val_w = max(len(str(r["value"])) for r in rows) + 2
    name_w = max(name_w, 22)
    val_w = max(val_w, 30)

    sep = "+" + "-" * (name_w + 2) + "+" + "-" * (val_w + 2) + "+--------+"
    print(sep)
    print(f"| {'CHECK'.ljust(name_w)} | {'VALUE'.ljust(val_w)} | STATUS |")
    print(sep)
    for r in rows:
        status = "  OK  " if r["ok"] else " FAIL "
        print(f"| {r['name'].ljust(name_w)} | {str(r['value']).ljust(val_w)} | {status} |")
        if r["detail"]:
            print(f"| {''.ljust(name_w)} | -> {r['detail'][:val_w-3].ljust(val_w-3)} |        |")
    print(sep)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> int:
    """Run all environment checks and return non-zero on hard failure."""
    print("=" * 70)
    print("Multilingual Chatbot — Environment Check")
    print("=" * 70)

    rows: list[dict[str, Any]] = []
    rows.append(check_python_version())
    rows.append(check_platform())
    rows.append(check_pip())
    rows.extend(check_torch_and_cuda())
    rows.append(check_ram())
    rows.append(check_disk_space("."))

    print_table(rows)

    failures = [r for r in rows if not r["ok"]]
    if failures:
        print("\nHard failures:")
        for r in failures:
            print(f"  - {r['name']}: {r['value']}  ({r['detail']})")
        # Only Python version + pip are truly blocking. Torch missing is OK.
        blocking = [r for r in failures if r["name"] in ("Python version", "pip")]
        if blocking:
            print("\nBlocking issues found. Please fix before continuing.")
            return 1
        print("\nNon-blocking issues only — safe to continue to install.py.")
    else:
        print("\nAll checks passed.")

    print("\nSummary:")
    print(f"  - Working dir : {os.getcwd()}")
    print(f"  - Python exe  : {sys.executable}")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\nAborted by user.")
        sys.exit(130)