"""
meshrun — process wrapper that applies per-process optimizations before exec.

Usage:
    python -m mesh.scripts.meshrun [options] -- <command> [args...]
    python -m mesh.scripts.meshrun --auto -- python train.py
    python -m mesh.scripts.meshrun --cpu-pin 0-7 --numa-node 0 --priority -5 -- ./my_benchmark

Wraps a command with numactl, taskset, nice, ionice, and environment variables
for optimal execution based on mesh/atlas recommendations.
"""

import argparse
import json
import logging
import os
import shutil
import subprocess
import sys
from typing import Dict, List, Optional, Tuple

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

logger = logging.getLogger(__name__)

MESH_AGENT_URL = "http://localhost:8400"
MESH_CONTROLLER_URL = "http://localhost:8401"

GPU_KEYWORDS = {"python", "pytorch", "cuda", "hip", "rocm", "torchrun", "deepspeed", "accelerate"}


def check_tool(name: str) -> bool:
    """Check if a system tool is available."""
    return shutil.which(name) is not None


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


def get_numa_cpulist(node: int) -> Optional[str]:
    """Get the CPU list for a NUMA node."""
    path = f"/sys/devices/system/node/node{node}/cpulist"
    try:
        with open(path) as f:
            return f.read().strip()
    except FileNotFoundError:
        return None


def get_gpu_numa_node() -> Optional[int]:
    """Detect the NUMA node of the first AMD GPU (renderD128+) or NVIDIA GPU."""
    # Check AMD GPUs first
    for i in range(8):
        path = f"/sys/class/drm/card{i}/device/numa_node"
        try:
            with open(path) as f:
                node = int(f.read().strip())
                if node >= 0:
                    return node
        except (FileNotFoundError, ValueError):
            continue
    # Check NVIDIA GPUs
    for i in range(8):
        path = f"/sys/bus/pci/drivers/nvidia/{i:04d}:*/numa_node"
        import glob
        matches = glob.glob(path)
        for m in matches:
            try:
                with open(m) as f:
                    node = int(f.read().strip())
                    if node >= 0:
                        return node
            except (FileNotFoundError, ValueError):
                continue
    return None


def is_gpu_workload(command: List[str]) -> bool:
    """Heuristic: does this command look like a GPU workload?"""
    cmd_str = " ".join(command).lower()
    return any(kw in cmd_str for kw in GPU_KEYWORDS)


def fetch_auto_config() -> Dict:
    """Contact local mesh agent and controller for auto-configuration."""
    config: Dict = {}
    if requests is None:
        logger.warning("requests not available; --auto will use local heuristics only")
        return config

    # Try local agent for hardware info
    try:
        resp = requests.get(f"{MESH_AGENT_URL}/hardware", timeout=2)
        if resp.ok:
            config["hardware"] = resp.json()
    except Exception:
        logger.debug("Could not reach local mesh agent at %s", MESH_AGENT_URL)

    # Try controller for JEPA recommendations
    try:
        node_name = os.uname().nodename
        resp = requests.get(
            f"{MESH_CONTROLLER_URL}/nodes/{node_name}/recommendations", timeout=2
        )
        if resp.ok:
            config["recommendations"] = resp.json()
    except Exception:
        logger.debug("Could not reach controller at %s", MESH_CONTROLLER_URL)

    return config


def build_auto_settings(command: List[str]) -> Dict:
    """
    Auto-detect best settings for the given command.

    Returns a dict with keys: cpu_pin, numa_node, hugepages, priority, io_class,
    omp_threads, malloc_arenas.
    """
    settings: Dict = {}
    remote_config = fetch_auto_config()
    rec = remote_config.get("recommendations", {})
    hw = remote_config.get("hardware", {})

    numa_count = get_numa_nodes()
    gpu_workload = is_gpu_workload(command)

    # NUMA pinning
    if gpu_workload and numa_count > 1:
        gpu_node = get_gpu_numa_node()
        if gpu_node is not None:
            settings["numa_node"] = gpu_node
            logger.info("Auto: GPU workload detected, pinning to NUMA node %d", gpu_node)
    elif numa_count > 1:
        # Default to node 0 for non-GPU workloads
        settings["numa_node"] = rec.get("preferred_numa_node", 0)

    # CPU pinning — if we have a NUMA node, pin to those CPUs
    if "numa_node" in settings:
        cpulist = get_numa_cpulist(settings["numa_node"])
        if cpulist:
            settings["cpu_pin"] = cpulist

    # Hugepages — use madvise-based via environment
    settings["hugepages"] = True

    # Priority — slightly elevated for GPU workloads
    if gpu_workload:
        settings["priority"] = -5
    else:
        settings["priority"] = 0

    # IO class — best-effort for most things
    settings["io_class"] = "best-effort"

    # Thread counts
    if "numa_node" in settings:
        cpulist = get_numa_cpulist(settings["numa_node"])
        if cpulist:
            # Count CPUs in the list (e.g. "0-7" = 8, "0,2,4,6" = 4)
            count = _count_cpus(cpulist)
            settings["omp_threads"] = count
    else:
        # Use all CPUs but leave 2 for system
        total = os.cpu_count() or 4
        settings["omp_threads"] = max(1, total - 2)

    settings["malloc_arenas"] = min(8, settings.get("omp_threads", 4))

    return settings


def _count_cpus(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 _ionice_class_flag(io_class: str) -> List[str]:
    """Convert IO class name to ionice flags."""
    mapping = {
        "realtime": ["-c", "1"],
        "best-effort": ["-c", "2"],
        "idle": ["-c", "3"],
    }
    return mapping.get(io_class.lower(), ["-c", "2"])


def build_command(
    command: List[str],
    cpu_pin: Optional[str] = None,
    numa_node: Optional[int] = None,
    hugepages: bool = False,
    priority: Optional[int] = None,
    io_class: Optional[str] = None,
    omp_threads: Optional[int] = None,
    malloc_arenas: Optional[int] = None,
) -> Tuple[List[str], Dict[str, str]]:
    """
    Build the final command line and environment variables.

    Returns (command_parts, env_updates).
    """
    parts: List[str] = []
    env: Dict[str, str] = {}
    warnings: List[str] = []

    # numactl (must come first as it wraps the whole command)
    if numa_node is not None:
        if check_tool("numactl"):
            parts.extend(["numactl", f"--cpunodebind={numa_node}", f"--membind={numa_node}"])
        else:
            warnings.append("numactl not found; skipping NUMA binding")

    # taskset for CPU pinning (only if not already using numactl, which handles CPU binding)
    if cpu_pin is not None and numa_node is None:
        if check_tool("taskset"):
            parts.extend(["taskset", "-c", cpu_pin])
        else:
            warnings.append("taskset not found; skipping CPU pinning")

    # ionice
    if io_class is not None:
        if check_tool("ionice"):
            parts.extend(["ionice"] + _ionice_class_flag(io_class))
        else:
            warnings.append("ionice not found; skipping IO class")

    # nice
    if priority is not None and priority != 0:
        if check_tool("nice"):
            parts.extend(["nice", "-n", str(priority)])
        else:
            warnings.append("nice not found; skipping priority")

    # Environment variables
    if omp_threads is not None:
        env["OMP_NUM_THREADS"] = str(omp_threads)
        env["MKL_NUM_THREADS"] = str(omp_threads)

    if malloc_arenas is not None:
        env["MALLOC_ARENA_MAX"] = str(malloc_arenas)

    if hugepages:
        # Try libhugetlbfs LD_PRELOAD, else set MADV hints via env
        libhugetlbfs = None
        for candidate in [
            "/usr/lib/x86_64-linux-gnu/libhugetlbfs.so",
            "/usr/lib64/libhugetlbfs.so",
            "/usr/lib/libhugetlbfs.so",
        ]:
            if os.path.exists(candidate):
                libhugetlbfs = candidate
                break

        if libhugetlbfs:
            existing = os.environ.get("LD_PRELOAD", "")
            if existing:
                env["LD_PRELOAD"] = f"{existing}:{libhugetlbfs}"
            else:
                env["LD_PRELOAD"] = libhugetlbfs
            env["HUGETLB_MORECORE"] = "yes"
        else:
            # No libhugetlbfs — set glibc malloc to use hugepages via madvise
            env["GLIBC_TUNABLES"] = "glibc.malloc.hugetlb=1"

    # Append the actual command
    parts.extend(command)

    for w in warnings:
        logger.warning(w)

    return parts, env


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

    parser = argparse.ArgumentParser(
        description="meshrun — per-process optimizer wrapper",
        usage="python -m mesh.scripts.meshrun [options] -- <command> [args...]",
    )
    parser.add_argument("--cpu-pin", type=str, default=None, help="CPU core range (e.g. '0-7')")
    parser.add_argument("--numa-node", type=int, default=None, help="NUMA node to bind to")
    parser.add_argument("--hugepages", action="store_true", help="Enable hugepage support")
    parser.add_argument("--priority", type=int, default=None, help="Nice priority value")
    parser.add_argument(
        "--io-class",
        type=str,
        default=None,
        choices=["best-effort", "realtime", "idle"],
        help="IO scheduling class",
    )
    parser.add_argument(
        "--auto",
        action="store_true",
        help="Auto-detect optimal settings from mesh agent / local hardware",
    )
    parser.add_argument("--dry-run", action="store_true", help="Print command without executing")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    parser.add_argument("command", nargs=argparse.REMAINDER, help="Command to run (after --)")

    args = parser.parse_args()

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

    # Extract the command (strip leading '--' if present)
    command = args.command
    if command and command[0] == "--":
        command = command[1:]
    if not command:
        parser.error("No command specified. Usage: meshrun [options] -- <command>")

    # Determine settings
    if args.auto:
        auto = build_auto_settings(command)
        # CLI args override auto settings
        cpu_pin = args.cpu_pin or auto.get("cpu_pin")
        numa_node = args.numa_node if args.numa_node is not None else auto.get("numa_node")
        hugepages = args.hugepages or auto.get("hugepages", False)
        priority = args.priority if args.priority is not None else auto.get("priority")
        io_class = args.io_class or auto.get("io_class")
        omp_threads = auto.get("omp_threads")
        malloc_arenas = auto.get("malloc_arenas")
    else:
        cpu_pin = args.cpu_pin
        numa_node = args.numa_node
        hugepages = args.hugepages
        priority = args.priority
        io_class = args.io_class
        omp_threads = None
        malloc_arenas = None

    # Build the command
    final_cmd, env_updates = build_command(
        command,
        cpu_pin=cpu_pin,
        numa_node=numa_node,
        hugepages=hugepages,
        priority=priority,
        io_class=io_class,
        omp_threads=omp_threads,
        malloc_arenas=malloc_arenas,
    )

    if args.dry_run or args.verbose:
        env_str = " ".join(f"{k}={v}" for k, v in env_updates.items())
        cmd_str = " ".join(final_cmd)
        full = f"{env_str} {cmd_str}".strip() if env_str else cmd_str
        if args.dry_run:
            print(f"meshrun dry-run: {full}")
            return
        else:
            logger.info("Executing: %s", full)

    # Apply environment
    for k, v in env_updates.items():
        os.environ[k] = v

    # Execute — replace this process
    try:
        os.execvp(final_cmd[0], final_cmd)
    except FileNotFoundError:
        logger.error("Command not found: %s", final_cmd[0])
        sys.exit(127)
    except PermissionError:
        logger.error("Permission denied: %s", final_cmd[0])
        sys.exit(126)


if __name__ == "__main__":
    main()
