#!/usr/bin/env python3
"""Process monitoring for AI/ML workloads with performance optimizations."""

import logging
import os
import re
import threading
import time
from collections import OrderedDict
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, List, Optional

from ...application.process_topology import annotate_job_topology
from ...application.configuration import load_process_patterns
from ...application.process_inventory import ProcessInventoryGateway
from ...infrastructure.file_config_source import FileConfigSource
from ...infrastructure.process_inventory import (
    PsutilProcessInventoryGateway,
    psutil as _psutil_module,
)

# Backward-compatible alias for tests that monkeypatch monitor_module.psutil.
psutil = _psutil_module


class AIProcessMonitor:
    """Enhanced process monitor for AI/ML workloads with caching."""

    _CMDLINE_SENSITIVE_PREFIXES = (
        "python",
        "node",
        "java",
        "ruby",
        "perl",
        "php",
        "rscript",
    )
    _CMDLINE_SENSITIVE_NAMES = {"r", "sh", "bash", "zsh", "fish"}
    _GPU_ACTIVITY_FIELDS = (
        "sm_util",
        "gpu_util",
        "cu_occupancy",
        "mem_util",
        "vram_percent",
        "memory",
        "vram_used",
    )
    _NON_AI_GPU_NAMES = (
        "chrome",
        "chromium",
        "firefox",
        "xorg",
        "xwayland",
        "mutter",
        "kwin",
        "plasmashell",
        "gnome-shell",
        "steam",
        "discord",
        "obs",
        "mpv",
        "vlc",
    )
    _AI_GPU_HINTS = (
        "train",
        "finetune",
        "inference",
        "infer",
        "asr",
        "stt",
        "transcrib",
        "diariz",
        "predict",
        "vllm",
        "lmdeploy",
        "xinference",
        "api_server",
        "text-generation",
        "tabbyapi",
        "koboldcpp",
        "llamafile",
        "llamafactory",
        "litgpt",
        "axolotl",
        "deepspeed",
        "transformers",
        "torch",
        "whisper",
        "faster-whisper",
        "faster_whisper",
        "whisperx",
        "whisper-cpp",
        "ctranslate2",
        "tensorflow",
        "jax",
        "llama",
        "ollama",
        "invokeai",
        "comfyui",
        "fooocus",
        "triton",
        "huggingface",
        "model",
        "embed",
        "diffusion",
        "voxtype",
    )
    _PYTHON_CMDLINE_HINT_FLAGS = (
        "--api",
        "--backend",
        "--chat",
        "--device",
        "--host",
        "--listen",
        "--loader",
        "--model",
        "--task",
    )
    _PYTHON_CMDLINE_HINT_WORDS = (
        "api",
        "chat",
        "inference",
        "infer",
        "serve",
        "server",
        "train",
        "finetune",
        "transcribe",
        "diarize",
        "webui",
    )
    _HOST_CPU_ACTIVE_THRESHOLD = 0.5
    _HOST_IO_ACTIVE_THRESHOLD_MB_S = 0.05

    def __init__(
        self,
        patterns_file: Optional[Path] = None,
        cache_timeout: int = 120,
        process_inventory: Optional[ProcessInventoryGateway] = None,
    ):
        """Initialize the AI process monitor with enhanced caching.

        Args:
            patterns_file: Path to JSON file containing process patterns.
                         If None, uses default patterns file in config directory.
            cache_timeout: How long to cache compiled regexes and matches in seconds
        """
        # Logger setup
        self.logger = logging.getLogger(self.__class__.__name__)
        self._config_source = FileConfigSource()
        self._process_inventory = process_inventory or PsutilProcessInventoryGateway(
            psutil_module=psutil
        )
        self._process_errors = self._process_inventory.process_errors()

        # Thread safety - single lock protects all shared state
        # Lock ordering: Acquire this BEFORE DataCollector._cache_lock
        # to prevent deadlocks
        self._lock = threading.RLock()

        # Load and compile patterns - optimize regex compilation
        self.patterns = self._load_and_compile_patterns(patterns_file)

        # Advanced caching for process detection
        self._known_ai_processes: "OrderedDict[str, bool]" = OrderedDict()
        self._known_non_ai_processes: "OrderedDict[str, bool]" = OrderedDict()
        self._max_name_cache_size = 10000
        self._last_cache_clear = time.time()
        self._cache_timeout = cache_timeout

        # Process information cache with optimized TTL strategy
        self._process_info_cache: Dict[int, Dict[str, Any]] = {}
        self._process_info_timestamps: Dict[int, float] = {}
        self._process_info_ttl = 5.0  # Extended TTL for better performance
        self._max_cache_size = 5000  # Prevent unbounded memory growth
        self._process_io_cache: Dict[int, Dict[str, float]] = {}

        # Track stale processes for efficient cleanup
        self._last_seen_pids: set[int] = set()
        self._current_pids: set[int] = set()
        self._pid_timestamp: Dict[int, float] = {}
        self._pid_staleness_threshold = (
            30.0  # Time before considering a PID stale (seconds)
        )

        # Keep recently seen AI processes visible across short telemetry gaps.
        self._visibility_grace_seconds = 120.0
        self._ai_visibility_cache: Dict[int, Dict[str, Any]] = {}
        self._ai_visibility_timestamps: Dict[int, float] = {}

        # Performance metrics
        self._collection_stats: Dict[str, Any] = {
            "total_processes": 0,
            "ai_processes": 0,
            "grace_retained": 0,
            "job_groups": 0,
            "cache_hits": 0,
            "cache_misses": 0,
            "last_collection_time": 0,
            "pattern_matches": 0,
            "cache_size": 0,
            "cache_evictions": 0,
        }

        # Last time full process scan was performed
        self._last_full_scan = 0.0
        self._full_scan_interval = 15.0  # Seconds between full process scans

        self.logger.debug(f"Initialized with {len(self.patterns)} compiled patterns")

    def _is_cmdline_sensitive_name(self, proc_name: str) -> bool:
        """Return True for interpreters where name-only matching is too coarse."""
        name = str(proc_name).strip().lower()
        if not name:
            return False
        return name.startswith(self._CMDLINE_SENSITIVE_PREFIXES) or (
            name in self._CMDLINE_SENSITIVE_NAMES
        )

    def _build_cache_key(self, proc_name: str, cmdline: str = "") -> str:
        """Build a stable cache key with cmdline granularity when available."""
        normalized_name = str(proc_name).strip().lower()
        if not normalized_name:
            return ""

        normalized_cmdline = str(cmdline).strip().lower()
        if normalized_cmdline:
            # Cap key size to keep cache memory bounded.
            return f"{normalized_name}|{normalized_cmdline[:256]}"

        return normalized_name

    @staticmethod
    def _coerce_positive_float(value: Any) -> float:
        """Convert metric payloads to non-negative float values."""
        try:
            numeric = float(value)
        except (TypeError, ValueError):
            return 0.0
        return numeric if numeric > 0.0 else 0.0

    def _summarize_gpu_activity(
        self, gpu_processes: List[Dict[str, Any]]
    ) -> Dict[int, Dict[str, float]]:
        """Aggregate GPU activity metrics by PID."""
        aggregated: Dict[int, Dict[str, float]] = {}
        for gpu_proc in gpu_processes:
            pid_raw = gpu_proc.get("pid")
            if pid_raw is None:
                continue
            try:
                pid = int(pid_raw)
            except (TypeError, ValueError):
                continue
            if pid <= 0:
                continue

            entry = aggregated.setdefault(pid, {})
            for field in self._GPU_ACTIVITY_FIELDS:
                value = self._coerce_positive_float(gpu_proc.get(field))
                if value <= 0.0:
                    continue
                entry[field] = max(entry.get(field, 0.0), value)
        return aggregated

    def _has_meaningful_gpu_activity(self, activity: Dict[str, float]) -> bool:
        """Check whether process telemetry indicates meaningful accelerator usage."""
        if not activity:
            return False
        util_peak = max(
            activity.get("sm_util", 0.0),
            activity.get("gpu_util", 0.0),
            activity.get("cu_occupancy", 0.0),
            activity.get("mem_util", 0.0),
            activity.get("vram_percent", 0.0),
        )
        if util_peak >= 1.0:
            return True

        memory_peak_mb = max(
            activity.get("memory", 0.0), activity.get("vram_used", 0.0)
        )
        return memory_peak_mb >= 64.0

    @staticmethod
    def _normalize_process_status(raw_status: Any) -> str:
        """Normalize raw process status labels into stable display-oriented values."""
        normalized = str(raw_status).strip().lower()
        if not normalized:
            return ""
        if normalized.startswith("zomb"):
            return "zombie"
        if normalized.startswith("dead"):
            return "dead"
        if normalized.startswith("trace") or normalized.startswith("stop"):
            return "stopped"
        if normalized.startswith("disk"):
            return "diskwait"
        if normalized.startswith("run"):
            return "running"
        if "sleep" in normalized:
            return "sleeping"
        return normalized

    def _has_meaningful_host_activity(
        self,
        cpu_percent: Any,
        io_read_mb_s: Any,
        io_write_mb_s: Any,
        net_connections: Any,
    ) -> bool:
        """Return True when host-side process metrics indicate active work."""
        cpu_active = (
            self._coerce_positive_float(cpu_percent) >= self._HOST_CPU_ACTIVE_THRESHOLD
        )
        io_active = (
            self._coerce_positive_float(io_read_mb_s)
            + self._coerce_positive_float(io_write_mb_s)
            >= self._HOST_IO_ACTIVE_THRESHOLD_MB_S
        )
        try:
            net_active = int(net_connections) > 0
        except (TypeError, ValueError):
            net_active = False
        return cpu_active or io_active or net_active

    def _derive_effective_status(
        self,
        raw_status: Any,
        cpu_percent: Any,
        io_read_mb_s: Any,
        io_write_mb_s: Any,
        net_connections: Any,
        gpu_activity: Optional[Dict[str, float]] = None,
    ) -> str:
        """Derive canonical process status from raw psutil state and activity hints."""
        normalized = self._normalize_process_status(raw_status)
        if normalized in {"zombie", "dead", "stopped", "diskwait"}:
            return normalized

        if self._has_meaningful_host_activity(
            cpu_percent=cpu_percent,
            io_read_mb_s=io_read_mb_s,
            io_write_mb_s=io_write_mb_s,
            net_connections=net_connections,
        ) or self._has_meaningful_gpu_activity(gpu_activity or {}):
            return "running"

        if normalized:
            return normalized
        return "sleeping"

    def _looks_like_non_ai_gpu_process(self, proc_name: str, cmdline: str) -> bool:
        """Best-effort guardrail for common graphics/media workloads."""
        blob = f"{proc_name} {cmdline}".strip().lower()
        if not blob:
            return False
        return any(token in blob for token in self._NON_AI_GPU_NAMES)

    def _is_likely_ai_gpu_process(
        self, proc_name: str, cmdline: str, activity: Dict[str, float]
    ) -> bool:
        """Classify unmatched GPU workloads that still look like AI/ML."""
        if not self._has_meaningful_gpu_activity(activity):
            return False

        normalized_name = str(proc_name).strip().lower()
        normalized_cmdline = str(cmdline).strip().lower()
        if self._looks_like_non_ai_gpu_process(normalized_name, normalized_cmdline):
            return False

        # Interpreter-backed compute jobs are common for AI workloads.
        if self._is_cmdline_sensitive_name(normalized_name):
            return True

        blob = f"{normalized_name} {normalized_cmdline}".strip()
        if any(hint in blob for hint in self._AI_GPU_HINTS):
            return True

        return False

    def _name_cache_contains(self, cache: "OrderedDict[str, bool]", key: str) -> bool:
        """Check membership while updating LRU recency."""
        if not key:
            return False
        if key not in cache:
            return False
        cache.move_to_end(key)
        return True

    def _evict_name_caches_if_needed(self) -> None:
        """Bound process-name caches to avoid unbounded growth."""
        while (
            len(self._known_ai_processes) + len(self._known_non_ai_processes)
            > self._max_name_cache_size
        ):
            # Prefer evicting from the larger cache to keep balance.
            if len(self._known_non_ai_processes) >= len(self._known_ai_processes):
                target = self._known_non_ai_processes or self._known_ai_processes
            else:
                target = self._known_ai_processes or self._known_non_ai_processes
            if not target:
                break
            target.popitem(last=False)
            self._collection_stats["cache_evictions"] += 1

    def _name_cache_store(self, cache: "OrderedDict[str, bool]", key: str) -> None:
        """Insert/update a cache key and enforce global size bounds."""
        if not key:
            return
        cache[key] = True
        cache.move_to_end(key)
        self._evict_name_caches_if_needed()

    def _evict_pid_caches_locked(self, pid: int) -> None:
        """Remove all per-PID cache entries (caller must hold ``self._lock``)."""
        self._process_info_cache.pop(pid, None)
        self._process_info_timestamps.pop(pid, None)
        self._pid_timestamp.pop(pid, None)
        self._process_io_cache.pop(pid, None)
        self._last_seen_pids.discard(pid)

    def _load_and_compile_patterns(
        self, patterns_file: Optional[Path]
    ) -> List[re.Pattern[str]]:
        """Load and compile regex patterns efficiently with better error handling."""
        try:
            patterns = load_process_patterns(patterns_file, source=self._config_source)
            # Optimize regex compilation with flags for performance
            compiled_patterns: List[re.Pattern[str]] = []
            for pattern in patterns:
                try:
                    # Use case-insensitive matching and ASCII character classes
                    # for predictable process-name/cmdline matching.
                    compiled = re.compile(pattern, re.IGNORECASE | re.ASCII)
                    compiled_patterns.append(compiled)
                except re.error as e:
                    self.logger.warning(f"Invalid regex pattern '{pattern}': {e}")

            if not compiled_patterns:
                raise RuntimeError("No valid patterns found")

            return compiled_patterns

        except (RuntimeError, FileNotFoundError, ValueError, TypeError) as e:
            self.logger.error(f"Failed to load patterns: {e}")
            # Use default patterns if file loading fails
            # Keep fallback signatures strict to avoid noisy false positives.
            default_patterns = [
                r"\bvllm\b(?:\s+serve\b|$)",
                r"\btext-generation-launcher\b",
                r"\btritonserver\b",
                r"\btrtllm-serve\b",
                r"\bllama-server\b",
                r"\bllamafile\b(?:\s+--server\b|$)",
                r"\bollama\b(?:\s+(?:serve|run)\b|$)",
                r"\bsglang\b(?:\s+launch_server\b|$)",
                r"\blmdeploy\b\s+serve\s+api_server\b",
                r"\bxinference(?:-local)?\b",
                r"\btext-generation-router\b",
                r"\binvokeai(?:-web)?\b",
                r"\btabbyapi\b",
                r"\bkoboldcpp(?:\.exe)?\b",
                r"\bllamafactory-cli\b",
                r"\baxolotl\b\s+(?:train|inference|serve|fetch)\b",
                r"\blitgpt\b\s+(?:serve|chat|finetune|pretrain)\b",
                r"\bvoxtype\b",
                r"\binsanely[-_]?fast[-_]?whisper\b",
                (
                    r"\b(?:faster[-_]?whisper|whisper(?:x|[-_]?cpp|\.cpp)?|"
                    r"ctranslate2|speechbrain|vosk)\b"
                ),
                r"\bdeepspeed\b",
                r"\baccelerate\b\s+launch\b",
                r"\btorchrun\b",
                r"\bpython:(?:[^ ]*/)?text[-_]?generation[-_]?webui/server\.py\b",
                r"\bpython:(?:[^ ]*/)?fooocus/entry_with_update\.py\b",
                r"\bpython:(?:[^ ]*/)?tabbyapi/main\.py\b",
                (
                    r"\bpython:(?:[^ ]*/)?server\.py\b"
                    r"(?=.*\b--api\b)"
                    r"(?=.*\b--(?:loader|model)\b).*"
                ),
                r"(?=.*\b(?:asr|speech[-_]?to[-_]?text|stt|transcrib(?:e|er|ing)|"
                r"diari[sz](?:e|ation))\b)"
                r"(?=.*\b(?:cuda|cublas|cudnn|torch|onnxruntime|ctranslate2|"
                r"whisper)\b).*",
                r"\bpython:[^ ]*(?:train|finetune|infer|inference|llm|diffusion)"
                r"[^ ]*\.py\b",
                r"(?=.*\b(?:train|finetune|inference|predict|serve|generate)\b)"
                r"(?=.*\b(?:torch|tensorflow|jax|transformers|vllm|llama|"
                r"deepspeed)\b).*",
            ]
            self.logger.info(f"Using {len(default_patterns)} default patterns")
            return [re.compile(p, re.IGNORECASE | re.ASCII) for p in default_patterns]

    @lru_cache(maxsize=2048)
    def _is_ai_name_cached(self, proc_name: str) -> bool:
        """LRU cached check for AI process name pattern matching.

        Uses increased cache size for better performance.

        Note: This function does NOT acquire locks because:
        1. self.patterns is read-only after __init__
        2. LRU cache decorator is thread-safe
        3. Pattern matching is a pure function (no state modification)
        4. Stats are updated by caller under lock
        """
        # Early return for empty or very short names
        if not proc_name or len(proc_name) < 2:
            return False

        # Check each pattern (patterns are read-only, no lock needed)
        for pattern in self.patterns:
            try:
                if pattern.search(proc_name):
                    return True
            except Exception as e:
                self.logger.warning(f"Pattern matching error for '{proc_name}': {e}")

        return False

    def is_ai_process(self, proc_name: str, cmdline: str = "") -> bool:
        """Check if a process name or command line matches AI/ML patterns.

        Includes enhanced caching.

        Args:
            proc_name: Name of the process
            cmdline: Full command line of the process (optional)

        Returns:
            bool: True if the process matches any AI/ML patterns
        """
        # Sanitize inputs
        if not proc_name:
            return False

        sanitized_name = str(proc_name).strip()
        sanitized_cmdline = str(cmdline).strip() if cmdline else ""
        cache_key = self._build_cache_key(sanitized_name, sanitized_cmdline)

        with self._lock:
            # Check if we should clear caches
            current_time = time.time()
            if current_time - self._last_cache_clear > self._cache_timeout:
                self._known_ai_processes.clear()
                self._known_non_ai_processes.clear()
                self._last_cache_clear = current_time
                self.logger.debug("Cleared process name cache")
                self._collection_stats["cache_evictions"] += 1

            # Check caches first for quick lookup
            if self._name_cache_contains(self._known_ai_processes, cache_key):
                return True

            if self._name_cache_contains(self._known_non_ai_processes, cache_key):
                return False

            # Update cache size metric
            cache_size = len(self._known_ai_processes) + len(
                self._known_non_ai_processes
            )
            self._collection_stats["cache_size"] = cache_size

            self._evict_name_caches_if_needed()

        # Check if the name matches AI patterns
        name_match = False
        try:
            name_match = self._is_ai_name_cached(sanitized_name)
        except Exception as e:
            self.logger.warning(f"Error checking process name '{sanitized_name}': {e}")

        if name_match:
            with self._lock:
                self._name_cache_store(self._known_ai_processes, cache_key)
                self._collection_stats["pattern_matches"] += 1
            return True

        # Check cmdline if provided
        if sanitized_cmdline:
            try:
                with self._lock:
                    cmdline_match = any(
                        pattern.search(sanitized_cmdline) for pattern in self.patterns
                    )
                if cmdline_match:
                    with self._lock:
                        self._name_cache_store(self._known_ai_processes, cache_key)
                    return True
            except Exception as e:
                self.logger.warning(
                    f"Error checking cmdline '{sanitized_cmdline[:50]}...': {e}"
                )

        # Not an AI process
        with self._lock:
            self._name_cache_store(self._known_non_ai_processes, cache_key)
        return False

    def _get_process_cmdline(self, proc: Any) -> str:
        """Efficiently get process command line with improved error handling."""
        try:
            with self._process_inventory.process_oneshot(
                proc
            ):  # Batch system calls for efficiency
                try:
                    cmdline = self._process_inventory.process_cmdline(proc)
                    proc_name = self._process_inventory.process_name(proc)
                except self._process_errors:
                    return ""  # Process no longer exists or not accessible

                # Handle empty command line
                if not cmdline:
                    return proc_name

                # Get the script name if it's a Python process
                if proc_name.lower().startswith("python"):
                    # Find the .py file in the command line
                    script_name = ""
                    for arg in cmdline:
                        if arg.endswith(".py"):
                            script_name = arg
                            break

                    if script_name:
                        normalized_script = script_name.replace("\\", "/")
                        script_basename = os.path.basename(normalized_script)
                        script_tokens = [
                            token for token in normalized_script.split("/") if token
                        ]
                        if len(script_tokens) >= 2:
                            script_display = f"{script_tokens[-2]}/{script_basename}"
                        else:
                            script_display = script_basename

                        # Keep only a few high-signal tokens to avoid giant cache keys.
                        # This preserves enough context for project-specific signatures.
                        seen_hints: set[str] = set()
                        cmdline_hints: list[str] = []
                        for arg in cmdline:
                            lowered = str(arg).strip().lower()
                            if not lowered or lowered == script_name.lower():
                                continue
                            if lowered in self._PYTHON_CMDLINE_HINT_WORDS or (
                                lowered.startswith("--")
                                and lowered in self._PYTHON_CMDLINE_HINT_FLAGS
                            ):
                                if lowered not in seen_hints:
                                    cmdline_hints.append(lowered)
                                    seen_hints.add(lowered)
                            if len(cmdline_hints) >= 3:
                                break

                        if cmdline_hints:
                            return f"python:{script_display} {' '.join(cmdline_hints)}"
                        return f"python:{script_display}"

                # Join command line with spaces, limiting length for performance
                full_cmd = " ".join(cmdline)
                if len(full_cmd) > 500:  # Limit very long command lines
                    return full_cmd[:497] + "..."
                return full_cmd

        except self._process_errors:
            try:
                return self._process_inventory.process_name(proc)
            except Exception:
                # Catch all exceptions as process may have terminated during access
                return ""  # Fallback for any other errors
        except Exception as e:
            self.logger.debug(f"Error getting cmdline: {e}")
            try:
                return self._process_inventory.process_name(proc)
            except Exception:
                # Catch all exceptions as process may have terminated during access
                return ""

    def _compute_io_rates(
        self, pid: int, current_time: float, read_bytes: float, write_bytes: float
    ) -> Dict[str, float]:
        """Compute process IO throughput in MB/s from cumulative counters."""
        with self._lock:
            previous = self._process_io_cache.get(pid)
            self._process_io_cache[pid] = {
                "timestamp": current_time,
                "read_bytes": read_bytes,
                "write_bytes": write_bytes,
            }

        if not previous:
            return {"io_read_mb_s": 0.0, "io_write_mb_s": 0.0}

        elapsed = max(0.001, current_time - previous["timestamp"])
        read_delta = max(0.0, read_bytes - previous["read_bytes"])
        write_delta = max(0.0, write_bytes - previous["write_bytes"])

        return {
            "io_read_mb_s": read_delta / elapsed / (1024 * 1024),
            "io_write_mb_s": write_delta / elapsed / (1024 * 1024),
        }

    def _get_process_info(
        self, proc: Any, current_time: float, include_extended: bool = True
    ) -> Dict[str, Any]:
        """Get detailed information about a process with enhanced caching.

        Args:
            proc: Process to get information for
            current_time: Current timestamp
            include_extended: Whether to include extended metrics
                (can be disabled for performance)

        Returns:
            Dict containing process information
        """
        try:
            # Get PID first with error handling
            pid = self._process_inventory.process_pid(proc)
        except self._process_errors:
            # Process no longer exists or can't be accessed
            return {}

        # Check if we have cached info for this PID
        with self._lock:
            if pid in self._process_info_cache:
                last_update = self._process_info_timestamps.get(pid, 0)
                # Use cache if it's fresh enough
                if current_time - last_update < self._process_info_ttl:
                    # Update seen PIDs
                    self._last_seen_pids.add(pid)
                    self._pid_timestamp[pid] = current_time
                    return self._process_info_cache[pid]

        # Collect new info
        try:
            with self._process_inventory.process_oneshot(
                proc
            ):  # Batch system calls for efficiency
                try:
                    # Basic info - always collected
                    raw_status = self._process_inventory.process_status(proc)
                    info = {
                        "pid": pid,
                        "ppid": self._process_inventory.process_ppid(proc),
                        "name": self._process_inventory.process_name(proc),
                        "username": self._process_inventory.process_username(proc),
                        "cmdline": self._get_process_cmdline(proc),
                        "cpu_percent": self._process_inventory.process_cpu_percent(
                            proc
                        ),
                        "memory_percent": self._process_inventory.process_memory_percent(
                            proc
                        ),
                        "raw_status": raw_status,
                        "create_time": self._process_inventory.process_create_time(
                            proc
                        ),
                    }

                    # Extended info - optional
                    if include_extended:
                        try:
                            cpu_times = self._process_inventory.process_cpu_times(proc)
                            memory_info = self._process_inventory.process_memory_info(
                                proc
                            )

                            # Add extended metrics
                            info.update(
                                {
                                    "cpu_time_user": cpu_times.user,
                                    "cpu_time_system": cpu_times.system,
                                    "memory_rss": memory_info.rss,
                                    "memory_vms": memory_info.vms,
                                    "num_threads": self._process_inventory.process_num_threads(
                                        proc
                                    ),
                                }
                            )

                            # Process disk IO throughput (MB/s)
                            try:
                                io_counters = (
                                    self._process_inventory.process_io_counters(proc)
                                )
                                read_bytes = float(io_counters.read_bytes)
                                write_bytes = float(io_counters.write_bytes)
                                info.update(
                                    self._compute_io_rates(
                                        pid, current_time, read_bytes, write_bytes
                                    )
                                )
                            except self._process_errors + (AttributeError,):
                                info["io_read_mb_s"] = 0.0
                                info["io_write_mb_s"] = 0.0

                            # Lightweight network activity proxy:
                            # active INET connection count.
                            try:
                                info[
                                    "net_connections"
                                ] = self._process_inventory.process_net_connection_count(
                                    proc
                                )
                            except self._process_errors + (
                                AttributeError,
                                ValueError,
                                OSError,
                            ):
                                info["net_connections"] = 0

                            # Add CPU affinity if available
                            try:
                                info[
                                    "cpu_affinity"
                                ] = self._process_inventory.process_cpu_affinity(proc)
                            except self._process_errors + (AttributeError,):
                                # CPU affinity might not be available on all platforms
                                info["cpu_affinity"] = []
                        except self._process_errors:
                            # Skip extended info if not accessible
                            pass

                    # Ensure optional metrics always exist for UI rendering.
                    info.setdefault("io_read_mb_s", 0.0)
                    info.setdefault("io_write_mb_s", 0.0)
                    info.setdefault("net_connections", 0)
                    info["status"] = self._derive_effective_status(
                        raw_status=info.get("raw_status", ""),
                        cpu_percent=info.get("cpu_percent", 0.0),
                        io_read_mb_s=info.get("io_read_mb_s", 0.0),
                        io_write_mb_s=info.get("io_write_mb_s", 0.0),
                        net_connections=info.get("net_connections", 0),
                        gpu_activity={},
                    )
                except self._process_errors:
                    return {}

                # Update cache
                with self._lock:
                    # Check cache size and evict if necessary
                    if len(self._process_info_cache) >= self._max_cache_size:
                        # Remove oldest entries
                        eviction_batch = max(1, self._max_cache_size // 10)
                        oldest_entries = sorted(
                            self._process_info_timestamps.items(), key=lambda x: x[1]
                        )[:eviction_batch]

                        for old_pid, _ in oldest_entries:
                            self._evict_pid_caches_locked(old_pid)
                            self._collection_stats["cache_evictions"] += 1

                    # Update cache
                    self._process_info_cache[pid] = info
                    self._process_info_timestamps[pid] = current_time

                    # Update tracking
                    self._last_seen_pids.add(pid)
                    self._pid_timestamp[pid] = current_time

                return info
        except self._process_errors:
            return {}
        except Exception as e:
            pid_str = getattr(proc, "pid", "unknown")
            self.logger.warning(f"Error getting process info for PID {pid_str}: {e}")
            return {}

    def _prune_stale_data(self, current_time: float) -> None:
        """Prune stale data from caches efficiently."""
        try:
            current_pids = self._process_inventory.snapshot_current_pids()
        except self._process_errors + (OSError, RuntimeError):
            # Skip prune cycle on transient psutil failures.
            return

        with self._lock:
            # Track current PIDs for staleness detection
            self._current_pids = current_pids

            # Find PIDs that no longer exist
            stale_pids = set(self._process_info_cache.keys()) - current_pids

            # Also find PIDs that haven't been seen recently
            for pid in list(self._pid_timestamp.keys()):
                if (
                    current_time - self._pid_timestamp.get(pid, 0)
                    > self._pid_staleness_threshold
                ):
                    stale_pids.add(pid)

            # Remove stale entries
            for pid in stale_pids:
                self._evict_pid_caches_locked(pid)

            # Clear last seen PIDs set periodically
            if current_time - self._last_cache_clear > self._cache_timeout:
                self._last_seen_pids.clear()

    def _is_cached_process_alive(self, cached_process: Dict[str, Any]) -> bool:
        """Check whether a cached process snapshot still refers to a live PID."""
        pid_raw = cached_process.get("pid")
        if not isinstance(pid_raw, int):
            return False

        expected_create_time = cached_process.get("create_time")
        try:
            proc = self._process_inventory.get_process(pid_raw)
            if not self._process_inventory.process_is_running(proc):
                return False
            if expected_create_time is None:
                return True
            current_create_time = self._process_inventory.process_create_time(proc)
            return abs(float(current_create_time) - float(expected_create_time)) < 1.0
        except self._process_errors + (ValueError, OSError, TypeError):
            return False

    def _apply_visibility_grace_window(
        self, current_processes: List[Dict[str, Any]], current_time: float
    ) -> List[Dict[str, Any]]:
        """Retain recently seen AI processes to smooth transient collector gaps."""
        merged: List[Dict[str, Any]] = []
        current_by_pid: Dict[int, Dict[str, Any]] = {}
        passthrough: List[Dict[str, Any]] = []
        for process in current_processes:
            pid = process.get("pid")
            if isinstance(pid, int):
                current_by_pid[pid] = process
            else:
                passthrough.append(process.copy())

        with self._lock:
            merged.extend(passthrough)

            # Refresh visibility cache with fresh processes.
            for pid, process in current_by_pid.items():
                process_copy = process.copy()
                process_copy.setdefault(
                    "raw_status", str(process_copy.get("status", "")).strip()
                )
                process_copy["grace_retained"] = False
                process_copy["grace_age_s"] = 0.0
                self._ai_visibility_cache[pid] = process_copy
                self._ai_visibility_timestamps[pid] = current_time
                merged.append(process_copy)

            # Add recently seen-but-missing processes while grace is valid.
            retained_count = 0
            for pid in list(self._ai_visibility_cache.keys()):
                if pid in current_by_pid:
                    continue

                cached = self._ai_visibility_cache.get(pid)
                if not cached:
                    continue

                age = current_time - self._ai_visibility_timestamps.get(pid, 0.0)
                if age > self._visibility_grace_seconds:
                    self._ai_visibility_cache.pop(pid, None)
                    self._ai_visibility_timestamps.pop(pid, None)
                    continue

                if not self._is_cached_process_alive(cached):
                    self._ai_visibility_cache.pop(pid, None)
                    self._ai_visibility_timestamps.pop(pid, None)
                    continue

                retained = cached.copy()
                retained["cpu_percent"] = 0.0
                retained["io_read_mb_s"] = 0.0
                retained["io_write_mb_s"] = 0.0
                retained["net_connections"] = 0
                retained["grace_retained"] = True
                retained["grace_age_s"] = round(max(0.0, age), 1)
                retained.setdefault(
                    "raw_status",
                    str(retained.get("status", "")).strip(),
                )
                if not str(retained.get("status", "")).strip():
                    retained["status"] = self._derive_effective_status(
                        raw_status=retained.get("raw_status", ""),
                        cpu_percent=retained.get("cpu_percent", 0.0),
                        io_read_mb_s=retained.get("io_read_mb_s", 0.0),
                        io_write_mb_s=retained.get("io_write_mb_s", 0.0),
                        net_connections=retained.get("net_connections", 0),
                        gpu_activity={},
                    )
                merged.append(retained)
                retained_count += 1

            self._collection_stats["grace_retained"] = retained_count

        return merged

    def get_ai_processes(
        self,
        gpu_processes: Optional[List[Dict[str, Any]]] = None,
        force_refresh: bool = False,
    ) -> List[Dict[str, Any]]:
        """Get information about AI-related processes with advanced caching.

        Args:
            gpu_processes: Optional list of GPU processes from vendors
            force_refresh: Whether to force refresh of process cache

        Returns:
            List of AI process information dictionaries
        """
        start_time = time.time()
        ai_processes: List[Dict[str, Any]] = []
        gpu_pids: set[int] = set()
        gpu_process_list = gpu_processes or []
        gpu_activity_by_pid = self._summarize_gpu_activity(gpu_process_list)
        cache_hits = 0
        cache_misses = 0

        # Check if we're due for full scan
        perform_full_scan = (
            force_refresh
            or time.time() - self._last_full_scan >= self._full_scan_interval
        )

        # First collect GPU process PIDs
        if gpu_process_list:
            gpu_pids = set(gpu_activity_by_pid.keys())

        # Prune stale data
        if perform_full_scan:
            self._prune_stale_data(start_time)
            self._last_full_scan = start_time

            # Force refresh if requested
            if force_refresh:
                with self._lock:
                    self._process_info_cache.clear()
                    self._process_info_timestamps.clear()
                    self._process_io_cache.clear()
                    self._pid_timestamp.clear()
                    self._last_seen_pids.clear()
                    self.logger.debug("Forced cache refresh")

        # Process iterator with batched system calls
        try:
            for ps_proc in self._process_inventory.iter_processes(["pid", "name"]):
                try:
                    # Get basic info from iterator
                    pid = self._process_inventory.process_info_pid(ps_proc)
                    name = self._process_inventory.process_info_name(ps_proc)

                    # Skip non-GPU processes already known to be non-AI
                    # (unless doing a full scan)
                    with self._lock:
                        cache_key = self._build_cache_key(name)
                        is_known_non_ai = (
                            not self._is_cmdline_sensitive_name(name)
                            and cache_key in self._known_non_ai_processes
                        )
                        skip_process = (
                            not perform_full_scan
                            and pid not in gpu_pids
                            and is_known_non_ai
                        )
                        if skip_process:
                            continue

                    # Check if it's an AI process
                    is_ai = False
                    cmdline = ""

                    # First check known AI processes cache
                    with self._lock:
                        cache_key = self._build_cache_key(name)
                        if (
                            not self._is_cmdline_sensitive_name(name)
                            and cache_key in self._known_ai_processes
                        ):
                            is_ai = True

                    # If unknown and not flagged as AI, check against patterns
                    should_check_patterns = (
                        pid in gpu_pids
                        or perform_full_scan
                        or self._is_cmdline_sensitive_name(name)
                    )
                    if not is_ai and should_check_patterns:
                        # Get command line for more accurate matching
                        cmdline = self._get_process_cmdline(ps_proc)
                        is_ai = self.is_ai_process(name, cmdline)

                    # If matching missed but the workload is clearly active on GPU,
                    # classify likely AI workloads conservatively.
                    if not is_ai and pid in gpu_pids:
                        activity = gpu_activity_by_pid.get(pid, {})
                        if self._is_likely_ai_gpu_process(name, cmdline, activity):
                            is_ai = True

                    # Collect info for matched AI processes only.
                    if is_ai:
                        # Check cache first
                        info = None
                        with self._lock:
                            if pid in self._process_info_cache:
                                last_update = self._process_info_timestamps.get(pid, 0)
                                if start_time - last_update < self._process_info_ttl:
                                    info = self._process_info_cache[pid]
                                    cache_hits += 1

                        # Cache miss or force refresh - collect new info
                        if info is None:
                            cache_misses += 1
                            # Get full info for GPU processes, basic info for others
                            # to improve performance on systems with many
                            # processes
                            include_extended = True
                            info = self._get_process_info(
                                ps_proc, start_time, include_extended
                            )

                        # Add to results list
                        if info:
                            activity = gpu_activity_by_pid.get(pid, {})
                            raw_status = str(
                                info.get("raw_status", info.get("status", ""))
                            ).strip()
                            info["raw_status"] = raw_status
                            info["status"] = self._derive_effective_status(
                                raw_status=raw_status,
                                cpu_percent=info.get("cpu_percent", 0.0),
                                io_read_mb_s=info.get("io_read_mb_s", 0.0),
                                io_write_mb_s=info.get("io_write_mb_s", 0.0),
                                net_connections=info.get("net_connections", 0),
                                gpu_activity=activity,
                            )

                            # Process command line intelligently
                            if "cmdline" in info and info["cmdline"]:
                                # Get the original cmdline for possible truncation
                                original_cmdline = info["cmdline"]
                                name_display = original_cmdline

                                # Handle different process types
                                if name_display.startswith("python "):
                                    # For Python: extract script name
                                    script_parts = [
                                        part
                                        for part in name_display.split()
                                        if part.endswith(".py")
                                    ]
                                    if script_parts:
                                        # Extract just basename + first argument if any
                                        script_name = os.path.basename(script_parts[0])
                                        args = name_display.split(script_parts[0], 1)
                                        arg_summary = ""
                                        if len(args) > 1 and args[1].strip():
                                            # Include just the first argument
                                            first_arg = (
                                                args[1].strip().split()[0]
                                                if args[1].strip().split()
                                                else ""
                                            )
                                            if first_arg:
                                                arg_summary = f" {first_arg}"
                                        name_display = (
                                            f"python:{script_name}{arg_summary}"
                                        )
                                else:
                                    # For other processes: extract base command
                                    # and first argument
                                    parts = name_display.split()
                                    if parts:
                                        # Get base command without path
                                        base_cmd = os.path.basename(parts[0])
                                        # Include just first meaningful argument
                                        if len(parts) > 1:
                                            name_display = f"{base_cmd} {parts[1]}"
                                        else:
                                            name_display = base_cmd

                                # Truncate if still too long
                                max_display_length = 60  # A reasonable compromise
                                if len(name_display) > max_display_length:
                                    name_display = (
                                        name_display[: max_display_length - 3] + "..."
                                    )

                                # Store both in the info dict
                                info["cmdline_summary"] = name_display
                                info["cmdline"] = original_cmdline

                            # Add a flag for GPU processes
                            info["is_gpu_process"] = pid in gpu_pids
                            ai_processes.append(info)

                except self._process_errors:
                    continue
                except Exception as e:
                    pid_str = "unknown"
                    try:
                        pid_str = str(self._process_inventory.process_pid(ps_proc))
                    except Exception:
                        pass
                    self.logger.warning(f"Error processing PID {pid_str}: {e}")
                    continue
        except self._process_errors + (OSError, RuntimeError) as e:
            self.logger.warning(f"Process iteration failed: {e}")

        # Keep recently seen AI processes visible across short detection gaps,
        # then annotate grouped topology + straggler metadata.
        ai_processes = self._apply_visibility_grace_window(ai_processes, start_time)
        try:
            ai_processes = annotate_job_topology(ai_processes, gpu_process_list)
        except Exception as e:
            self.logger.warning(f"Job topology annotation failed: {e}")

        # Update collection stats
        elapsed = time.time() - start_time
        try:
            total_processes = len(self._process_inventory.list_pids())
        except self._process_errors + (OSError, RuntimeError):
            total_processes = len(self._current_pids)

        with self._lock:
            self._collection_stats.update(
                {
                    "total_processes": total_processes,
                    "ai_processes": len(ai_processes),
                    "job_groups": len(
                        {
                            str(proc.get("job_id", "")).strip()
                            for proc in ai_processes
                            if proc.get("job_id")
                        }
                    ),
                    "cache_hits": cache_hits,
                    "cache_misses": cache_misses,
                    "last_collection_time": elapsed,
                    "cache_size": len(self._process_info_cache),
                }
            )

        self.logger.debug(
            f"Process collection: {len(ai_processes)} AI processes, "
            f"{cache_hits} cache hits, {cache_misses} misses, "
            f"{elapsed:.3f}s elapsed"
        )

        return ai_processes

    def get_process_gpu_usage(
        self, pid: int, gpu_processes: List[Dict[str, Any]]
    ) -> Dict[str, float]:
        """Get GPU usage information for a specific process efficiently."""
        # Create a lookup table for quick access if we have many processes
        if len(gpu_processes) > 10:
            lookup = {}
            for proc in gpu_processes:
                if proc.get("pid") == pid:
                    # Return first match with non-zero values
                    gpu_util = proc.get("gpu_util", 0.0)
                    vram_percent = proc.get("vram_percent", 0.0)

                    if gpu_util > 0 or vram_percent > 0:
                        return {"gpu_util": gpu_util, "vram_percent": vram_percent}

                    # Otherwise store in lookup
                    lookup[pid] = {"gpu_util": gpu_util, "vram_percent": vram_percent}

            # Return from lookup if we found any match
            if pid in lookup:
                return lookup[pid]
        else:
            # For small lists, just iterate
            for proc in gpu_processes:
                if proc.get("pid") == pid:
                    return {
                        "gpu_util": proc.get("gpu_util", 0.0),
                        "vram_percent": proc.get("vram_percent", 0.0),
                    }

        # No match found
        return {"gpu_util": 0.0, "vram_percent": 0.0}

    def clear_caches(self) -> None:
        """Clear all caches - useful when system state changes significantly."""
        with self._lock:
            self._known_ai_processes.clear()
            self._known_non_ai_processes.clear()
            self._process_info_cache.clear()
            self._process_info_timestamps.clear()
            self._ai_visibility_cache.clear()
            self._ai_visibility_timestamps.clear()
            self._last_cache_clear = time.time()

        # Also clear the LRU cache
        self._is_ai_name_cached.cache_clear()

        self.logger.debug("All process caches cleared")

    def get_collection_stats(self) -> Dict[str, Any]:
        """Get statistics about the last collection operation."""
        with self._lock:
            return self._collection_stats.copy()
