#!/usr/bin/env python3
"""
elphel-perfd - fleet performance-management daemon (CPU + GPU), so that
unprivileged programs (Eyesis/imagej-elphel, agents, scripts) can request
deterministic or maximum-performance clock states WITHOUT sudo.

Design ruled by Andrey 07/17/2026 (lives in the imagej-elphel codebase: the
capability follows the hardware the project runs on - 224/137 - while other
programs use it through the same socket/CLI). Implemented by Claude.

  - Runs as a root systemd service (one-time `sudo perfd/install.sh`).
  - Control = line protocol on 127.0.0.1:48890 (localhost only; TCP so that
    Java 8/11 clients need no extra libraries):
        PING                                   -> OK pong
        STATUS                                 -> OK {json}
        HOLD measure [cpu=KHZ] [gpu=MHZ]       -> OK held measure ...
        HOLD max                               -> OK held max ...
        RELEASE                                -> OK released
    A hold is TIED TO THE CONNECTION: client exit/crash closes the socket and
    the hold is released automatically (crash-safe, like a media player's
    sleep inhibitor). Multiple concurrent holders are refcounted; the
    strongest profile wins (measure > max); lock values come from the newest
    holder of the winning profile. cpu=/gpu= overrides are bounds-checked and
    exist so lock-value calibration needs no daemon reconfiguration.
  - Profiles:
        measure : CPU frequency PINNED (intel_pstate scaling_min=max=KHZ) +
                  GPU persistence ON + GPU clocks LOCKED at MHZ
                  -> run-to-run deterministic timing for per-item A/B
                  evaluation. PINNED, not 'performance': the i9-11900K at
                  free boost hits 4.9 GHz -> 100 C -> throttle wobble
                  (3.2-3.5 GHz) within seconds of an all-core phase - the
                  boost bias itself is a nondeterminism source (measured
                  07/17/2026).
        max     : CPU 'performance' profile (power-profiles-daemon hold,
                  free boost) + GPU persistence ON, clocks unlocked
                  -> production throughput.
        (idle)  : everything released; also enforced on daemon stop.
  - GPU ops shell out to nvidia-smi with fixed argument lists; CPU pinning
    writes sysfs cpufreq directly (root). No client input ever reaches a
    command line. Supersedes attic/elphel-agent-tools/bin/cpu_perf_hold.py.

Config via environment (systemd unit):
    PERFD_PORT          control port         (default 48890)
    PERFD_GPU_LOCK_MHZ  'measure' GPU lock   (default 2550; RTX 5060 Ti burst
                        boost measured ~2595 @ 23 W / 49 C - thermally trivial)
    PERFD_CPU_LOCK_KHZ  'measure' CPU pin    (default 3200000; i9-11900K
                        all-core THROTTLE equilibrium measured 3.2-3.5 GHz
                        @ 100 C - the pin must sit at/below it to stay
                        deterministic AND cooler)
    PERFD_GPU_INDEX     GPU index            (default 0)
"""
import glob
import json
import os
import signal
import socketserver
import subprocess
import sys
import threading

PORT         = int(os.environ.get("PERFD_PORT", "48890"))
GPU_LOCK_MHZ = int(os.environ.get("PERFD_GPU_LOCK_MHZ", "2550"))
CPU_LOCK_KHZ = int(os.environ.get("PERFD_CPU_LOCK_KHZ", "3200000"))
GPU_INDEX    = os.environ.get("PERFD_GPU_INDEX", "0")
NVIDIA_SMI   = "/usr/bin/nvidia-smi"

PROFILES = ("measure", "max")  # strongest first
GPU_MHZ_BOUNDS = (180, 3090)
CPU_KHZ_BOUNDS = (800000, 5300000)


def log(msg):
    print("[elphel-perfd] " + msg, flush=True)


# ---------------------------------------------------------------- GPU control
def _smi(*args):
    """Run nvidia-smi with a FIXED argument list; returns (rc, output)."""
    try:
        r = subprocess.run((NVIDIA_SMI, "-i", GPU_INDEX) + args,
                           capture_output=True, text=True, timeout=20)
        out = (r.stdout + r.stderr).strip()
        return r.returncode, out
    except Exception as e:  # missing binary, timeout
        return -1, str(e)


class GpuControl:
    def __init__(self):
        self.locked_mhz = None

    def apply(self, profile, gpu_mhz):
        errs = []
        rc, out = _smi("-pm", "1")
        if rc != 0:
            errs.append("persistence: " + out)
        if profile == "measure":
            rc, out = _smi("-lgc", "{0},{0}".format(gpu_mhz))
            if rc != 0:
                errs.append("lgc: " + out)
                self.locked_mhz = None
            else:
                self.locked_mhz = gpu_mhz
        else:  # max - no stale lock may survive a profile downgrade
            self.release_lock()
        return errs

    def release_lock(self):
        if self.locked_mhz is not None:
            rc, out = _smi("-rgc")
            if rc != 0:
                log("WARNING: -rgc failed: " + out)
            self.locked_mhz = None

    def release_all(self):
        # persistence mode is left ON (harmless, avoids driver re-init cost)
        self.release_lock()

    def status(self):
        rc, out = _smi("--query-gpu=clocks.sm,temperature.gpu,power.draw",
                       "--format=csv,noheader")
        return {"locked_mhz": self.locked_mhz,
                "telemetry": out if rc == 0 else "unavailable"}


# ---------------------------------------------------------------- CPU control
class CpuControl:
    """measure: PIN the frequency (scaling_min=max=KHZ on every cpufreq
    policy) - deterministic and below the thermal wall. max: hold the
    'performance' profile via power-profiles-daemon (free boost), the same
    call the retired cpu_perf_hold.py used."""
    def __init__(self):
        self.policies = sorted(glob.glob(
            "/sys/devices/system/cpu/cpufreq/policy*"))
        self.hw_bounds = {}
        for p in self.policies:
            try:
                lo = int(open(p + "/cpuinfo_min_freq").read())
                hi = int(open(p + "/cpuinfo_max_freq").read())
                self.hw_bounds[p] = (lo, hi)
            except OSError:
                pass
        self.pinned_khz = None
        self.cookie = None
        self.iface = None
        try:
            import dbus
            bus = dbus.SystemBus()
            obj = bus.get_object("org.freedesktop.UPower.PowerProfiles",
                                 "/org/freedesktop/UPower/PowerProfiles")
            self.iface = dbus.Interface(obj,
                                        "org.freedesktop.UPower.PowerProfiles")
        except Exception as e:
            log("WARNING: power-profiles-daemon unavailable (" + str(e) +
                ") - 'max' CPU holds will be no-ops")

    @staticmethod
    def _write(path, value):
        with open(path, "w") as f:
            f.write(str(value))

    def pin(self, khz):
        errs = []
        for p, (lo, hi) in self.hw_bounds.items():
            k = max(lo, min(hi, khz))
            try:
                # safe order: open the floor, set the ceiling, raise the floor
                self._write(p + "/scaling_min_freq", lo)
                self._write(p + "/scaling_max_freq", k)
                self._write(p + "/scaling_min_freq", k)
            except OSError as e:
                errs.append(os.path.basename(p) + ": " + str(e))
        if not errs:
            self.pinned_khz = khz
        return errs

    def unpin(self):
        if self.pinned_khz is None:
            return
        for p, (lo, hi) in self.hw_bounds.items():
            try:
                self._write(p + "/scaling_min_freq", lo)
                self._write(p + "/scaling_max_freq", hi)
            except OSError as e:
                log("WARNING: unpin " + os.path.basename(p) + ": " + str(e))
        self.pinned_khz = None

    def hold_ppd(self):
        if (self.iface is not None) and (self.cookie is None):
            try:
                self.cookie = self.iface.HoldProfile(
                    "performance", "elphel-perfd client hold",
                    "com.elphel.perfd")
            except Exception as e:
                log("WARNING: CPU HoldProfile failed: " + str(e))

    def release_ppd(self):
        if (self.iface is not None) and (self.cookie is not None):
            try:
                self.iface.ReleaseProfile(self.cookie)
            except Exception as e:
                log("WARNING: CPU ReleaseProfile failed: " + str(e))
            self.cookie = None

    def apply(self, profile, cpu_khz):
        if profile == "measure":
            self.release_ppd()
            return self.pin(cpu_khz)
        else:  # max
            self.unpin()
            self.hold_ppd()
            return []

    def release_all(self):
        self.unpin()
        self.release_ppd()

    def status(self):
        return {"pinned_khz": self.pinned_khz,
                "ppd_held": self.cookie is not None,
                "policies": len(self.hw_bounds)}


# ------------------------------------------------------------------ hold book
class HoldBook:
    """Per-connection holds; the strongest profile wins; lock values come
    from the NEWEST holder of the winning profile."""
    def __init__(self):
        self.lock = threading.Lock()
        self.holds = {}  # conn_id -> (seq, profile, cpu_khz, gpu_mhz)
        self.seq = 0
        self.applied = None  # (profile, cpu_khz, gpu_mhz) or None
        self.gpu = GpuControl()
        self.cpu = CpuControl()

    def _target(self):
        for prof in PROFILES:  # strongest first
            best = None
            for (seq, p, ck, gm) in self.holds.values():
                if (p == prof) and ((best is None) or (seq > best[0])):
                    best = (seq, p, ck, gm)
            if best is not None:
                return best[1:]
        return None

    def _apply_target(self):
        target = self._target()
        if target == self.applied:
            return []
        errs = []
        if target is None:
            self.cpu.release_all()
            self.gpu.release_all()
            log("all holds released - default power management")
        else:
            profile, cpu_khz, gpu_mhz = target
            errs = self.cpu.apply(profile, cpu_khz) + \
                   self.gpu.apply(profile, gpu_mhz)
            log("profile '%s' applied (cpu=%d kHz gpu=%d MHz; errors: %s)" %
                (profile, cpu_khz, gpu_mhz, errs if errs else "none"))
        self.applied = target
        return errs

    def acquire(self, conn_id, profile, cpu_khz, gpu_mhz):
        with self.lock:
            self.seq += 1
            self.holds[conn_id] = (self.seq, profile, cpu_khz, gpu_mhz)
            return self._apply_target()

    def release(self, conn_id):
        with self.lock:
            self.holds.pop(conn_id, None)
            self._apply_target()

    def release_everything(self):
        with self.lock:
            self.holds.clear()
            self._apply_target()

    def status(self):
        with self.lock:
            return {"applied": (None if self.applied is None else {
                        "profile": self.applied[0],
                        "cpu_khz": self.applied[1],
                        "gpu_mhz": self.applied[2]}),
                    "holders": len(self.holds),
                    "cpu": self.cpu.status(),
                    "gpu": self.gpu.status(),
                    "defaults": {"cpu_khz": CPU_LOCK_KHZ,
                                 "gpu_mhz": GPU_LOCK_MHZ}}


BOOK = HoldBook()


# ------------------------------------------------------------------ TCP layer
def _parse_hold(parts):
    """HOLD measure [cpu=KHZ] [gpu=MHZ] -> (profile, cpu_khz, gpu_mhz) or None."""
    if (len(parts) < 2) or (parts[1].lower() not in PROFILES):
        return None
    profile, cpu_khz, gpu_mhz = parts[1].lower(), CPU_LOCK_KHZ, GPU_LOCK_MHZ
    for tok in parts[2:]:
        if "=" not in tok:
            return None
        key, _, val = tok.partition("=")
        try:
            v = int(val)
        except ValueError:
            return None
        if key == "cpu" and CPU_KHZ_BOUNDS[0] <= v <= CPU_KHZ_BOUNDS[1]:
            cpu_khz = v
        elif key == "gpu" and GPU_MHZ_BOUNDS[0] <= v <= GPU_MHZ_BOUNDS[1]:
            gpu_mhz = v
        else:
            return None
    return profile, cpu_khz, gpu_mhz


class Handler(socketserver.StreamRequestHandler):
    def handle(self):
        conn_id = id(self)
        holding = False
        try:
            for raw in self.rfile:
                parts = raw.decode("utf-8", "replace").strip().split()
                if not parts:
                    continue
                cmd = parts[0].upper()
                if cmd == "PING":
                    self._reply("OK pong")
                elif cmd == "STATUS":
                    self._reply("OK " + json.dumps(BOOK.status()))
                elif cmd == "HOLD":
                    parsed = _parse_hold(parts)
                    if parsed is None:
                        self._reply("ERR usage: HOLD measure|max [cpu=KHZ] [gpu=MHZ] (bounds cpu %d..%d, gpu %d..%d)" %
                                    (CPU_KHZ_BOUNDS + GPU_MHZ_BOUNDS))
                        continue
                    errs = BOOK.acquire(conn_id, *parsed)
                    holding = True
                    self._reply("OK held %s cpu=%d gpu=%d%s" %
                                (parsed[0], parsed[1], parsed[2],
                                 (" errors=" + ";".join(errs)) if errs else ""))
                elif cmd == "RELEASE":
                    BOOK.release(conn_id)
                    holding = False
                    self._reply("OK released")
                elif cmd == "QUIT":
                    break
                else:
                    self._reply("ERR unknown command")
        except Exception:
            pass
        finally:
            if holding:  # crash-safe auto-release
                BOOK.release(conn_id)

    def _reply(self, line):
        self.wfile.write((line + "\n").encode("utf-8"))
        self.wfile.flush()


class Server(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True


def main():
    def shutdown(signum, frame):
        log("signal %d - releasing everything and exiting" % signum)
        BOOK.release_everything()
        sys.exit(0)
    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)
    srv = Server(("127.0.0.1", PORT), Handler)  # localhost ONLY
    log("listening on 127.0.0.1:%d (measure defaults: cpu=%d kHz, gpu=%d MHz)" %
        (PORT, CPU_LOCK_KHZ, GPU_LOCK_MHZ))
    try:
        srv.serve_forever()
    finally:
        BOOK.release_everything()


if __name__ == "__main__":
    main()
