#!/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   -> OK held measure ...   (until connection closes)
        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).
  - Profiles:
        measure : CPU 'performance' (power-profiles-daemon hold) +
                  GPU persistence ON + clocks LOCKED at GPU_LOCK_MHZ
                  -> run-to-run deterministic timing (per-item A/B evaluation)
        max     : CPU 'performance' + 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 (no client
    input ever reaches a command line); CPU hold uses the same D-Bus call the
    old attic/elphel-agent-tools/bin/cpu_perf_hold.py used (which this daemon
    supersedes).

Config via environment (set in the systemd unit or /etc/default/elphel-perfd):
    PERFD_PORT          control port      (default 48890)
    PERFD_GPU_LOCK_MHZ  'measure' lock    (default 2550; RTX 5060 Ti burst
                        boost measured ~2595 @ 23 W / 49 C on 07/17/2026 -
                        thermally trivial, so the value is about
                        reproducibility, not speed)
    PERFD_GPU_INDEX     GPU index         (default 0)
"""
import json
import os
import signal
import socket
import socketserver
import subprocess
import sys
import threading
import time

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

PROFILES = ("measure", "max")  # strongest first


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 = False

    def apply(self, profile):
        errs = []
        rc, out = _smi("-pm", "1")
        if rc != 0:
            errs.append("persistence: " + out)
        if profile == "measure":
            rc, out = _smi("-lgc", "{0},{0}".format(GPU_LOCK_MHZ))
            if rc != 0:
                errs.append("lgc: " + out)
            else:
                self.locked = True
        else:  # max - make sure no stale lock survives a profile downgrade
            self.release_lock()
        return errs

    def release_lock(self):
        if self.locked:
            rc, out = _smi("-rgc")
            if rc != 0:
                log("WARNING: -rgc failed: " + out)
            self.locked = False

    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": self.locked,
                "lock_mhz": GPU_LOCK_MHZ if self.locked else None,
                "telemetry": out if rc == 0 else "unavailable"}


# ---------------------------------------------------------------- CPU control
class CpuControl:
    """Hold the 'performance' profile via power-profiles-daemon (D-Bus) -
    the same mechanism cpu_perf_hold.py used; falls back to a no-op with a
    warning if the D-Bus service is unavailable."""
    def __init__(self):
        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) +
                ") - CPU holds will be no-ops")

    def hold(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(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 status(self):
        return {"held": self.cookie is not None}


# ------------------------------------------------------------------ hold book
class HoldBook:
    """Refcounted holds; the strongest requested profile is applied."""
    def __init__(self):
        self.lock = threading.Lock()
        self.counts = {p: 0 for p in PROFILES}
        self.applied = None
        self.gpu = GpuControl()
        self.cpu = CpuControl()

    def _target(self):
        for p in PROFILES:  # strongest first
            if self.counts[p] > 0:
                return p
        return None

    def _apply_target(self):
        target = self._target()
        if target == self.applied:
            return []
        errs = []
        if target is None:
            self.cpu.release()
            self.gpu.release_all()
            log("all holds released - default power management")
        else:
            self.cpu.hold()
            errs = self.gpu.apply(target)
            log("profile '%s' applied (gpu errors: %s)" %
                (target, errs if errs else "none"))
        self.applied = target
        return errs

    def acquire(self, profile):
        with self.lock:
            self.counts[profile] += 1
            return self._apply_target()

    def release(self, profile):
        with self.lock:
            if self.counts[profile] > 0:
                self.counts[profile] -= 1
            self._apply_target()

    def release_everything(self):
        with self.lock:
            for p in PROFILES:
                self.counts[p] = 0
            self._apply_target()

    def status(self):
        with self.lock:
            return {"applied": self.applied,
                    "holds": dict(self.counts),
                    "cpu": self.cpu.status(),
                    "gpu": self.gpu.status(),
                    "gpu_lock_mhz_configured": GPU_LOCK_MHZ}


BOOK = HoldBook()


# ------------------------------------------------------------------ TCP layer
class Handler(socketserver.StreamRequestHandler):
    def handle(self):
        held = None  # this connection's profile
        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" and (len(parts) == 2) and \
                        (parts[1].lower() in PROFILES):
                    prof = parts[1].lower()
                    if held is not None:       # one hold per connection
                        BOOK.release(held)
                    errs = BOOK.acquire(prof)
                    held = prof
                    self._reply("OK held %s%s" % (prof,
                                (" gpu_errors=" + ";".join(errs)) if errs else ""))
                elif cmd == "RELEASE":
                    if held is not None:
                        BOOK.release(held)
                        held = None
                    self._reply("OK released")
                elif cmd == "QUIT":
                    break
                else:
                    self._reply("ERR unknown command")
        except Exception:
            pass
        finally:
            if held is not None:               # crash-safe auto-release
                BOOK.release(held)

    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 (gpu lock for 'measure' = %d MHz)" %
        (PORT, GPU_LOCK_MHZ))
    try:
        srv.serve_forever()
    finally:
        BOOK.release_everything()


if __name__ == "__main__":
    main()
