"""ARIS Kernel Probes — Deep hardware and OS awareness.

Provides ARIS with visibility into layers 0-4 of the computing stack:
  L0: Hardware (SMART, RAPL power, thermal sensors)
  L1: HAL (NVML GPU direct, ACPI, DMI/SMBIOS)
  L2: Kernel (eBPF syscall tracing, scheduler stats)
  L3: Subsystems (filesystem I/O, network connections, block devices)
  L4: Libraries (memory allocator stats, shared library usage)

This is what makes ARIS different from every other AI tool:
they operate at L5-8, we see L0-4.
"""

import asyncio
import json
import os
import re
from pathlib import Path
from typing import Any, Optional


# ============================================================
# LAYER 0: HARDWARE DIRECT
# ============================================================

class HardwareProbe:
    """Direct hardware sensor access."""

    @staticmethod
    async def cpu_info() -> dict[str, Any]:
        """Read CPU details from /proc/cpuinfo."""
        try:
            text = Path("/proc/cpuinfo").read_text()
            model = ""
            cores = 0
            mhz = 0.0
            for line in text.split("\n"):
                if "model name" in line and not model:
                    model = line.split(":")[1].strip()
                if "processor" in line:
                    cores += 1
                if "cpu MHz" in line and mhz == 0:
                    mhz = float(line.split(":")[1].strip())
            return {"model": model, "cores": cores, "mhz_current": mhz}
        except Exception as e:
            return {"error": str(e)}

    @staticmethod
    async def memory_detailed() -> dict[str, Any]:
        """Read detailed memory from /proc/meminfo — not just free -m."""
        try:
            text = Path("/proc/meminfo").read_text()
            info = {}
            for line in text.split("\n"):
                if ":" in line:
                    key, val = line.split(":", 1)
                    # Convert kB to MB
                    num = re.search(r"(\d+)", val)
                    if num:
                        info[key.strip()] = int(num.group(1)) // 1024  # MB
            return {
                "total_mb": info.get("MemTotal", 0),
                "free_mb": info.get("MemFree", 0),
                "available_mb": info.get("MemAvailable", 0),
                "buffers_mb": info.get("Buffers", 0),
                "cached_mb": info.get("Cached", 0),
                "swap_total_mb": info.get("SwapTotal", 0),
                "swap_free_mb": info.get("SwapFree", 0),
                "dirty_mb": info.get("Dirty", 0),
                "hugepages_total": info.get("HugePages_Total", 0),
                "slab_mb": info.get("Slab", 0),
            }
        except Exception as e:
            return {"error": str(e)}

    @staticmethod
    async def smart_health(device: str = "/dev/sda") -> dict[str, Any]:
        """Read SMART disk health — predict failures before they happen."""
        try:
            proc = await asyncio.create_subprocess_exec(
                "smartctl", "-A", "-H", device, "--json",
                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
            )
            stdout, _ = await proc.communicate()
            data = json.loads(stdout.decode())

            health = data.get("smart_status", {}).get("passed", None)
            attrs = {}
            for attr in data.get("ata_smart_attributes", {}).get("table", []):
                name = attr.get("name", "")
                if name in ("Reallocated_Sector_Ct", "Current_Pending_Sector",
                            "Power_On_Hours", "Temperature_Celsius",
                            "Wear_Leveling_Count", "Total_LBAs_Written"):
                    attrs[name] = attr.get("raw", {}).get("value", 0)

            return {"device": device, "healthy": health, "attributes": attrs}
        except FileNotFoundError:
            return {"error": "smartctl not installed (apt install smartmontools)"}
        except Exception as e:
            return {"error": str(e)}

    @staticmethod
    async def thermal_zones() -> list[dict[str, Any]]:
        """Read all thermal sensors — not just thermal_zone0."""
        zones = []
        thermal_base = Path("/sys/class/thermal")
        if thermal_base.exists():
            for zone in sorted(thermal_base.iterdir()):
                if zone.name.startswith("thermal_zone"):
                    try:
                        temp = int((zone / "temp").read_text().strip()) / 1000
                        zone_type = (zone / "type").read_text().strip()
                        zones.append({"zone": zone.name, "type": zone_type, "temp_c": temp})
                    except Exception:
                        pass
        return zones

    @staticmethod
    async def power_rapl() -> dict[str, Any]:
        """Read CPU power consumption via RAPL (Running Average Power Limit)."""
        rapl_base = Path("/sys/class/powercap")
        if not rapl_base.exists():
            return {"error": "RAPL not available (needs root or powercap access)"}

        domains = {}
        for entry in rapl_base.iterdir():
            if entry.name.startswith("intel-rapl"):
                try:
                    name = (entry / "name").read_text().strip()
                    energy_uj = int((entry / "energy_uj").read_text().strip())
                    domains[name] = {"energy_microjoules": energy_uj}
                except Exception:
                    pass
        return domains if domains else {"error": "No RAPL domains found"}


# ============================================================
# LAYER 1: HAL — Hardware Abstraction
# ============================================================

class HALProbe:
    """Hardware Abstraction Layer — direct device access."""

    @staticmethod
    async def gpu_nvml() -> dict[str, Any]:
        """Read GPU via NVML (pynvml) — 10x faster than nvidia-smi."""
        try:
            import pynvml
            pynvml.nvmlInit()
            count = pynvml.nvmlDeviceGetCount()
            gpus = []
            for i in range(count):
                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                name = pynvml.nvmlDeviceGetName(handle)
                if isinstance(name, bytes):
                    name = name.decode()
                mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
                temp = pynvml.nvmlDeviceGetTemperature(handle, 0)
                power = pynvml.nvmlDeviceGetPowerUsage(handle) / 1000  # mW to W
                util = pynvml.nvmlDeviceGetUtilizationRates(handle)
                clocks = {
                    "graphics_mhz": pynvml.nvmlDeviceGetClockInfo(handle, 0),
                    "mem_mhz": pynvml.nvmlDeviceGetClockInfo(handle, 2),
                }
                gpus.append({
                    "index": i, "name": name,
                    "memory_used_mb": mem.used // (1024 * 1024),
                    "memory_total_mb": mem.total // (1024 * 1024),
                    "memory_free_mb": mem.free // (1024 * 1024),
                    "temperature_c": temp,
                    "power_watts": round(power, 1),
                    "gpu_utilization_pct": util.gpu,
                    "memory_utilization_pct": util.memory,
                    "clocks": clocks,
                })
            pynvml.nvmlShutdown()
            return {"gpu_count": count, "gpus": gpus}
        except ImportError:
            return {"error": "pynvml not installed (pip install pynvml)"}
        except Exception as e:
            return {"error": str(e)}

    @staticmethod
    async def dmi_info() -> dict[str, Any]:
        """Read DMI/SMBIOS — motherboard, BIOS, chassis info."""
        dmi = {}
        dmi_base = Path("/sys/class/dmi/id")
        if dmi_base.exists():
            for field in ["board_name", "board_vendor", "bios_version",
                          "bios_date", "chassis_type", "product_name"]:
                try:
                    dmi[field] = (dmi_base / field).read_text().strip()
                except Exception:
                    pass
        return dmi if dmi else {"error": "DMI not available"}

    @staticmethod
    async def pci_devices() -> list[dict[str, Any]]:
        """List PCI devices — find GPUs, NICs, NVMe controllers."""
        devices = []
        try:
            proc = await asyncio.create_subprocess_exec(
                "lspci", "-mm",
                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
            )
            stdout, _ = await proc.communicate()
            for line in stdout.decode().strip().split("\n"):
                if any(kw in line.lower() for kw in ["vga", "3d", "nvidia", "ethernet", "nvme", "usb"]):
                    devices.append({"raw": line.strip()})
        except Exception as e:
            devices.append({"error": str(e)})
        return devices


# ============================================================
# LAYER 2: KERNEL
# ============================================================

class KernelProbe:
    """Kernel-level probes — processes, scheduler, syscalls."""

    @staticmethod
    async def process_tree_top(n: int = 10) -> list[dict[str, Any]]:
        """Top N processes by CPU/RAM — from /proc directly, not ps."""
        procs = []
        proc_path = Path("/proc")
        page_size = os.sysconf("SC_PAGE_SIZE")
        clk_tck = os.sysconf("SC_CLK_TCK")

        for entry in proc_path.iterdir():
            if entry.name.isdigit():
                try:
                    stat = (entry / "stat").read_text().split()
                    comm = stat[1].strip("()")
                    state = stat[2]
                    utime = int(stat[13])
                    stime = int(stat[14])
                    rss_pages = int(stat[23])
                    rss_mb = (rss_pages * page_size) // (1024 * 1024)
                    cpu_ticks = utime + stime

                    procs.append({
                        "pid": int(entry.name),
                        "name": comm,
                        "state": state,
                        "cpu_ticks": cpu_ticks,
                        "rss_mb": rss_mb,
                    })
                except Exception:
                    pass

        # Sort by RSS (memory usage)
        procs.sort(key=lambda p: -p["rss_mb"])
        return procs[:n]

    @staticmethod
    async def scheduler_stats() -> dict[str, Any]:
        """Read scheduler statistics — how well is the CPU being used?"""
        try:
            stat = Path("/proc/stat").read_text()
            first_line = stat.split("\n")[0]
            parts = first_line.split()
            # cpu user nice system idle iowait irq softirq steal
            return {
                "user": int(parts[1]),
                "nice": int(parts[2]),
                "system": int(parts[3]),
                "idle": int(parts[4]),
                "iowait": int(parts[5]),
                "irq": int(parts[6]),
                "softirq": int(parts[7]),
                "steal": int(parts[8]) if len(parts) > 8 else 0,
            }
        except Exception as e:
            return {"error": str(e)}

    @staticmethod
    async def context_switches() -> dict[str, Any]:
        """Context switches and process forks — indicators of system stress."""
        try:
            stat = Path("/proc/stat").read_text()
            ctxt = 0
            processes = 0
            for line in stat.split("\n"):
                if line.startswith("ctxt "):
                    ctxt = int(line.split()[1])
                if line.startswith("processes "):
                    processes = int(line.split()[1])
            return {"context_switches": ctxt, "processes_created": processes}
        except Exception as e:
            return {"error": str(e)}

    @staticmethod
    async def interrupts_summary() -> dict[str, int]:
        """Interrupt counts — which devices are busiest?"""
        try:
            text = Path("/proc/interrupts").read_text()
            lines = text.strip().split("\n")
            summary = {}
            for line in lines[1:]:  # Skip header
                parts = line.split()
                if len(parts) > 1:
                    irq_name = parts[-1] if not parts[-1].isdigit() else parts[0].rstrip(":")
                    count = sum(int(p) for p in parts[1:-1] if p.isdigit())
                    if count > 10000:  # Only significant IRQs
                        summary[irq_name] = count
            return dict(sorted(summary.items(), key=lambda x: -x[1])[:10])
        except Exception as e:
            return {"error": str(e)}


# ============================================================
# LAYER 3: KERNEL SUBSYSTEMS
# ============================================================

class SubsystemProbe:
    """Kernel subsystem probes — filesystem, network, block I/O."""

    @staticmethod
    async def disk_io_stats() -> dict[str, Any]:
        """Read /proc/diskstats — actual I/O operations, not just df."""
        try:
            text = Path("/proc/diskstats").read_text()
            disks = {}
            for line in text.strip().split("\n"):
                parts = line.split()
                if len(parts) >= 14:
                    name = parts[2]
                    if name.startswith("sd") or name.startswith("nvme"):
                        if not any(c.isdigit() for c in name[3:]) or "nvme" in name:
                            disks[name] = {
                                "reads_completed": int(parts[3]),
                                "reads_merged": int(parts[4]),
                                "sectors_read": int(parts[5]),
                                "read_time_ms": int(parts[6]),
                                "writes_completed": int(parts[7]),
                                "writes_merged": int(parts[8]),
                                "sectors_written": int(parts[9]),
                                "write_time_ms": int(parts[10]),
                                "io_in_progress": int(parts[11]),
                            }
            return disks
        except Exception as e:
            return {"error": str(e)}

    @staticmethod
    async def network_connections() -> dict[str, Any]:
        """Read /proc/net/tcp — active connections without netstat."""
        try:
            connections = {"established": 0, "listen": 0, "time_wait": 0,
                          "close_wait": 0, "total": 0}
            state_map = {"01": "established", "0A": "listen",
                         "06": "time_wait", "08": "close_wait"}

            for proto in ["tcp", "tcp6"]:
                path = Path(f"/proc/net/{proto}")
                if path.exists():
                    for line in path.read_text().strip().split("\n")[1:]:
                        parts = line.split()
                        if len(parts) >= 4:
                            state = parts[3]
                            connections["total"] += 1
                            if state in state_map:
                                connections[state_map[state]] += 1

            return connections
        except Exception as e:
            return {"error": str(e)}

    @staticmethod
    async def filesystem_mounts() -> list[dict[str, Any]]:
        """Read mounted filesystems with usage — from /proc/mounts + statvfs."""
        mounts = []
        try:
            text = Path("/proc/mounts").read_text()
            seen = set()
            for line in text.strip().split("\n"):
                parts = line.split()
                if len(parts) >= 3:
                    device, mountpoint, fstype = parts[0], parts[1], parts[2]
                    if fstype in ("ext4", "xfs", "btrfs", "ntfs", "vfat", "nvme",
                                  "tmpfs", "overlay") and mountpoint not in seen:
                        seen.add(mountpoint)
                        try:
                            st = os.statvfs(mountpoint)
                            total = (st.f_blocks * st.f_frsize) // (1024 * 1024 * 1024)
                            free = (st.f_bfree * st.f_frsize) // (1024 * 1024 * 1024)
                            if total > 0:
                                mounts.append({
                                    "device": device, "mount": mountpoint,
                                    "fstype": fstype, "total_gb": total,
                                    "free_gb": free, "used_pct": round((1 - free / total) * 100, 1),
                                })
                        except Exception:
                            pass
        except Exception as e:
            mounts.append({"error": str(e)})
        return mounts


# ============================================================
# LAYER 4: SYSTEM LIBRARIES
# ============================================================

class LibraryProbe:
    """System library inspection."""

    @staticmethod
    async def memory_allocator() -> dict[str, Any]:
        """Read glibc malloc stats from /proc/self/status."""
        try:
            text = Path("/proc/self/status").read_text()
            info = {}
            for line in text.split("\n"):
                if any(line.startswith(k) for k in ["VmSize", "VmRSS", "VmPeak",
                                                      "VmSwap", "VmLib", "Threads"]):
                    key, val = line.split(":", 1)
                    num = re.search(r"(\d+)", val)
                    if num:
                        info[key.strip()] = int(num.group(1))
            return info
        except Exception as e:
            return {"error": str(e)}

    @staticmethod
    async def loaded_libraries(pid: Optional[int] = None) -> list[str]:
        """List shared libraries loaded by a process."""
        target_pid = pid or os.getpid()
        try:
            maps = Path(f"/proc/{target_pid}/maps").read_text()
            libs = set()
            for line in maps.split("\n"):
                if ".so" in line:
                    parts = line.split()
                    if len(parts) >= 6:
                        lib = parts[-1]
                        if ".so" in lib:
                            libs.add(lib.split("/")[-1].split(".so")[0] + ".so")
            return sorted(libs)[:20]
        except Exception as e:
            return [f"error: {e}"]


# ============================================================
# UNIFIED DEEP SCAN
# ============================================================

async def deep_scan(include_gpu: bool = True) -> dict[str, Any]:
    """Full 8-layer deep scan of the system.

    This is what makes ARIS different from every other AI tool.
    They see Layer 5-8. We see Layer 0-4.
    """
    results = {}

    # L0: Hardware
    hw = HardwareProbe()
    results["L0_hardware"] = {
        "cpu": await hw.cpu_info(),
        "memory": await hw.memory_detailed(),
        "thermal": await hw.thermal_zones(),
    }

    # L1: HAL
    hal = HALProbe()
    results["L1_hal"] = {
        "dmi": await hal.dmi_info(),
    }
    if include_gpu:
        results["L1_hal"]["gpu"] = await hal.gpu_nvml()

    # L2: Kernel
    kern = KernelProbe()
    results["L2_kernel"] = {
        "top_processes": await kern.process_tree_top(5),
        "scheduler": await kern.scheduler_stats(),
        "context_switches": await kern.context_switches(),
    }

    # L3: Subsystems
    sub = SubsystemProbe()
    results["L3_subsystems"] = {
        "disk_io": await sub.disk_io_stats(),
        "network": await sub.network_connections(),
        "filesystems": await sub.filesystem_mounts(),
    }

    # L4: Libraries
    lib = LibraryProbe()
    results["L4_libraries"] = {
        "memory_allocator": await lib.memory_allocator(),
    }

    return results
