"""
wivexx_optimize.py — on-demand PC optimizer  (v1.0)
═══════════════════════════════════════════════════════════════════════
Our OWN optimizer (no external apps, no admin). Honest about what's actually
possible from user space:

  • RAM  → the REAL lever. EmptyWorkingSet() trims each process's working set,
           pushing idle pages out of physical RAM onto the standby list, so
           "available" memory goes up measurably. Apps re-fault pages as needed.
  • CPU  → you can't "free" CPU on demand; the honest help is to REPORT the top
           background CPU hogs so you know what to close.
  • GPU  → likewise can't be trimmed from user space; we REPORT what's holding
           VRAM (nvidia-smi compute/graphics apps).

Public API (called from wivexx_api):
    ram_trim() -> dict                       # {trimmed, freed_mb, before_mb, after_mb}
    top_cpu(n=3) -> list[(name, pct)]
    top_vram(n=3) -> list[(name, mb)]
    optimize(field="all") -> dict            # field: ram|cpu|gpu|all
    summary_line(result) -> str              # a spoken one-liner
"""

import ctypes
import os
import re
import subprocess
import time
from ctypes import wintypes
from typing import Optional

_NOWIN = getattr(subprocess, "CREATE_NO_WINDOW", 0)

# kernel/idle pseudo-processes that aren't real "CPU usage" — never report them
_CPU_SKIP = {"system idle process", "idle", "system", "memcompression",
             "memory compression", "registry", "secure system"}

# Win32: EmptyWorkingSet needs a handle opened with these rights (granted for
# same-user processes without elevation; system procs fail → skipped silently).
_PROCESS_SET_QUOTA = 0x0100
_PROCESS_QUERY_LIMITED_INFORMATION = 0x1000

try:
    _k32 = ctypes.WinDLL("kernel32", use_last_error=True)
    _psapi = ctypes.WinDLL("psapi", use_last_error=True)
    _k32.OpenProcess.restype = wintypes.HANDLE
    _k32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
    _k32.CloseHandle.argtypes = [wintypes.HANDLE]
    _k32.CloseHandle.restype = wintypes.BOOL
    _psapi.EmptyWorkingSet.restype = wintypes.BOOL
    _psapi.EmptyWorkingSet.argtypes = [wintypes.HANDLE]
    _WIN_OK = True
except Exception:
    _WIN_OK = False

# never bother trimming the kernel/idle pseudo-processes
_SKIP_PIDS = {0, 4}


def _trim_pid(pid: int) -> bool:
    if not _WIN_OK:
        return False
    h = _k32.OpenProcess(_PROCESS_SET_QUOTA | _PROCESS_QUERY_LIMITED_INFORMATION,
                         False, pid)
    if not h:
        return False
    try:
        return bool(_psapi.EmptyWorkingSet(h))
    finally:
        _k32.CloseHandle(h)


def ram_trim() -> dict:
    """Trim every accessible process's working set → free physical RAM. Returns
    the freed amount (available-memory delta, MB)."""
    import psutil
    before = psutil.virtual_memory().available
    trimmed = 0
    for p in psutil.process_iter(["pid"]):
        pid = p.info.get("pid", 0)
        if pid in _SKIP_PIDS:
            continue
        try:
            if _trim_pid(pid):
                trimmed += 1
        except Exception:
            pass
    time.sleep(0.4)                       # let the memory manager settle
    after = psutil.virtual_memory().available
    freed_mb = max(0.0, (after - before) / (1024 * 1024))
    return {"trimmed": trimmed, "freed_mb": round(freed_mb),
            "before_mb": round(before / 1024 / 1024),
            "after_mb": round(after / 1024 / 1024)}


def top_cpu(n: int = 3) -> list:
    """The n busiest processes by CPU% (sampled over ~0.6 s), aggregated by name."""
    import psutil
    for p in psutil.process_iter():
        try:
            p.cpu_percent(None)           # prime the counters
        except Exception:
            pass
    time.sleep(0.6)
    agg: dict = {}
    ncpu = psutil.cpu_count() or 1
    for p in psutil.process_iter(["name"]):
        try:
            nm = (p.info.get("name") or "?")
            if nm.lower() in _CPU_SKIP:        # skip idle/system pseudo-processes
                continue
            c = p.cpu_percent(None) / ncpu     # normalize to whole-CPU %
            if c > 0:
                agg[nm] = agg.get(nm, 0.0) + c
        except Exception:
            pass
    rows = sorted(agg.items(), key=lambda kv: kv[1], reverse=True)
    return [(nm, round(c, 1)) for nm, c in rows[:n] if c >= 1.0]


import shutil
_NVSMI = ""


def _nvidia_smi() -> str:
    """Full path to nvidia-smi, robust to the env-stripped VBS launch (where PATH
    is gone). Tries PATH, then %SystemRoot%\\System32 (where the driver installs
    it), then the legacy NVSMI dir. Never caches a MISS — a later launch with a
    proper env can still find it."""
    global _NVSMI
    if _NVSMI:
        return _NVSMI
    cand = shutil.which("nvidia-smi") or ""
    if not cand:
        root = os.environ.get("SystemRoot") or os.environ.get("windir") or r"C:\Windows"
        for p in (os.path.join(root, "System32", "nvidia-smi.exe"),
                  r"C:\Program Files\NVIDIA Corporation\NVSMI\nvidia-smi.exe"):
            if os.path.isfile(p):
                cand = p
                break
    if cand:
        _NVSMI = cand          # cache only a successful find
    return cand


# desktop/shell processes that are ALWAYS on the GPU — not worth flagging to close
_GPU_NOISE = {"explorer.exe", "dwm.exe", "shellexperiencehost.exe",
             "systemsettings.exe", "textinputhost.exe", "searchhost.exe",
             "startmenuexperiencehost.exe", "modernflyoutshost.exe",
             "translucenttb.exe", "applicationframehost.exe", "pythonw.exe"}


def gpu_apps(limit: int = 4) -> list:
    """Distinct notable apps currently using the GPU. Per-process VRAM is N/A on
    consumer GeForce in WDDM mode, so we can only list WHO's on the GPU (handy to
    know what to close), not how much each holds."""
    names, seen = [], set()
    smi = _nvidia_smi()
    if not smi:
        return names
    try:
        r = subprocess.run([smi], capture_output=True, timeout=5,
                           creationflags=_NOWIN)
        in_proc = False
        for ln in r.stdout.decode(errors="ignore").splitlines():
            if "Processes:" in ln:
                in_proc = True
                continue
            if not in_proc:
                continue
            m = re.search(r"\b\d+\s+[A-Z]\+?[A-Z]?\s+(.+?\.exe)\b", ln, re.IGNORECASE)
            if m:
                nm = os.path.basename(m.group(1).strip())
                low = nm.lower()
                if low in seen or low in _GPU_NOISE:
                    continue
                seen.add(low)
                names.append(nm)
    except Exception:
        pass
    return names[:limit]


def gpu_load() -> Optional[dict]:
    """Total VRAM usage + the notable apps on the GPU (per-app VRAM unavailable on
    consumer GPUs in WDDM, so this is the honest GPU picture)."""
    smi = _nvidia_smi()
    if not smi:
        return None
    try:
        r = subprocess.run(
            [smi, "--query-gpu=memory.used,memory.total",
             "--format=csv,noheader,nounits"],
            capture_output=True, timeout=4, creationflags=_NOWIN)
        u, t = [float(x.strip()) for x in
                r.stdout.decode(errors="ignore").strip().splitlines()[0].split(",")]
        return {"used_mb": round(u), "total_mb": round(t),
                "pct": round(100.0 * u / t, 1) if t else None,
                "apps": gpu_apps()}
    except Exception:
        return None


def optimize(field: str = "all") -> dict:
    """Run the requested optimization. field: ram | cpu | gpu | all."""
    field = (field or "all").lower()
    res: dict = {"field": field}
    if field in ("ram", "all", "cpu"):     # the RAM trim helps CPU pressure too
        try:
            res["ram"] = ram_trim()
        except Exception:
            res["ram"] = None
    if field in ("cpu", "all"):
        try:
            res["cpu"] = top_cpu()
        except Exception:
            res["cpu"] = []
    if field in ("gpu", "all"):
        try:
            res["gpu"] = gpu_load()
        except Exception:
            res["gpu"] = None
    return res


def summary_line(result: dict) -> str:
    """A natural spoken/console summary of an optimize() result."""
    bits = []
    ram = result.get("ram")
    if ram:
        if ram.get("freed_mb"):
            bits.append("freed about %d megabytes of RAM" % ram["freed_mb"])
        else:
            bits.append("RAM was already tight — little to reclaim")
    cpu = result.get("cpu")
    if cpu:
        nm, pct = cpu[0]
        bits.append("top CPU is %s at %.0f%%" % (nm, pct))
    gpu = result.get("gpu")
    if gpu and gpu.get("pct") is not None:
        line = "GPU memory is %.0f%% used" % gpu["pct"]
        if gpu.get("apps"):
            line += " (apps on it: %s)" % ", ".join(gpu["apps"][:3])
        bits.append(line)
    if not bits:
        return "Nothing to optimize right now, sir."
    return "Done, sir — " + "; ".join(bits) + "."


if __name__ == "__main__":
    import json
    r = optimize("all")
    print(json.dumps(r, indent=2, default=str))
    print(summary_line(r))
