import numpy as np
import psutil
import subprocess
import os
import time
from grub9 import HDGL_VM_32

# -------------------------
# Hardware Abstraction Mapper
# -------------------------
def map_hardware_to_vm(vm):
    # CPU cores → D1-D4
    cpu_loads = psutil.cpu_percent(percpu=True)
    for i in range(4):
        vm.D[i] = min(cpu_loads[i] / 100.0, 1.0)

    # Memory → D5-D12
    mem = psutil.virtual_memory()
    for i in range(8):
        vm.D[5 + i] = min(mem.total / (1024**3) / 32, 1.0)

    # Storage → D13-D20
    disks = psutil.disk_partitions()
    for i in range(min(8, len(disks))):
        usage = psutil.disk_usage(disks[i].mountpoint)
        vm.D[13 + i] = min(usage.percent / 100.0, 1.0)

    # Network → D21-D28
    net_io = psutil.net_io_counters(pernic=True)
    nets = list(net_io.keys())
    for i in range(min(8, len(nets))):
        vm.D[21 + i] = min(net_io[nets[i]].bytes_sent / (1024**2), 1.0)

    # GPU → D29-D32 (placeholder)
    for i in range(4):
        vm.D[29 + i] = 0.5

    # Omega tuning
    vm.omega = np.mean(vm.D)

# -------------------------
# GRUB Menu Builder
# -------------------------
def build_grub_entry(vm, kernel_path="/boot/vmlinuz-linux", initrd_path="/boot/initramfs-linux.img"):
    maxcpus = max(1, int(vm.D[0] * 32))
    mem_gb = max(1, int(np.sum(vm.D[5:13]) * 32))
    root_disk = "/dev/sda"
    kernel_params = f"root={root_disk} maxcpus={maxcpus} mem={mem_gb}G"

    vm.fs.setdefault('/boot', {})
    grub_cfg = vm.fs['/boot'].setdefault('grub.cfg', {})
    grub_cfg['menu'] = [("Auto-Mapped Linux Kernel", 0.12345)]
    grub_cfg['kernel_params'] = kernel_params
    vm.print(f"GRUB menu built with params: {kernel_params}")

    return kernel_path, initrd_path, kernel_params

# -------------------------
# Live HDGL REPL (Pre-Boot)
# -------------------------
def hdgl_live_repl(vm):
    vm.print("\n=== HDGL Live Lattice REPL ===")
    vm.print("Adjust D1-D32 or omega before booting Linux.")
    vm.print("Commands: set Dn val | omega val | step n | show | boot | exit")
    while True:
        try:
            cmd = input("HDGL> ").strip().split()
            if not cmd: continue
            instr = cmd[0].lower()
            if instr == "exit": break
            elif instr == "show":
                vm.print(f"D1-D32: {np.round(vm.D, 5)} Omega: {vm.omega:.5f}")
            elif instr == "set" and len(cmd) == 3:
                idx = int(cmd[1][1:]) - 1
                val = float(cmd[2])
                vm.D[idx] = val
            elif instr == "omega" and len(cmd) == 2:
                vm.omega = float(cmd[1])
            elif instr == "step":
                n = int(cmd[1]) if len(cmd) > 1 else 1
                for _ in range(n): vm.step()
            elif instr == "boot":
                vm.print("Exiting REPL and booting Linux...")
                break
            else:
                vm.print("Unknown command. Use set, omega, step, show, boot, exit")
        except Exception as e:
            vm.print(f"Error: {e}")

# -------------------------
# Kernel Launcher via kexec
# -------------------------
def boot_linux(vm):
    kernel, initrd, params = build_grub_entry(vm)

    vm.fs['/boot']['last_boot_params'] = params

    vm.print("Seeding lattice before boot...")
    for _ in range(20):
        vm.step()

    vm.print("Launching Linux kernel via kexec...")
    subprocess.run([
        "sudo", "kexec", "-l", kernel,
        "--initrd=" + initrd,
        "--append=" + params
    ], check=True)

    subprocess.run(["sudo", "kexec", "-e"], check=True)

# -------------------------
# End-to-End Auto-Boot
# -------------------------
if __name__ == "__main__":
    vm = HDGL_VM_32()
    map_hardware_to_vm(vm)
    hdgl_live_repl(vm)
    boot_linux(vm)
