#!/usr/bin/env python3
"""
GPU Utilities for OCR Provenance MCP System

Provides universal device detection, VRAM monitoring, and GPU validation.
Supports CUDA (any NVIDIA GPU), MPS (Apple Silicon), and CPU fallback.

Device Priority: CUDA > MPS > CPU
- CUDA: Any NVIDIA GPU with CUDA support (no minimum compute capability)
- MPS: Apple Silicon (M1/M2/M3/M4+) via Metal Performance Shaders
- CPU: Always available as fallback

All functions log which device was chosen and why. Explicit device requests
('cuda', 'mps') fail fast with clear errors if the device is unavailable.
CPU never fails.
"""

import argparse
import json
import logging
import os
import sys
import warnings
from pathlib import Path
from typing import Optional, TypedDict

# Suppress pynvml FutureWarning from torch.cuda (emitted to stderr, breaks JSON-RPC)
warnings.filterwarnings("ignore", category=FutureWarning, message=".*pynvml.*")

# Configure logging with detailed format for debugging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=sys.stderr,
)
logger = logging.getLogger(__name__)


# =============================================================================
# Error Classes (CS-ERR-001 compliant)
# =============================================================================


class GPUError(Exception):
    """Base class for GPU-related errors."""

    pass


class GPUNotAvailableError(GPUError):
    """Raised when a specifically requested device is not available."""

    def __init__(
        self,
        message: str = "Requested device not available.",
    ):
        self.message = message
        super().__init__(self.message)


class GPUOutOfMemoryError(GPUError):
    """Raised when GPU runs out of memory."""

    def __init__(
        self, message: str, vram_required: Optional[float] = None, vram_available: Optional[float] = None
    ):
        self.message = message
        self.vram_required = vram_required
        self.vram_available = vram_available
        super().__init__(self.message)


class EmbeddingModelError(Exception):
    """Raised when embedding model fails to load or run."""

    def __init__(self, message: str, model_path: Optional[str] = None, cause: Optional[Exception] = None):
        self.message = message
        self.model_path = model_path
        self.cause = cause
        super().__init__(self.message)


# =============================================================================
# Type Definitions
# =============================================================================


class GPUInfo(TypedDict):
    """Device information structure. Works for CUDA, MPS, and CPU."""

    available: bool
    name: str
    vram_gb: float
    vram_used_gb: float
    vram_free_gb: float
    cuda_version: str
    compute_capability: str
    driver_version: str


class VRAMUsage(TypedDict):
    """VRAM usage statistics. Returns zeros on non-CUDA devices."""

    allocated_gb: float
    reserved_gb: float
    free_gb: float
    total_gb: float


class ModelInfo(TypedDict):
    """Embedding model information."""

    success: bool
    model_path: str
    error: Optional[str]
    model_info: dict


# =============================================================================
# Device Detection Functions
# =============================================================================


def _available_devices_summary() -> str:
    """Build a human-readable summary of all available devices for error messages."""
    parts = []
    try:
        import torch
        if torch.cuda.is_available():
            count = torch.cuda.device_count()
            names = [torch.cuda.get_device_name(i) for i in range(count)]
            parts.append(f"CUDA: {count} device(s) ({', '.join(names)})")
        else:
            parts.append("CUDA: not available")

        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            parts.append("MPS: available (Apple Silicon)")
        else:
            parts.append("MPS: not available")
    except ImportError:
        parts.append("PyTorch: not installed")

    parts.append("CPU: always available")
    return "; ".join(parts)


def detect_best_device() -> str:
    """
    Auto-detect the best available compute device.

    Priority: CUDA > MPS > CPU. NEVER raises an error -- always returns
    a valid device string. Logs the chosen device and reason.

    Returns:
        Device string: 'cuda:0', 'mps', or 'cpu'
    """
    try:
        import torch
    except ImportError:
        logger.warning("PyTorch not installed. Falling back to CPU.")
        return "cpu"

    # CUDA (any NVIDIA GPU)
    if torch.cuda.is_available():
        device_name = torch.cuda.get_device_name(0)
        logger.info("Device auto-detected: cuda:0 (%s)", device_name)
        return "cuda:0"

    # MPS (Apple Silicon)
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        logger.info("Device auto-detected: mps (Apple Silicon)")
        return "mps"

    # CPU fallback
    logger.info("Device auto-detected: cpu (no GPU available)")
    return "cpu"


def resolve_device(requested: str) -> str:
    """
    Resolve a device request to a validated device string.

    Args:
        requested: One of 'auto', 'cuda', 'cuda:N', 'mps', or 'cpu'

    Returns:
        Validated device string

    Raises:
        GPUNotAvailableError: If an explicitly requested device is not available.
            Error includes: what was requested, what was found, available devices.
    """
    if requested == "auto":
        return detect_best_device()

    if requested == "cpu":
        logger.info("Device explicitly set: cpu")
        return "cpu"

    if requested == "mps":
        try:
            import torch
        except ImportError:
            raise GPUNotAvailableError(
                f"Device 'mps' requested but PyTorch is not installed. "
                f"Available: {_available_devices_summary()}"
            )
        if not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()):
            raise GPUNotAvailableError(
                f"Device 'mps' requested but MPS is not available. "
                f"MPS requires Apple Silicon (M1+) and macOS 12.3+. "
                f"Available: {_available_devices_summary()}"
            )
        logger.info("Device explicitly set: mps")
        return "mps"

    if requested == "cuda" or requested.startswith("cuda:"):
        try:
            import torch
        except ImportError:
            raise GPUNotAvailableError(
                f"Device '{requested}' requested but PyTorch is not installed. "
                f"Available: {_available_devices_summary()}"
            )
        if not torch.cuda.is_available():
            raise GPUNotAvailableError(
                f"Device '{requested}' requested but CUDA is not available. "
                f"No NVIDIA GPU detected or CUDA drivers not installed. "
                f"Available: {_available_devices_summary()}"
            )
        # Validate specific device index
        if requested.startswith("cuda:"):
            try:
                idx = int(requested.split(":")[1])
            except (ValueError, IndexError):
                raise GPUNotAvailableError(
                    f"Invalid CUDA device format: '{requested}'. Expected 'cuda:N' where N is an integer."
                )
            if idx >= torch.cuda.device_count():
                raise GPUNotAvailableError(
                    f"Device '{requested}' requested but only {torch.cuda.device_count()} CUDA device(s) available. "
                    f"Valid range: cuda:0 to cuda:{torch.cuda.device_count() - 1}. "
                    f"Available: {_available_devices_summary()}"
                )
        device = requested if ":" in requested else "cuda:0"
        device_name = torch.cuda.get_device_name(int(device.split(":")[1]))
        logger.info("Device explicitly set: %s (%s)", device, device_name)
        return device

    raise GPUNotAvailableError(
        f"Unknown device '{requested}'. Valid options: 'auto', 'cuda', 'cuda:N', 'mps', 'cpu'. "
        f"Available: {_available_devices_summary()}"
    )


# =============================================================================
# Environment Validation
# =============================================================================


def _get_driver_version() -> str:
    """
    Get NVIDIA driver version via pynvml, falling back to torch internal API.

    Suppresses the pynvml FutureWarning about nvidia-ml-py package rename.
    Returns 'N/A' on non-CUDA systems.
    """
    import warnings

    try:
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=FutureWarning, module="pynvml")
            import pynvml
        pynvml.nvmlInit()
        version = pynvml.nvmlSystemGetDriverVersion()
        pynvml.nvmlShutdown()
        return version
    except Exception:
        pass

    try:
        import torch
        return str(torch._C._cuda_getDriverVersion())
    except Exception:
        return "N/A"


def _get_cuda_toolkit_version() -> str:
    """
    Get the installed CUDA toolkit version from nvcc or version.json.
    Returns version string like '13.2' or 'N/A'.
    """
    import re
    import subprocess

    # Try nvcc -- prefer /usr/local/cuda symlink (alternatives-managed), then $CUDA_HOME, then PATH
    nvcc_paths = ["/usr/local/cuda/bin/nvcc"]
    cuda_home = os.environ.get("CUDA_HOME", "")
    if cuda_home:
        nvcc_paths.append(cuda_home + "/bin/nvcc")
    nvcc_paths.append("nvcc")

    for nvcc_path in nvcc_paths:
        try:
            result = subprocess.run(
                [nvcc_path, "--version"],
                capture_output=True, text=True, timeout=5,
            )
            for line in result.stdout.split("\n"):
                if "release" in line.lower():
                    match = re.search(r"release\s+(\d+\.\d+)", line)
                    if match:
                        return match.group(1)
        except Exception:
            continue

    # Try version.json
    version_paths = ["/usr/local/cuda/version.json"]
    if cuda_home:
        version_paths.append(cuda_home + "/version.json")

    for vpath in version_paths:
        try:
            with open(vpath) as f:
                data = json.load(f)
            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:
            continue

    return "N/A"


def validate_gpu_environment() -> dict:
    """
    Validate and log the compute environment for the best available device.

    For CUDA: Reports GPU name, VRAM, compute capability, driver version,
    toolkit version, PyTorch version, BF16 support, torch.compile availability.
    For MPS: Reports MPS availability and PyTorch version.
    For CPU: Reports CPU-only mode and PyTorch version.

    NEVER errors on low compute capability. ANY NVIDIA GPU works.
    NEVER requires BF16 or specific SM architecture.

    Returns:
        dict with environment information keyed by device type
    """
    device = detect_best_device()

    try:
        import torch
        pytorch_version = torch.__version__
    except ImportError:
        logger.warning("PyTorch not installed. Reporting CPU-only environment.")
        env_info = {
            "device": "cpu",
            "device_type": "cpu",
            "pytorch_installed": False,
            "pytorch_version": "not installed",
        }
        logger.info("=== Compute Environment ===")
        logger.info("  Device:          CPU (PyTorch not installed)")
        logger.info("===========================")
        return env_info

    if device.startswith("cuda"):
        dev_idx = int(device.split(":")[1]) if ":" in device else 0
        props = torch.cuda.get_device_properties(dev_idx)
        compute_cap = f"{props.major}.{props.minor}"
        total_vram_gb = round(props.total_memory / (1024**3), 2)
        torch_compile_available = hasattr(torch, "compile")
        cuda_alloc_conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "not set")
        driver_version = _get_driver_version()
        toolkit_version = _get_cuda_toolkit_version()

        # Check BF16 support (informational only, not required)
        bf16_supported = False
        try:
            bf16_supported = torch.cuda.is_bf16_supported()
        except Exception:
            pass

        env_info = {
            "device": device,
            "device_type": "cuda",
            "cuda_toolkit_version": toolkit_version,
            "cuda_runtime_version": torch.version.cuda or "unknown",
            "driver_version": driver_version,
            "gpu_name": props.name,
            "compute_capability": compute_cap,
            "sm_count": props.multi_processor_count,
            "vram_total_gb": total_vram_gb,
            "bf16_supported": bf16_supported,
            "torch_compile_available": torch_compile_available,
            "pytorch_version": pytorch_version,
            "pytorch_installed": True,
            "cuda_alloc_conf": cuda_alloc_conf,
        }

        logger.info("=== Compute Environment (CUDA) ===")
        logger.info("  PyTorch version:      %s", env_info["pytorch_version"])
        logger.info("  CUDA toolkit (nvcc):  %s", env_info["cuda_toolkit_version"])
        logger.info("  CUDA runtime (torch): %s", env_info["cuda_runtime_version"])
        logger.info("  Driver version:       %s", env_info["driver_version"])
        logger.info("  GPU:                  %s (%d SMs)", env_info["gpu_name"], env_info["sm_count"])
        logger.info("  Compute capability:   %s", env_info["compute_capability"])
        logger.info("  VRAM total:           %.2f GB", env_info["vram_total_gb"])
        logger.info("  BF16 support:         %s", "yes" if env_info["bf16_supported"] else "no")
        logger.info("  torch.compile:        %s", env_info["torch_compile_available"])
        logger.info("  CUDA alloc config:    %s", env_info["cuda_alloc_conf"])
        if toolkit_version not in ("N/A",) and toolkit_version != env_info["cuda_runtime_version"]:
            logger.info(
                "  Note: toolkit=%s, PyTorch runtime=%s (forward-compatible).",
                toolkit_version, env_info["cuda_runtime_version"],
            )
        logger.info("==================================")

    elif device == "mps":
        env_info = {
            "device": "mps",
            "device_type": "mps",
            "pytorch_version": pytorch_version,
            "pytorch_installed": True,
            "gpu_name": "Apple Silicon (MPS)",
            "mps_available": True,
        }

        logger.info("=== Compute Environment (MPS) ===")
        logger.info("  PyTorch version:  %s", env_info["pytorch_version"])
        logger.info("  Device:           MPS (Apple Silicon)")
        logger.info("=================================")

    else:
        env_info = {
            "device": "cpu",
            "device_type": "cpu",
            "pytorch_version": pytorch_version,
            "pytorch_installed": True,
            "gpu_name": "CPU",
        }

        logger.info("=== Compute Environment (CPU) ===")
        logger.info("  PyTorch version:  %s", env_info["pytorch_version"])
        logger.info("  Device:           CPU (no GPU available)")
        logger.info("=================================")

    return env_info


# =============================================================================
# Startup and Configuration
# =============================================================================


def configure_cuda_memory_pool() -> None:
    """
    Configure CUDA memory allocator for optimal performance.

    Sets PYTORCH_CUDA_ALLOC_CONF with expandable_segments and garbage_collection_threshold.
    No-op when CUDA is not available.
    """
    try:
        import torch
    except ImportError:
        return

    if not torch.cuda.is_available():
        logger.debug("CUDA not available, skipping memory pool configuration.")
        return

    key = "PYTORCH_CUDA_ALLOC_CONF"
    current = os.environ.get(key)
    if current:
        logger.info("CUDA memory pool already configured: %s", current)
        return

    value = "expandable_segments:True,garbage_collection_threshold:0.8,max_split_size_mb:512"
    os.environ[key] = value
    logger.info("CUDA memory pool configured: %s=%s", key, value)


def startup_gpu(worker_name: str) -> dict:
    """
    Run the standard device startup sequence for a worker.

    - Resolves TORCH_DEVICE env var from 'auto' to actual device (cpu/cuda:0/mps)
    - Configures CUDA memory pool (only when CUDA is available)
    - Validates the compute environment
    - Updates os.environ["TORCH_DEVICE"] so downstream libraries (Marker-pdf) get a valid device string

    NEVER raises GPUNotAvailableError -- CPU is always a valid device.

    Args:
        worker_name: Name of the worker (for log messages)

    Returns:
        Environment info dict from validate_gpu_environment()
    """
    # Resolve TORCH_DEVICE env var BEFORE any model loading.
    # Marker-pdf reads TORCH_DEVICE directly from env and passes it to PyTorch.
    # PyTorch does NOT accept "auto" — it must be "cpu", "cuda:0", "mps", etc.
    raw_device = os.environ.get("TORCH_DEVICE", "auto")
    resolved_device = resolve_device(raw_device)
    if raw_device != resolved_device:
        logger.info(
            "Resolved TORCH_DEVICE: '%s' → '%s' (updating env var for downstream libraries)",
            raw_device, resolved_device,
        )
    os.environ["TORCH_DEVICE"] = resolved_device

    # Also resolve EMBEDDING_DEVICE if set to "auto"
    raw_embed_device = os.environ.get("EMBEDDING_DEVICE", "")
    if raw_embed_device == "auto":
        resolved_embed = resolve_device("auto")
        os.environ["EMBEDDING_DEVICE"] = resolved_embed
        logger.info("Resolved EMBEDDING_DEVICE: 'auto' → '%s'", resolved_embed)

    configure_cuda_memory_pool()
    env_info = validate_gpu_environment()
    device = env_info.get("device", "cpu")
    device_name = env_info.get("gpu_name", device)
    logger.info("%s compute environment validated: %s (device=%s)", worker_name, device_name, device)
    return env_info


def auto_tune_batch_size(device: str, vram_gb: Optional[float] = None) -> int:
    """
    Auto-tune batch size based on device type and available VRAM.

    Args:
        device: Device string ('cuda:0', 'mps', 'cpu', etc.)
        vram_gb: Override VRAM in GB. If None, auto-detected for CUDA devices.

    Returns:
        Recommended batch size: 512, 256, 128, 64, or 32
    """
    if device == "cpu" or device == "":
        logger.info("Batch size auto-tuned: 32 (CPU device)")
        return 32

    if vram_gb is None and device.startswith("cuda"):
        try:
            import torch
            if torch.cuda.is_available():
                dev_idx = int(device.split(":")[1]) if ":" in device else 0
                props = torch.cuda.get_device_properties(dev_idx)
                vram_gb = props.total_memory / (1024**3)
        except Exception as e:
            logger.warning("Could not detect VRAM for %s: %s. Using default batch size.", device, e)

    if vram_gb is None:
        # MPS or CUDA where detection failed
        batch_size = 64
        logger.info("Batch size auto-tuned: %d (unknown VRAM for device=%s)", batch_size, device)
        return batch_size

    if vram_gb >= 24:
        batch_size = 512
    elif vram_gb >= 16:
        batch_size = 256
    elif vram_gb >= 8:
        batch_size = 128
    elif vram_gb >= 4:
        batch_size = 64
    else:
        batch_size = 32

    logger.info("Batch size auto-tuned: %d (%.1f GB VRAM on %s)", batch_size, vram_gb, device)
    return batch_size


# =============================================================================
# GPU Verification Functions
# =============================================================================


def verify_gpu() -> GPUInfo:
    """
    Verify device availability and return detailed information.

    Works for CUDA, MPS, and CPU:
    - CUDA: Full GPU info (name, VRAM, compute capability, CUDA version, driver)
    - MPS: Basic info (Apple Silicon, no VRAM stats)
    - CPU: Minimal info (CPU-only mode)

    NEVER raises GPUNotAvailableError. Always returns valid info for some device.

    Returns:
        GPUInfo dict with device specifications
    """
    logger.info("Starting device verification...")

    try:
        import torch
    except ImportError:
        logger.warning("PyTorch not installed. Reporting CPU-only device.")
        return GPUInfo(
            available=False,
            name="CPU (PyTorch not installed)",
            vram_gb=0.0,
            vram_used_gb=0.0,
            vram_free_gb=0.0,
            cuda_version="N/A",
            compute_capability="N/A",
            driver_version="N/A",
        )

    # CUDA
    if torch.cuda.is_available():
        device_idx = torch.cuda.current_device()
        props = torch.cuda.get_device_properties(device_idx)
        total_memory = props.total_memory / (1024**3)
        allocated = torch.cuda.memory_allocated(device_idx) / (1024**3)
        free = total_memory - allocated
        compute_cap = f"{props.major}.{props.minor}"
        driver_version = _get_driver_version()

        gpu_info = GPUInfo(
            available=True,
            name=props.name,
            vram_gb=round(total_memory, 2),
            vram_used_gb=round(allocated, 2),
            vram_free_gb=round(free, 2),
            cuda_version=torch.version.cuda or "unknown",
            compute_capability=compute_cap,
            driver_version=driver_version,
        )

        logger.info(
            "Device verified: CUDA - %s, VRAM: %.2f GB, CUDA: %s, Compute: %s, Driver: %s",
            gpu_info["name"],
            gpu_info["vram_gb"],
            gpu_info["cuda_version"],
            gpu_info["compute_capability"],
            gpu_info["driver_version"],
        )
        return gpu_info

    # MPS
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        gpu_info = GPUInfo(
            available=True,
            name="Apple Silicon (MPS)",
            vram_gb=0.0,  # MPS shares system memory, no dedicated VRAM reporting
            vram_used_gb=0.0,
            vram_free_gb=0.0,
            cuda_version="N/A",
            compute_capability="N/A",
            driver_version="N/A",
        )

        logger.info("Device verified: MPS (Apple Silicon)")
        return gpu_info

    # CPU
    gpu_info = GPUInfo(
        available=False,
        name="CPU",
        vram_gb=0.0,
        vram_used_gb=0.0,
        vram_free_gb=0.0,
        cuda_version="N/A",
        compute_capability="N/A",
        driver_version="N/A",
    )

    logger.info("Device verified: CPU only (no GPU available)")
    return gpu_info


def get_vram_usage() -> VRAMUsage:
    """
    Get current VRAM usage statistics.

    Returns zeros gracefully on non-CUDA devices (MPS, CPU) -- never errors.

    Returns:
        VRAMUsage dict with allocated, reserved, and free memory in GB.
        All zeros on non-CUDA devices.
    """
    logger.debug("Querying VRAM usage...")

    zero_usage = VRAMUsage(
        allocated_gb=0.0,
        reserved_gb=0.0,
        free_gb=0.0,
        total_gb=0.0,
    )

    try:
        import torch
    except ImportError:
        logger.debug("PyTorch not installed. Returning zero VRAM usage.")
        return zero_usage

    if not torch.cuda.is_available():
        logger.debug("CUDA not available. Returning zero VRAM usage.")
        return zero_usage

    device = torch.cuda.current_device()

    allocated = torch.cuda.memory_allocated(device) / (1024**3)
    reserved = torch.cuda.memory_reserved(device) / (1024**3)
    total = torch.cuda.get_device_properties(device).total_memory / (1024**3)

    usage = VRAMUsage(
        allocated_gb=round(allocated, 3),
        reserved_gb=round(reserved, 3),
        free_gb=round(total - reserved, 3),
        total_gb=round(total, 2),
    )

    logger.debug(
        "VRAM: allocated=%.3f GB, reserved=%.3f GB, free=%.3f GB, total=%.2f GB",
        usage["allocated_gb"],
        usage["reserved_gb"],
        usage["free_gb"],
        usage["total_gb"],
    )

    return usage


def verify_model_loading(model_path: str = "./models/nomic-embed-text-v1.5") -> ModelInfo:
    """
    Verify the embedding model can be loaded on the best available device.

    Uses detect_best_device() to find the optimal device, then loads the model there.
    Works on CUDA, MPS, and CPU.

    Args:
        model_path: Path to the nomic-embed-text-v1.5 model

    Returns:
        ModelInfo dict with loading status and model info

    Raises:
        EmbeddingModelError: If model files are missing or loading fails
    """
    logger.info("Verifying model loading from: %s", model_path)

    model_dir = Path(model_path)
    if not model_dir.exists():
        error_msg = f"Model directory not found: {model_path}"
        logger.error(error_msg)
        raise EmbeddingModelError(error_msg, model_path=model_path)

    # Check for required model files
    required_files = ["config.json", "model.safetensors", "tokenizer.json"]
    missing_files = [f for f in required_files if not (model_dir / f).exists()]
    if missing_files:
        error_msg = f"Missing required model files: {missing_files}"
        logger.error(error_msg)
        raise EmbeddingModelError(error_msg, model_path=model_path)

    try:
        import torch
        from sentence_transformers import SentenceTransformer

        device = detect_best_device()
        logger.info("Loading model to device: %s", device)

        # M-7: Pin revision to prevent supply-chain attacks with trust_remote_code=True
        model = SentenceTransformer(model_path, device=device, trust_remote_code=True, local_files_only=True, revision="0ef289df6812e6e0e532352f5174b161cbfb30db")

        # Get model info
        embedding_dim = model.get_sentence_embedding_dimension()
        max_seq_len = model.max_seq_length

        model_info = ModelInfo(
            success=True,
            model_path=model_path,
            error=None,
            model_info={
                "device": device,
                "embedding_dimension": embedding_dim,
                "max_seq_length": max_seq_len,
            },
        )

        logger.info(
            "Model loaded successfully: dim=%d, max_seq=%d, device=%s",
            embedding_dim,
            max_seq_len,
            device,
        )

        # Cleanup
        del model
        if device.startswith("cuda") and torch.cuda.is_available():
            torch.cuda.empty_cache()
        logger.debug("Model unloaded and cache cleared")

        return model_info

    except EmbeddingModelError:
        raise
    except Exception as e:
        logger.error("Model loading failed: %s", e)
        raise EmbeddingModelError(
            f"Failed to load embedding model: {e}", model_path=model_path, cause=e
        ) from e


def clear_gpu_memory() -> None:
    """
    Clear GPU memory cache. Safe to call on any device. Never raises.

    On CUDA: calls torch.cuda.empty_cache()
    On MPS/CPU: no-op (safe)

    NOTE: Does NOT call gc.collect() or torch.cuda.synchronize() -- those are
    expensive operations that callers should invoke explicitly when needed
    (e.g., after deleting large objects).
    """
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            logger.debug("GPU memory cleared")
    except ImportError:
        logger.warning("PyTorch not installed -- cannot free GPU memory")
    except Exception as e:
        logger.warning("Failed to clear GPU memory: %s", e)


def test_embedding_generation(model_path: str = "./models/nomic-embed-text-v1.5") -> dict:
    """
    Test that embedding generation works end-to-end on the best available device.

    Uses detect_best_device() to find the optimal device, then generates
    test embeddings. Works on CUDA, MPS, and CPU.

    Args:
        model_path: Path to the embedding model

    Returns:
        dict with test results including embedding shape, timing, and device

    Raises:
        EmbeddingModelError: If embedding generation fails
    """
    logger.info("Testing embedding generation...")

    try:
        import time

        import torch
        from sentence_transformers import SentenceTransformer

        device = detect_best_device()
        logger.info("Testing embeddings on device: %s", device)

        # M-7: Pin revision to prevent supply-chain attacks with trust_remote_code=True
        model = SentenceTransformer(model_path, device=device, trust_remote_code=True, local_files_only=True, revision="0ef289df6812e6e0e532352f5174b161cbfb30db")

        # Test with sample text
        test_texts = [
            "This is a test document for OCR provenance verification.",
            "The system must maintain complete data lineage.",
            "Every embedding must trace back to its source file.",
        ]

        # Time the embedding generation
        start_time = time.perf_counter()
        embeddings = model.encode(test_texts, device=device, convert_to_numpy=True)
        end_time = time.perf_counter()

        elapsed_ms = (end_time - start_time) * 1000

        result = {
            "success": True,
            "embedding_shape": list(embeddings.shape),
            "embedding_dimension": embeddings.shape[1],
            "num_texts": len(test_texts),
            "elapsed_ms": round(elapsed_ms, 2),
            "ms_per_text": round(elapsed_ms / len(test_texts), 2),
            "device": device,
        }

        logger.info(
            "Embedding test passed: shape=%s, time=%.2fms, device=%s",
            embeddings.shape,
            elapsed_ms,
            device,
        )

        # Cleanup
        del model
        del embeddings
        if device.startswith("cuda") and torch.cuda.is_available():
            torch.cuda.empty_cache()

        return result

    except EmbeddingModelError:
        raise
    except Exception as e:
        logger.error("Embedding generation test failed: %s", e)
        raise EmbeddingModelError(
            f"Embedding generation test failed: {e}", model_path=model_path, cause=e
        ) from e


# =============================================================================
# CLI Entry Point
# =============================================================================


def main() -> None:
    """CLI entry point for device verification and diagnostics."""
    parser = argparse.ArgumentParser(
        description="GPU/Device Utilities for OCR Provenance MCP System",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python gpu_utils.py --verify                 # Verify device availability
  python gpu_utils.py --verify --json          # Output as JSON
  python gpu_utils.py --vram                   # Show VRAM usage (CUDA only)
  python gpu_utils.py --verify-model           # Verify model loading
  python gpu_utils.py --test-embedding         # Test embedding generation
  python gpu_utils.py --device cuda            # Explicitly request CUDA
  python gpu_utils.py --device auto            # Auto-detect best device
        """,
    )
    parser.add_argument(
        "--verify", action="store_true", help="Verify device availability and capabilities"
    )
    parser.add_argument("--vram", action="store_true", help="Show current VRAM usage (CUDA only)")
    parser.add_argument(
        "--model",
        type=str,
        default="./models/nomic-embed-text-v1.5",
        help="Path to embedding model (default: ./models/nomic-embed-text-v1.5)",
    )
    parser.add_argument(
        "--verify-model", action="store_true", help="Verify model can be loaded on device"
    )
    parser.add_argument(
        "--test-embedding", action="store_true", help="Test embedding generation end-to-end"
    )
    parser.add_argument(
        "--device",
        type=str,
        default="auto",
        help="Device to use: auto, cuda, cuda:N, mps, cpu (default: auto)",
    )
    parser.add_argument("--json", action="store_true", help="Output results as JSON")
    parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose logging")

    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # Default to --verify if no arguments
    if not any([args.verify, args.vram, args.verify_model, args.test_embedding]):
        args.verify = True

    results: dict = {}
    exit_code = 0

    try:
        # Resolve device first if explicitly set
        if args.device != "auto":
            resolved = resolve_device(args.device)
            results["resolved_device"] = resolved

        if args.verify:
            gpu_info = verify_gpu()
            results["gpu"] = dict(gpu_info)

            if not args.json:
                print("=" * 60)
                print("Device Verification Results")
                print("=" * 60)
                device_type = "CUDA" if gpu_info["cuda_version"] != "N/A" else ("MPS" if gpu_info["name"] == "Apple Silicon (MPS)" else "CPU")
                print(f"  Device Type:         {device_type}")
                print(f"  Device Name:         {gpu_info['name']}")
                if gpu_info["vram_gb"] > 0:
                    print(f"  VRAM Total:          {gpu_info['vram_gb']} GB")
                    print(f"  VRAM Free:           {gpu_info['vram_free_gb']} GB")
                if gpu_info["cuda_version"] != "N/A":
                    print(f"  CUDA Version:        {gpu_info['cuda_version']}")
                    print(f"  Compute Capability:  {gpu_info['compute_capability']}")
                    print(f"  Driver Version:      {gpu_info['driver_version']}")
                print("=" * 60)

        if args.vram:
            vram_info = get_vram_usage()
            results["vram"] = dict(vram_info)

            if not args.json:
                if vram_info["total_gb"] > 0:
                    print("\nVRAM Usage:")
                    print(f"  Allocated: {vram_info['allocated_gb']} GB")
                    print(f"  Reserved:  {vram_info['reserved_gb']} GB")
                    print(f"  Free:      {vram_info['free_gb']} GB")
                    print(f"  Total:     {vram_info['total_gb']} GB")
                else:
                    print("\nVRAM Usage: N/A (no CUDA device)")

        if args.verify_model:
            model_result = verify_model_loading(args.model)
            results["model"] = dict(model_result)

            if not args.json:
                print("\nModel Verification:")
                print("  Status:    SUCCESS")
                print(f"  Path:      {model_result['model_path']}")
                print(f"  Device:    {model_result['model_info'].get('device', 'N/A')}")
                print(
                    f"  Dimension: {model_result['model_info'].get('embedding_dimension', 'N/A')}"
                )
                print(f"  Max Seq:   {model_result['model_info'].get('max_seq_length', 'N/A')}")

        if args.test_embedding:
            embed_result = test_embedding_generation(args.model)
            results["embedding_test"] = embed_result

            if not args.json:
                print("\nEmbedding Generation Test:")
                print("  Status:     SUCCESS")
                print(f"  Shape:      {embed_result['embedding_shape']}")
                print(f"  Dimension:  {embed_result['embedding_dimension']}")
                print(f"  Time:       {embed_result['elapsed_ms']} ms")
                print(f"  Per Text:   {embed_result['ms_per_text']} ms")
                print(f"  Device:     {embed_result['device']}")

    except GPUNotAvailableError as e:
        logger.error("Device Error: %s", e)
        results["error"] = {"type": "GPU_NOT_AVAILABLE", "message": str(e)}
        if not args.json:
            print(f"\nERROR: {e}")
        exit_code = 1

    except EmbeddingModelError as e:
        logger.error("Model Error: %s", e)
        results["error"] = {
            "type": "EMBEDDING_MODEL_ERROR",
            "message": str(e),
            "model_path": e.model_path,
        }
        if not args.json:
            print(f"\nERROR: {e}")
        exit_code = 1

    except ImportError as e:
        logger.error("Import Error: %s", e)
        results["error"] = {"type": "IMPORT_ERROR", "message": str(e)}
        if not args.json:
            print(f"\nERROR: {e}")
        exit_code = 1

    except Exception as e:
        logger.exception("Unexpected error: %s", e)
        results["error"] = {"type": "UNEXPECTED_ERROR", "message": str(e)}
        if not args.json:
            print(f"\nUNEXPECTED ERROR: {e}")
        exit_code = 1

    if args.json:
        print(json.dumps(results, indent=2))

    sys.exit(exit_code)


if __name__ == "__main__":
    main()
