"""
System Optimizer — translates JEPA/atlas knowledge into OS-level tuning.

Runs on each mesh node. Two modes:
1. System-wide: sysctl, governor, hugepages, etc. (runs once + every 6h)
2. Process-level: NUMA bind, IO class, nice (event-driven via proc connector,
   fallback to polling every 5s)

Most operations require root. If not running as root, attempts sudo.
"""

import glob
import json
import logging
import os
import platform
import re
import socket
import struct
import subprocess
import shutil
import threading
import time
from pathlib import Path
from typing import Callable, Dict, List, Optional, Set

IS_LINUX = platform.system() == "Linux"

try:
    import requests
except ImportError:
    requests = None  # type: ignore

logger = logging.getLogger(__name__)


class SystemOptimizer:
    """Applies OS-level tuning derived from atlas/JEPA recommendations."""

    def __init__(self, node_name: str, controller_url: str = "http://localhost:8401"):
        self.node_name = node_name
        self.controller_url = controller_url.rstrip("/")
        self._is_root = os.geteuid() == 0
        self._results: Dict[str, str] = {}
        logger.info(
            "SystemOptimizer initialized for node=%s root=%s controller=%s",
            node_name, self._is_root, controller_url,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _write_sysfs(self, path: str, value: str, dry_run: bool = False) -> str:
        """Write a value to a sysfs/procfs path. Returns status string."""
        if dry_run:
            return f"DRY_RUN: would write '{value}' -> {path}"
        try:
            if self._is_root:
                Path(path).write_text(value)
            else:
                subprocess.run(
                    ["sudo", "tee", path],
                    input=value.encode(),
                    stdout=subprocess.DEVNULL,
                    check=True,
                )
            return f"OK: wrote '{value}' -> {path}"
        except (PermissionError, subprocess.CalledProcessError) as exc:
            msg = f"FAIL: {path} <- '{value}': {exc}"
            logger.warning(msg)
            return msg
        except FileNotFoundError:
            msg = f"SKIP: {path} does not exist"
            logger.info(msg)
            return msg

    def _read_sysfs(self, path: str) -> Optional[str]:
        """Read a sysfs/procfs value. Returns None on error."""
        try:
            return Path(path).read_text().strip()
        except Exception:
            return None

    def _sysctl_write(self, key: str, value: str, dry_run: bool = False) -> str:
        """Write a sysctl parameter (key like 'vm.swappiness')."""
        proc_path = "/proc/sys/" + key.replace(".", "/")
        return self._write_sysfs(proc_path, value, dry_run=dry_run)

    def _run_cmd(self, cmd: List[str], dry_run: bool = False) -> str:
        """Run a command, optionally prefixed with sudo."""
        if dry_run:
            return f"DRY_RUN: would run {' '.join(cmd)}"
        if not self._is_root:
            cmd = ["sudo"] + cmd
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            return f"OK: {' '.join(cmd)}: {result.stdout.strip()}"
        except (subprocess.CalledProcessError, FileNotFoundError) as exc:
            msg = f"FAIL: {' '.join(cmd)}: {exc}"
            logger.warning(msg)
            return msg

    # ------------------------------------------------------------------
    # Fetch / generate recommendations
    # ------------------------------------------------------------------

    def fetch_recommendations(self) -> Dict:
        """GET recommendations from the controller, or return empty dict on failure."""
        if requests is None:
            logger.warning("requests library not available; cannot fetch recommendations")
            return {}
        url = f"{self.controller_url}/nodes/{self.node_name}/recommendations"
        try:
            resp = requests.get(url, timeout=5)
            resp.raise_for_status()
            return resp.json()
        except Exception as exc:
            logger.warning("Failed to fetch recommendations from %s: %s", url, exc)
            return {}

    def generate_recommendations(self, atlas_data: Dict, hardware: Dict) -> Dict:
        """
        Given atlas results and hardware info, compute optimal OS settings.

        hardware keys used:
            ram_gb, numa_nodes, storage (list of {device, type}),
            nic_speed_gbps, gpu_numa_node
        atlas_data keys used:
            workload_type ("memory_bound" | "compute_bound" | "mixed")
        """
        rec: Dict = {}
        ram_gb = hardware.get("ram_gb", 0)
        numa_nodes = hardware.get("numa_nodes", 1)
        storage = hardware.get("storage", [])
        nic_speed = hardware.get("nic_speed_gbps", 1)
        workload = atlas_data.get("workload_type", "mixed")

        # --- CPU governor ---
        rec["cpu_governor"] = "performance"

        # --- Hugepages ---
        if ram_gb > 64:
            rec["hugepages"] = max(2048, int(ram_gb * 16))  # ~2% of RAM in 2M pages
            rec["swappiness"] = 1
        else:
            rec["hugepages"] = 1024
            rec["swappiness"] = 10

        # --- NUMA ---
        if numa_nodes > 1:
            rec["numa_balancing"] = True
            rec["pin_irqs"] = True
        else:
            rec["numa_balancing"] = False
            rec["pin_irqs"] = False

        # --- IO scheduler per device ---
        io_scheds = {}
        for dev in storage:
            name = dev.get("device", "")
            dtype = dev.get("type", "unknown").lower()
            if dtype == "nvme":
                io_scheds[name] = "none"
            elif dtype in ("ssd", "flash"):
                io_scheds[name] = "mq-deadline"
            else:
                io_scheds[name] = "mq-deadline"
        rec["io_schedulers"] = io_scheds

        # --- Network ---
        if nic_speed >= 10:
            rec["rmem_max"] = 33554432
            rec["wmem_max"] = 33554432
            rec["tcp_rmem"] = "4096 262144 33554432"
            rec["tcp_wmem"] = "4096 131072 33554432"
            rec["somaxconn"] = 65535
        else:
            rec["rmem_max"] = 16777216
            rec["wmem_max"] = 16777216
            rec["tcp_rmem"] = "4096 131072 16777216"
            rec["tcp_wmem"] = "4096 65536 16777216"
            rec["somaxconn"] = 4096

        # --- Workload-specific ---
        if workload == "memory_bound":
            rec["hugepages"] = max(rec.get("hugepages", 1024), 4096)
            rec["vfs_cache_pressure"] = 30
            rec["dirty_ratio"] = 30
            rec["dirty_background_ratio"] = 10
            rec["transparent_hugepages"] = "always"
        elif workload == "compute_bound":
            rec["vfs_cache_pressure"] = 50
            rec["dirty_ratio"] = 20
            rec["dirty_background_ratio"] = 5
            rec["sched_min_granularity_ns"] = 3000000
            rec["sched_wakeup_granularity_ns"] = 4000000
            rec["transparent_hugepages"] = "madvise"
        else:
            rec["vfs_cache_pressure"] = 50
            rec["dirty_ratio"] = 20
            rec["dirty_background_ratio"] = 5
            rec["transparent_hugepages"] = "madvise"

        return rec

    # ------------------------------------------------------------------
    # Apply individual tunings
    # ------------------------------------------------------------------

    def apply_cpu_governor(self, governor: str = "performance", dry_run: bool = False) -> str:
        """Set CPU frequency governor for all cores."""
        results = []
        paths = sorted(glob.glob("/sys/devices/system/cpu/cpu*/cpufreq/scaling_governor"))
        if not paths:
            msg = "SKIP: no cpufreq governors found (VM or driver not loaded)"
            logger.info(msg)
            return msg
        for path in paths:
            results.append(self._write_sysfs(path, governor, dry_run=dry_run))
        failures = [r for r in results if r.startswith("FAIL")]
        if failures:
            return f"PARTIAL: {len(results) - len(failures)}/{len(results)} cores set to {governor}"
        return f"OK: {len(results)} cores set to governor={governor}"

    def apply_hugepages(self, pages: int = 1024, size: str = "2M", dry_run: bool = False) -> str:
        """Allocate hugepages."""
        if size == "2M":
            path = "/sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages"
        elif size == "1G":
            path = "/sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages"
        else:
            path = "/proc/sys/vm/nr_hugepages"
        return self._write_sysfs(path, str(pages), dry_run=dry_run)

    def apply_numa_balancing(self, enabled: bool = True, dry_run: bool = False) -> str:
        """Enable or disable automatic NUMA balancing."""
        value = "1" if enabled else "0"
        return self._write_sysfs("/proc/sys/kernel/numa_balancing", value, dry_run=dry_run)

    def apply_vm_tuning(
        self,
        swappiness: int = 10,
        dirty_ratio: int = 20,
        dirty_background_ratio: int = 5,
        vfs_cache_pressure: int = 50,
        dry_run: bool = False,
    ) -> str:
        """Apply VM / memory management tuning."""
        settings = {
            "vm.swappiness": str(swappiness),
            "vm.dirty_ratio": str(dirty_ratio),
            "vm.dirty_background_ratio": str(dirty_background_ratio),
            "vm.vfs_cache_pressure": str(vfs_cache_pressure),
        }
        results = []
        for key, val in settings.items():
            results.append(self._sysctl_write(key, val, dry_run=dry_run))
        return "; ".join(results)

    def apply_scheduler_tuning(
        self,
        autogroup: bool = False,
        util_clamp_min: int = 0,
        util_clamp_max: int = 1024,
        rr_timeslice_ms: int = 100,
        dry_run: bool = False,
    ) -> str:
        """Tune EEVDF/CFS scheduler parameters.

        Kernel 6.6+ uses EEVDF — the old CFS granularity knobs are gone.
        Available tunables:
          - sched_autogroup_enabled: group tasks by TTY (disable for servers)
          - sched_util_clamp_min/max: capacity hints (0-1024)
          - sched_rr_timeslice_ms: round-robin slice for SCHED_RR tasks
        """
        results = []
        # Disable autogroup for compute servers (prevents tty-based grouping
        # from interfering with explicit CPU affinity and priority)
        results.append(self._write_sysfs(
            "/proc/sys/kernel/sched_autogroup_enabled",
            "0" if not autogroup else "1",
            dry_run=dry_run,
        ))
        # Util clamp: set min > 0 to keep CPUs at higher frequency for
        # bursty workloads; max < 1024 to cap power on thermal-limited nodes
        results.append(self._write_sysfs(
            "/proc/sys/kernel/sched_util_clamp_min",
            str(util_clamp_min),
            dry_run=dry_run,
        ))
        results.append(self._write_sysfs(
            "/proc/sys/kernel/sched_util_clamp_max",
            str(util_clamp_max),
            dry_run=dry_run,
        ))
        results.append(self._write_sysfs(
            "/proc/sys/kernel/sched_rr_timeslice_ms",
            str(rr_timeslice_ms),
            dry_run=dry_run,
        ))
        return "; ".join(results)

    def apply_io_scheduler(self, device: str, scheduler: str = "mq-deadline", dry_run: bool = False) -> str:
        """Set IO scheduler for a block device (e.g. 'sda', 'nvme0n1')."""
        path = f"/sys/block/{device}/queue/scheduler"
        return self._write_sysfs(path, scheduler, dry_run=dry_run)

    def apply_irq_affinity(self, nic: str, numa_node: int, dry_run: bool = False) -> str:
        """Pin NIC IRQs to CPUs on the given NUMA node."""
        # Find CPUs on this NUMA node
        cpulist_path = f"/sys/devices/system/node/node{numa_node}/cpulist"
        cpulist = self._read_sysfs(cpulist_path)
        if cpulist is None:
            return f"SKIP: cannot read NUMA node {numa_node} CPU list"

        # Find IRQs for this NIC
        irq_results = []
        try:
            irq_dir = Path("/proc/irq")
            for irq_path in sorted(irq_dir.iterdir()):
                if not irq_path.is_dir():
                    continue
                # Check if any action file references this NIC
                actions_path = irq_path / "actions"
                if not actions_path.exists():
                    # Some kernels don't have 'actions', check smp_affinity_list exists
                    pass
                # Try the numeric IRQ directories that have the NIC name in their subdirs
                for child in irq_path.iterdir():
                    if child.is_dir() and nic in child.name:
                        affinity_path = str(irq_path / "smp_affinity_list")
                        irq_results.append(
                            self._write_sysfs(affinity_path, cpulist, dry_run=dry_run)
                        )
        except Exception as exc:
            return f"FAIL: IRQ affinity for {nic}: {exc}"

        if not irq_results:
            return f"SKIP: no IRQs found for NIC {nic}"
        failures = [r for r in irq_results if r.startswith("FAIL")]
        return f"OK: pinned {len(irq_results) - len(failures)}/{len(irq_results)} IRQs for {nic} to NUMA node {numa_node} (cpus {cpulist})"

    def apply_network_tuning(
        self,
        rmem_max: int = 16777216,
        wmem_max: int = 16777216,
        tcp_rmem: str = "4096 131072 16777216",
        tcp_wmem: str = "4096 65536 16777216",
        somaxconn: int = 65535,
        dry_run: bool = False,
    ) -> str:
        """Apply network buffer and TCP tuning."""
        settings = {
            "net.core.rmem_max": str(rmem_max),
            "net.core.wmem_max": str(wmem_max),
            "net.ipv4.tcp_rmem": tcp_rmem,
            "net.ipv4.tcp_wmem": tcp_wmem,
            "net.core.somaxconn": str(somaxconn),
        }
        results = []
        for key, val in settings.items():
            results.append(self._sysctl_write(key, val, dry_run=dry_run))
        return "; ".join(results)

    def apply_transparent_hugepages(self, mode: str = "madvise", dry_run: bool = False) -> str:
        """Set transparent hugepage mode: always, madvise, or never."""
        return self._write_sysfs("/sys/kernel/mm/transparent_hugepage/enabled", mode, dry_run=dry_run)

    # ------------------------------------------------------------------
    # Composite operations
    # ------------------------------------------------------------------

    def apply_all(self, dry_run: bool = False) -> Dict[str, str]:
        """Apply all tuning from fetched or default recommendations."""
        if not IS_LINUX:
            logger.info("System optimization skipped (not Linux: %s)", platform.system())
            return {"platform": f"SKIP: system tuning not supported on {platform.system()}"}
        rec = self.fetch_recommendations()
        results: Dict[str, str] = {}

        results["cpu_governor"] = self.apply_cpu_governor(
            rec.get("cpu_governor", "performance"), dry_run=dry_run
        )
        results["hugepages"] = self.apply_hugepages(
            rec.get("hugepages", 1024), dry_run=dry_run
        )
        results["numa_balancing"] = self.apply_numa_balancing(
            rec.get("numa_balancing", False), dry_run=dry_run
        )
        results["vm_tuning"] = self.apply_vm_tuning(
            swappiness=rec.get("swappiness", 10),
            dirty_ratio=rec.get("dirty_ratio", 20),
            dirty_background_ratio=rec.get("dirty_background_ratio", 5),
            vfs_cache_pressure=rec.get("vfs_cache_pressure", 50),
            dry_run=dry_run,
        )
        results["scheduler"] = self.apply_scheduler_tuning(
            autogroup=rec.get("sched_autogroup", False),
            util_clamp_min=rec.get("sched_util_clamp_min", 0),
            util_clamp_max=rec.get("sched_util_clamp_max", 1024),
            rr_timeslice_ms=rec.get("sched_rr_timeslice_ms", 100),
            dry_run=dry_run,
        )
        results["transparent_hugepages"] = self.apply_transparent_hugepages(
            rec.get("transparent_hugepages", "madvise"), dry_run=dry_run
        )
        results["network"] = self.apply_network_tuning(
            rmem_max=rec.get("rmem_max", 16777216),
            wmem_max=rec.get("wmem_max", 16777216),
            tcp_rmem=rec.get("tcp_rmem", "4096 131072 16777216"),
            tcp_wmem=rec.get("tcp_wmem", "4096 65536 16777216"),
            somaxconn=rec.get("somaxconn", 65535),
            dry_run=dry_run,
        )

        # IO schedulers per device
        for dev, sched in rec.get("io_schedulers", {}).items():
            results[f"io_sched_{dev}"] = self.apply_io_scheduler(dev, sched, dry_run=dry_run)

        # Process-level optimization
        proc_results = self.optimize_processes(dry_run=dry_run)
        results.update(proc_results)

        logger.info("apply_all complete: %d settings (%d process-level)",
                     len(results), len(proc_results))
        return results

    # ------------------------------------------------------------------
    # Process-level optimization (the missing piece)
    # ------------------------------------------------------------------

    def optimize_processes(self, dry_run: bool = False) -> Dict[str, str]:
        """Scan running processes and apply per-process tuning.

        This is the key difference from system-wide tuning: it looks at
        what's actually running, classifies each heavy process, and applies
        NUMA binding, CPU affinity, IO class, and nice corrections.

        Runs periodically (every 30s) from the agent optimization loop.
        """
        if not IS_LINUX:
            return {}
        import psutil

        results: Dict[str, str] = {}
        numa_count = self._get_numa_count()
        if numa_count <= 1:
            # Single NUMA node — only fix IO class/priority issues
            pass

        numa_cpus = {}
        for node in range(numa_count):
            cpus = self._read_sysfs(f"/sys/devices/system/node/node{node}/cpulist")
            if cpus:
                numa_cpus[node] = cpus

        # Track which NUMA nodes are being used for load balancing
        numa_load: Dict[int, int] = {n: 0 for n in range(numa_count)}

        for proc in psutil.process_iter(["pid", "name", "cmdline", "cpu_percent",
                                          "memory_info", "nice", "ionice",
                                          "num_threads", "username"]):
            try:
                info = proc.info
                pid = info["pid"]
                name = info["name"] or ""
                cmdline = info.get("cmdline") or []
                rss_mb = (info.get("memory_info") or psutil._common.pmem(0, 0, 0, 0, 0, 0, 0)).rss / (1024 * 1024)
                nice = info.get("nice", 0) or 0
                num_threads = info.get("num_threads", 1) or 1

                # Skip small processes, kernel threads, and ourselves
                if rss_mb < 100 or pid <= 2:
                    continue
                if "mesh" in name and "optimizer" in " ".join(cmdline):
                    continue

                cmd_str = " ".join(cmdline).lower()
                classification = self._classify_process(name, cmd_str, rss_mb, num_threads)

                if classification is None:
                    continue

                actions = []

                # --- NUMA binding ---
                if numa_count > 1 and classification.get("numa_bind"):
                    current_numa = self._get_process_numa(pid)
                    target_numa = classification.get("preferred_numa")

                    if target_numa is None:
                        # Pick least-loaded NUMA node
                        target_numa = min(numa_load, key=numa_load.get)

                    if current_numa != target_numa or not self._is_numa_bound(pid):
                        if target_numa in numa_cpus:
                            r = self._bind_process_numa(pid, target_numa, numa_cpus[target_numa], dry_run)
                            actions.append(r)

                    numa_load[target_numa] = numa_load.get(target_numa, 0) + num_threads

                # --- IO class fix ---
                target_io = classification.get("io_class", "best-effort")
                try:
                    current_io = proc.ionice()
                    # ioclass: 0=none, 1=realtime, 2=best-effort, 3=idle
                    io_map = {"realtime": 1, "best-effort": 2, "idle": 3}
                    target_io_num = io_map.get(target_io, 2)
                    if current_io.ioclass == 3 and target_io_num != 3:
                        # Process is at idle IO but shouldn't be
                        r = self._set_io_class(pid, target_io_num, dry_run)
                        actions.append(r)
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass

                # --- Nice fix ---
                target_nice = classification.get("nice", 0)
                if nice > target_nice and classification.get("fix_nice", False):
                    r = self._set_nice(pid, target_nice, dry_run)
                    actions.append(r)

                if actions:
                    results[f"pid_{pid}_{name}"] = "; ".join(actions)

            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                continue
            except Exception as e:
                logger.debug("Process optimization error for pid: %s", e)
                continue

        if results:
            logger.info("Process optimization: %d processes tuned", len(results))
        return results

    def _classify_process(self, name: str, cmd_str: str, rss_mb: float,
                           num_threads: int) -> Optional[Dict]:
        """Classify a process and return optimization recommendations.

        Returns None if the process doesn't need optimization, or a dict with:
          numa_bind: bool — should we NUMA-bind this process?
          preferred_numa: Optional[int] — specific NUMA node preference
          io_class: str — target IO class
          nice: int — target nice value
          fix_nice: bool — should we fix nice if too high?
        """
        # Vivado synthesis/implementation — heavy compute + memory + IO
        if "vivado" in name.lower() or "vivado" in cmd_str:
            return {
                "numa_bind": True,
                "preferred_numa": None,  # Auto-balance across NUMA nodes
                "io_class": "best-effort",
                "nice": 0,  # Vivado shouldn't be niced down
                "fix_nice": True,
                "reason": "vivado_synthesis",
            }

        # Python ML/training workloads
        if ("python" in name.lower() and any(kw in cmd_str for kw in
                ["train", "torch", "tensorflow", "jax", "cuda", "hip",
                 "rocm", "benchmark", "fuzzer", "fuzz"])):
            return {
                "numa_bind": True,
                "preferred_numa": None,
                "io_class": "best-effort",
                "nice": 0,
                "fix_nice": True,
                "reason": "ml_training",
            }

        # Compilers (gcc, g++, clang, rustc, javac)
        if any(comp in name.lower() for comp in ["gcc", "g++", "cc1", "clang", "rustc", "javac"]):
            return {
                "numa_bind": True,
                "preferred_numa": None,
                "io_class": "best-effort",
                "nice": 5,  # Slightly lower than interactive
                "fix_nice": False,
                "reason": "compiler",
            }

        # Database engines
        if any(db in name.lower() for db in ["postgres", "mysql", "mongod", "redis"]):
            return {
                "numa_bind": True,
                "preferred_numa": 0,  # Pin to node 0 for consistency
                "io_class": "best-effort",
                "nice": 0,
                "fix_nice": True,
                "reason": "database",
            }

        # Java processes (including Xilinx tools)
        if "java" in name.lower() and rss_mb > 500:
            return {
                "numa_bind": True,
                "preferred_numa": None,
                "io_class": "best-effort",
                "nice": 0,
                "fix_nice": True,
                "reason": "java_heavy",
            }

        # Any process using >2GB RSS and >4 threads — likely compute-heavy
        if rss_mb > 2048 and num_threads > 4:
            return {
                "numa_bind": True,
                "preferred_numa": None,
                "io_class": "best-effort",
                "nice": 0,
                "fix_nice": False,
                "reason": "heavy_process",
            }

        return None

    def _get_numa_count(self) -> int:
        """Get number of NUMA nodes."""
        try:
            nodes = [d for d in os.listdir("/sys/devices/system/node/")
                     if d.startswith("node") and d[4:].isdigit()]
            return len(nodes)
        except Exception:
            return 1

    def _get_process_numa(self, pid: int) -> Optional[int]:
        """Determine which NUMA node a process's memory is primarily on."""
        try:
            numa_maps = Path(f"/proc/{pid}/numa_maps").read_text()
            node_pages: Dict[int, int] = {}
            for line in numa_maps.split("\n"):
                for m in re.finditer(r"N(\d+)=(\d+)", line):
                    node = int(m.group(1))
                    pages = int(m.group(2))
                    node_pages[node] = node_pages.get(node, 0) + pages
            if node_pages:
                return max(node_pages, key=node_pages.get)
        except Exception:
            pass
        return None

    def _is_numa_bound(self, pid: int) -> bool:
        """Check if a process has CPU affinity restricted to one NUMA node."""
        try:
            cpus_allowed = self._read_sysfs(f"/proc/{pid}/status")
            if cpus_allowed:
                for line in cpus_allowed.split("\n"):
                    if line.startswith("Cpus_allowed_list:"):
                        cpulist = line.split(":", 1)[1].strip()
                        # If it contains all CPUs, it's not bound
                        all_cpus = os.cpu_count() or 1
                        bound_count = self._count_cpus(cpulist)
                        return bound_count < all_cpus
        except Exception:
            pass
        return False

    def _count_cpus(self, cpulist: str) -> int:
        """Count CPUs in a list like '0-7,16-23'."""
        count = 0
        for part in cpulist.split(","):
            part = part.strip()
            if "-" in part:
                lo, hi = part.split("-", 1)
                count += int(hi) - int(lo) + 1
            else:
                count += 1
        return count

    def _bind_process_numa(self, pid: int, numa_node: int, cpulist: str,
                            dry_run: bool = False) -> str:
        """Bind a process to a specific NUMA node using taskset + migrate_pages."""
        if dry_run:
            return f"DRY_RUN: would bind pid {pid} to NUMA node {numa_node} (cpus {cpulist})"

        results = []
        # Set CPU affinity
        r = self._run_cmd(["taskset", "-apc", cpulist, str(pid)], dry_run=dry_run)
        results.append(r)

        # Migrate memory pages to target NUMA node (if migratepages available)
        if shutil.which("migratepages"):
            all_nodes = ",".join(str(n) for n in range(self._get_numa_count()))
            r = self._run_cmd(["migratepages", str(pid), all_nodes, str(numa_node)], dry_run=dry_run)
            results.append(r)

        return f"NUMA bind pid {pid} -> node {numa_node}: " + "; ".join(results)

    def _set_io_class(self, pid: int, io_class: int, dry_run: bool = False) -> str:
        """Set IO scheduling class for a process."""
        class_names = {1: "realtime", 2: "best-effort", 3: "idle"}
        if dry_run:
            return f"DRY_RUN: would set pid {pid} IO class to {class_names.get(io_class, io_class)}"
        return self._run_cmd(["ionice", "-c", str(io_class), "-p", str(pid)], dry_run=dry_run)

    def _set_nice(self, pid: int, nice: int, dry_run: bool = False) -> str:
        """Set nice value for a process."""
        if dry_run:
            return f"DRY_RUN: would renice pid {pid} to {nice}"
        return self._run_cmd(["renice", "-n", str(nice), "-p", str(pid)], dry_run=dry_run)

    def get_current_state(self) -> Dict:
        """Read all current settings and return them as a dict."""
        state: Dict = {}

        # CPU governor (just check cpu0)
        state["cpu_governor"] = self._read_sysfs(
            "/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor"
        )

        # Hugepages
        state["hugepages_2M"] = self._read_sysfs(
            "/sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages"
        )
        state["hugepages_1G"] = self._read_sysfs(
            "/sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages"
        )

        # NUMA balancing
        state["numa_balancing"] = self._read_sysfs("/proc/sys/kernel/numa_balancing")

        # VM
        state["swappiness"] = self._read_sysfs("/proc/sys/vm/swappiness")
        state["dirty_ratio"] = self._read_sysfs("/proc/sys/vm/dirty_ratio")
        state["dirty_background_ratio"] = self._read_sysfs("/proc/sys/vm/dirty_background_ratio")
        state["vfs_cache_pressure"] = self._read_sysfs("/proc/sys/vm/vfs_cache_pressure")

        # Scheduler (EEVDF on 6.6+)
        state["sched_autogroup_enabled"] = self._read_sysfs(
            "/proc/sys/kernel/sched_autogroup_enabled"
        )
        state["sched_util_clamp_min"] = self._read_sysfs(
            "/proc/sys/kernel/sched_util_clamp_min"
        )
        state["sched_util_clamp_max"] = self._read_sysfs(
            "/proc/sys/kernel/sched_util_clamp_max"
        )
        state["sched_rr_timeslice_ms"] = self._read_sysfs(
            "/proc/sys/kernel/sched_rr_timeslice_ms"
        )

        # THP
        thp_raw = self._read_sysfs("/sys/kernel/mm/transparent_hugepage/enabled")
        if thp_raw:
            # Format is "always [madvise] never" — extract the bracketed one
            m = re.search(r"\[(\w+)\]", thp_raw)
            state["transparent_hugepages"] = m.group(1) if m else thp_raw
        else:
            state["transparent_hugepages"] = None

        # Network
        state["rmem_max"] = self._read_sysfs("/proc/sys/net/core/rmem_max")
        state["wmem_max"] = self._read_sysfs("/proc/sys/net/core/wmem_max")
        state["tcp_rmem"] = self._read_sysfs("/proc/sys/net/ipv4/tcp_rmem")
        state["tcp_wmem"] = self._read_sysfs("/proc/sys/net/ipv4/tcp_wmem")
        state["somaxconn"] = self._read_sysfs("/proc/sys/net/core/somaxconn")

        return state


# ------------------------------------------------------------------
# Process Monitor — event-driven process optimization
# ------------------------------------------------------------------

class ProcessMonitor:
    """Watches for new processes and optimizes them immediately.

    Uses Linux proc connector (netlink CN_PROC) for instant fork/exec
    notifications. Falls back to polling /proc every 5s if netlink
    is unavailable (needs root for proc connector).
    """

    # Netlink / proc connector constants
    NETLINK_CONNECTOR = 11
    CN_IDX_PROC = 1
    CN_VAL_PROC = 1
    PROC_CN_MCAST_LISTEN = 1
    PROC_EVENT_EXEC = 0x00000002
    PROC_EVENT_FORK = 0x00000001

    def __init__(self, optimizer: SystemOptimizer):
        self.optimizer = optimizer
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._known_pids: Set[int] = set()
        self._optimized_pids: Set[int] = set()
        # Debounce: don't re-optimize a PID within 60s
        self._optimize_times: Dict[int, float] = {}
        self._debounce_s = 60.0

    def start(self):
        """Start the process monitor in a background thread."""
        if not IS_LINUX:
            logger.info("Process monitor skipped (not Linux)")
            return
        self._running = True
        # Try proc connector first (instant), fall back to polling
        if self._can_use_proc_connector():
            self._thread = threading.Thread(
                target=self._proc_connector_loop,
                name="proc-monitor-netlink",
                daemon=True,
            )
            logger.info("Process monitor starting (netlink proc connector)")
        else:
            self._thread = threading.Thread(
                target=self._polling_loop,
                name="proc-monitor-poll",
                daemon=True,
            )
            logger.info("Process monitor starting (polling mode, 5s interval)")
        self._thread.start()

    def stop(self):
        self._running = False
        if self._thread:
            self._thread.join(timeout=5)

    def _can_use_proc_connector(self) -> bool:
        """Check if we can use netlink proc connector (needs root)."""
        if os.geteuid() != 0:
            return False
        try:
            sock = socket.socket(socket.AF_NETLINK, socket.SOCK_DGRAM,
                                 self.NETLINK_CONNECTOR)
            sock.close()
            return True
        except Exception:
            return False

    def _proc_connector_loop(self):
        """Event-driven: get notified on every exec() via netlink."""
        try:
            sock = socket.socket(socket.AF_NETLINK, socket.SOCK_DGRAM,
                                 self.NETLINK_CONNECTOR)
            sock.bind((os.getpid(), self.CN_IDX_PROC))

            # Subscribe to proc events
            # nlmsghdr (16 bytes) + cn_msg header (20 bytes) + enum (4 bytes)
            msg = struct.pack("=IHHII", 40, 0x1B, 0, 0, os.getpid())  # nlmsghdr-ish
            # Simpler: use the proc_connector subscription message
            # This is tricky to get right — fall back to /proc polling with
            # inotify on /proc for faster detection
            sock.close()
            logger.info("Proc connector subscribe failed, falling back to polling")
            self._polling_loop()
        except Exception as e:
            logger.warning("Proc connector failed: %s, falling back to polling", e)
            self._polling_loop()

    def _polling_loop(self):
        """Poll /proc every 5s for new processes."""
        while self._running:
            try:
                current_pids = set()
                for entry in os.listdir("/proc"):
                    if entry.isdigit():
                        current_pids.add(int(entry))

                # New processes since last scan
                new_pids = current_pids - self._known_pids
                self._known_pids = current_pids

                # Clean up old optimized PIDs that no longer exist
                self._optimized_pids &= current_pids
                dead_pids = [p for p in self._optimize_times if p not in current_pids]
                for p in dead_pids:
                    del self._optimize_times[p]

                # Check new processes
                now = time.time()
                for pid in new_pids:
                    if pid in self._optimized_pids:
                        continue
                    last_opt = self._optimize_times.get(pid, 0)
                    if now - last_opt < self._debounce_s:
                        continue

                    try:
                        self._check_and_optimize(pid)
                    except Exception:
                        pass

            except Exception as e:
                logger.debug("Process monitor poll error: %s", e)

            time.sleep(5)

    def _check_and_optimize(self, pid: int):
        """Check if a new process needs optimization and apply it."""
        try:
            # Read process info from /proc
            comm = Path(f"/proc/{pid}/comm").read_text().strip()
            try:
                cmdline = Path(f"/proc/{pid}/cmdline").read_text().replace("\0", " ").strip()
            except Exception:
                cmdline = comm

            # Quick RSS check
            try:
                status = Path(f"/proc/{pid}/status").read_text()
                rss_kb = 0
                threads = 1
                for line in status.split("\n"):
                    if line.startswith("VmRSS:"):
                        rss_kb = int(line.split()[1])
                    elif line.startswith("Threads:"):
                        threads = int(line.split()[1])
                rss_mb = rss_kb / 1024
            except Exception:
                return

            # Skip small processes
            if rss_mb < 50:  # Lower threshold for event-driven (catch earlier)
                return

            classification = self.optimizer._classify_process(
                comm, cmdline.lower(), rss_mb, threads
            )
            if classification is None:
                return

            # This process needs optimization
            logger.info("New process detected: pid=%d name=%s rss=%.0fMB reason=%s",
                        pid, comm, rss_mb, classification.get("reason", "unknown"))

            numa_count = self.optimizer._get_numa_count()
            numa_cpus = {}
            for node in range(numa_count):
                cpus = self.optimizer._read_sysfs(
                    f"/sys/devices/system/node/node{node}/cpulist"
                )
                if cpus:
                    numa_cpus[node] = cpus

            # NUMA binding
            if numa_count > 1 and classification.get("numa_bind"):
                target_numa = classification.get("preferred_numa")
                if target_numa is None:
                    # Distribute: use PID mod numa_count for simple balancing
                    target_numa = pid % numa_count
                if target_numa in numa_cpus:
                    self.optimizer._bind_process_numa(
                        pid, target_numa, numa_cpus[target_numa]
                    )

            # IO class
            target_io = classification.get("io_class", "best-effort")
            io_map = {"realtime": 1, "best-effort": 2, "idle": 3}
            self.optimizer._set_io_class(pid, io_map.get(target_io, 2))

            # Nice
            if classification.get("fix_nice", False):
                target_nice = classification.get("nice", 0)
                self.optimizer._set_nice(pid, target_nice)

            self._optimized_pids.add(pid)
            self._optimize_times[pid] = time.time()

        except FileNotFoundError:
            pass  # Process already exited
        except PermissionError:
            pass


# ------------------------------------------------------------------
# CLI entry point
# ------------------------------------------------------------------

def main():
    import argparse

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    parser = argparse.ArgumentParser(description="System Optimizer — apply OS tuning from atlas/JEPA")
    parser.add_argument("--node", default=os.uname().nodename, help="Node name")
    parser.add_argument("--controller", default="http://localhost:8401", help="Controller URL")
    parser.add_argument("--dry-run", action="store_true", help="Show what would change without applying")
    parser.add_argument("--state", action="store_true", help="Print current system state and exit")
    parser.add_argument("--apply", action="store_true", help="Apply all recommendations")
    args = parser.parse_args()

    opt = SystemOptimizer(args.node, args.controller)

    if args.state:
        state = opt.get_current_state()
        print(json.dumps(state, indent=2))
        return

    if args.apply or args.dry_run:
        results = opt.apply_all(dry_run=args.dry_run)
        for k, v in results.items():
            print(f"  {k}: {v}")
    else:
        parser.print_help()


if __name__ == "__main__":
    main()
