from __future__ import annotations

import importlib.util
import os
import platform
import shutil
import subprocess
import sys
from pathlib import Path

_LMM_GUARD_INSTALL_ATTEMPTED = False


def device_supports_llm_guard() -> bool:
    if str(os.getenv("ARMORER_DISABLE_LLM_GUARD", "")).strip().lower() in {"1", "true", "yes", "on"}:
        return False
    if str(os.getenv("ARMORER_FORCE_LLM_GUARD", "")).strip().lower() in {"1", "true", "yes", "on"}:
        return True

    if sys.platform == "darwin" and platform.machine().lower() == "arm64":
        return True

    visible_devices = str(os.getenv("NVIDIA_VISIBLE_DEVICES", "")).strip().lower()
    if visible_devices and visible_devices not in {"void", "none"}:
        return True

    if any(Path(path).exists() for path in ("/dev/nvidiactl", "/dev/nvidia0")):
        return True

    nvidia_smi = shutil.which("nvidia-smi")
    if not nvidia_smi:
        return False
    try:
        result = subprocess.run(
            [nvidia_smi, "-L"],
            capture_output=True,
            text=True,
            timeout=5,
        )
    except Exception:
        return False
    return result.returncode == 0 and bool((result.stdout or "").strip())


def llm_guard_installed() -> bool:
    return importlib.util.find_spec("llm_guard") is not None


def ensure_llm_guard_runtime_dependency() -> bool:
    global _LMM_GUARD_INSTALL_ATTEMPTED

    if str(os.getenv("ARMORER_DISABLE_LLM_GUARD", "")).strip().lower() in {"1", "true", "yes", "on"}:
        return False
    if llm_guard_installed():
        return True
    if not device_supports_llm_guard():
        return False
    if str(os.getenv("ARMORER_SKIP_LLM_GUARD_INSTALL", "")).strip().lower() in {"1", "true", "yes", "on"}:
        return False
    if _LMM_GUARD_INSTALL_ATTEMPTED:
        return False

    _LMM_GUARD_INSTALL_ATTEMPTED = True
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "install", "llm-guard>=0.3.15"],
            capture_output=True,
            text=True,
            timeout=600,
        )
    except Exception:
        return False
    return result.returncode == 0 and llm_guard_installed()
