"""predatord – privileged background service.

* line-based JSON protocol over a Unix socket (root and members of the allowed group only)
* writes a state snapshot to /run/predator-control/state.json every 1.5 s
  (read by the Plasma panel widget)
* restores settings at boot and after resume
* switches power profiles when the charger is plugged in / unplugged
* drives the fans along a user-defined temperature curve
"""

import copy
import grp
import json
import os
import pwd
import signal
import socket
import socketserver
import struct
import sys
import threading
import time
import traceback

from . import CONFIG_DIR, CONFIG_PATH, SOCKET_PATH, STATE_PATH, __version__
from . import hardware as hw
from . import updater
from .hotkey import HotkeyListener, launch_in_session
from .idle import IdleWatcher
from .keyboard import EFFECTS, NUM_CELLS, Keyboard

DEFAULT_CONFIG = {
    "allowed_group": "wheel",
    "profile": {"ac": "balanced", "battery": "quiet", "auto_switch": True},
    "fans": {
        "mode": "auto",           # auto | max | custom | curve
        "cpu": 50, "gpu": 50,
        "curve": {
            "cpu": [[45, 0], [60, 35], [72, 55], [82, 75], [90, 100]],
            "gpu": [[45, 0], [60, 35], [70, 55], [78, 75], [85, 100]],
        },
    },
    "keyboard": {
        "effect": "static", "color": "#00c8ff", "brightness": 200, "effect_id": 0,
        "frame": "",               # 384 bytes as hex (perkey)
    },
    "lightbar": {
        "mode": "static", "color": "#00c8ff", "brightness": 200, "speed": 5,
        "direction": 1, "sync_keyboard": True,
    },
    "battery": {"limit": None},   # None = leave the firmware state alone
    "keyboard_presets": {},
    "updates": {"auto_check": True, "interval_hours": 6},
    "predator_key": "open",       # open | cycle_profile | none
    "power": {
        "idle_off_ac": 0,         # seconds without input until lighting turns off, 0 = never
        "idle_off_battery": 120,
        "battery_brightness": 100,   # keyboard brightness on battery, % of the normal value
        "battery_lights_off": False,
        "battery_lightbar_off": False,
    },
}


def hex_to_rgb(s):
    s = s.lstrip("#")
    return tuple(int(s[i:i + 2], 16) for i in (0, 2, 4))


def merge(base, upd):
    for k, v in upd.items():
        if isinstance(v, dict) and isinstance(base.get(k), dict):
            merge(base[k], v)
        else:
            base[k] = v
    return base


def interp(points, x):
    pts = sorted(points)
    if not pts:
        return 0
    if x <= pts[0][0]:
        return pts[0][1]
    for (x0, y0), (x1, y1) in zip(pts, pts[1:]):
        if x <= x1:
            return y0 + (y1 - y0) * (x - x0) / max(1e-6, x1 - x0)
    return pts[-1][1]


def log(*a):
    print(*a, flush=True)


class Daemon:
    def __init__(self):
        self.lock = threading.RLock()
        self.kbd = Keyboard()
        self.cpuload = CpuLoadProxy()
        self.config = self._load_config()
        self.errors = {}
        self.last_ac = None
        self.curve_applied = None
        self.kbd_present = False
        self.state = {}
        self.battery_limit = None
        self._ticks = 0
        self.update_info = None
        self.update_error = None
        self._update_checking = False
        self.idle = IdleWatcher(self._activity)
        self.lights_idle = False
        self.on_ac = True

    # ---------------------------------------------------------- Configuration
    def _load_config(self):
        cfg = copy.deepcopy(DEFAULT_CONFIG)
        try:
            with open(CONFIG_PATH) as f:
                merge(cfg, json.load(f))
        except FileNotFoundError:
            pass
        except Exception as e:
            log("Invalid configuration, using defaults:", e)
        return cfg

    def save_config(self):
        os.makedirs(CONFIG_DIR, exist_ok=True)
        tmp = CONFIG_PATH + ".tmp"
        with open(tmp, "w") as f:
            json.dump(self.config, f, indent=2)
        os.replace(tmp, CONFIG_PATH)

    # ---------------------------------------------------------- Applying
    def _try(self, key, fn, *a, **kw):
        try:
            r = fn(*a, **kw)
            self.errors.pop(key, None)
            return r
        except Exception as e:
            if self.errors.get(key) != str(e):
                log("[%s] %s" % (key, e))
            self.errors[key] = str(e)
            return None

    # ---------------------------------------------------------- Power rules
    def keyboard_effective(self):
        k = dict(self.config["keyboard"])
        p = self.config["power"]
        if self.lights_idle or (not self.on_ac and p.get("battery_lights_off")):
            k["effect"] = "off"
        elif not self.on_ac and p.get("battery_brightness", 100) < 100:
            k["brightness"] = int(k["brightness"] * p["battery_brightness"] / 100)
        return k

    def idle_timeout(self):
        p = self.config["power"]
        return int(p.get("idle_off_ac" if self.on_ac else "idle_off_battery") or 0)

    def update_idle(self):
        t = self.idle_timeout()
        idle = t > 0 and self.idle.idle_seconds() >= t
        if idle != self.lights_idle:
            self.lights_idle = idle
            log("Lighting %s" % ("off (idle)" if idle else "restored"))
            self.apply_keyboard()

    def _activity(self):
        if self.lights_idle:
            with self.lock:
                self.update_idle()

    def apply_keyboard(self, cfg=None):
        k = cfg or self.keyboard_effective()
        frame = bytes.fromhex(k.get("frame") or "") if k["effect"] == "perkey" else None
        self._try("keyboard", self.kbd.apply, k["effect"], hex_to_rgb(k["color"]),
                  k["brightness"], frame, int(k.get("effect_id") or 0))
        if self.config["lightbar"].get("sync_keyboard") and cfg is None:
            self.apply_lightbar()

    def lightbar_effective(self):
        lb = dict(self.config["lightbar"])
        if self.lights_idle or (not self.on_ac and self.config["power"].get("battery_lightbar_off")):
            lb["mode"] = "off"
            return lb
        if lb.get("sync_keyboard"):
            k = self.keyboard_effective()
            lb["brightness"] = k["brightness"]
            if k["effect"] == "off":
                lb["mode"] = "off"
            elif k["effect"] in ("rainbow", "neon"):
                lb["mode"] = "rainbow" if k["effect"] == "rainbow" else "neon"
            elif k["effect"] == "breathing":
                lb["mode"], lb["color"] = "breathing", k["color"]
            elif k["effect"] == "perkey":
                pass  # keep the lightbar's own setting
            else:
                lb["mode"], lb["color"] = "static", k["color"]
        return lb

    def apply_lightbar(self):
        lb = self.lightbar_effective()
        self._try("lightbar", hw.set_lightbar, lb["mode"], hex_to_rgb(lb["color"]),
                  lb["brightness"], lb["speed"], lb["direction"])

    def apply_fans(self):
        f = self.config["fans"]
        self.curve_applied = None
        if f["mode"] == "curve":
            self.fan_curve_tick(force=True)
        else:
            self._try("fans", hw.set_fans, f["mode"], f["cpu"], f["gpu"])

    def apply_profile_for_power(self, ac):
        p = self.config["profile"]
        target = p["ac"] if ac else p["battery"]
        if target in hw.profile_choices():
            self._try("profile", hw.set_profile, target)

    def restore(self):
        log("Restoring settings")
        bat = hw.battery()
        self.on_ac = bat.get("ac_online", True)
        self.idle.last_input = time.monotonic()   # e.g. after resume: lights on
        self.lights_idle = False
        if self.config["profile"].get("auto_switch"):
            self.apply_profile_for_power(bat.get("ac_online", True))
        self.apply_fans()
        if self.kbd.available():
            self.apply_keyboard()
        else:
            self.apply_lightbar()
        lim = self.config["battery"].get("limit")
        if lim is not None and hw.wmi_available(hw.GUID_BATTERY):
            self._try("battery", hw.set_battery_limit, lim)
        self.refresh_battery_limit()

    def refresh_battery_limit(self):
        if hw.wmi_available(hw.GUID_BATTERY):
            try:
                self.battery_limit = hw.get_battery_limit()
            except Exception:
                self.battery_limit = None

    # ---------------------------------------------------------- Background
    def fan_curve_tick(self, force=False):
        f = self.config["fans"]
        if f["mode"] != "curve":
            return
        s = hw.sensors()
        cpu_t = s.get("cpu_temp") or 0
        gpu_t = s.get("gpu_temp") or cpu_t
        cpu = int(round(interp(f["curve"]["cpu"], cpu_t) / 5) * 5)
        gpu = int(round(interp(f["curve"]["gpu"], gpu_t) / 5) * 5)
        if s.get("dgpu_active") is False:
            gpu = min(gpu, cpu)
        target = (cpu, gpu)
        prev = self.curve_applied
        # hysteresis: ramp up immediately, ramp down only after a >=10 % drop
        if not force and prev is not None:
            if all(t <= p and p - t < 10 for t, p in zip(target, prev)) or target == prev:
                return
        self._try("fans", hw.set_fans, "custom", cpu, gpu)
        self.curve_applied = target

    def build_state(self):
        s = hw.sensors()
        bat = hw.battery(limit_80=bool(self.battery_limit))
        st = {
            "version": __version__,
            "time": time.time(),
            "sensors": s,
            "cpu_load": self.cpuload.value,
            "cpu_freq": hw.cpu_freq_mhz(),
            "memory": hw.mem_usage(),
            "battery": bat,
            "profile": hw.get_profile(),
            "profile_choices": hw.profile_choices(),
            "profile_labels": hw.PROFILE_LABELS,
            "config": self.config,
            "lightbar_effective": self.lightbar_effective(),
            "fan_curve_applied": self.curve_applied,
            "battery_limit": self.battery_limit,
            "lights_idle": self.lights_idle,
            "update": self.update_info,
            "update_error": self.update_error,
            "update_status": updater.status(),
            "available": {
                "keyboard": self.kbd_present,
                "wmi": hw.wmi_available(hw.GUID_GAMING),
                "battery_wmi": hw.wmi_available(hw.GUID_BATTERY),
            },
            "errors": dict(self.errors),
            "keyboard_effects": {k: [v[2], v[1]] for k, v in EFFECTS.items()},  # [label, uses colour]
            "lightbar_modes": {k: v[1] for k, v in hw.LIGHTBAR_MODES.items()},
        }
        return st

    def check_update_async(self):
        if self._update_checking:
            return
        self._update_checking = True

        def work():
            try:
                info = updater.check()
                if info["available"] and (self.update_info or {}).get("latest") != info["latest"]:
                    log("Update available: %s -> %s" % (info["current"], info["latest"]))
                self.update_info, self.update_error = info, None
            except Exception as e:
                self.update_error = str(e)
            finally:
                self._update_checking = False
        threading.Thread(target=work, daemon=True).start()

    def maybe_check_update(self):
        u = self.config.get("updates", {})
        if not u.get("auto_check", True):
            return
        last = (self.update_info or {}).get("checked", 0)
        if time.time() - last > float(u.get("interval_hours", 6)) * 3600:
            if self.update_error and time.time() - getattr(self, "_last_try", 0) < 900:
                return  # offline: retry every 15 min
            self._last_try = time.time()
            self.check_update_async()

    def predator_key(self):
        action = self.config.get("predator_key", "open")
        if action == "open":
            if not launch_in_session(["predator-control"]):
                log("Predator key: no active graphical session to open the app in")
        elif action == "cycle_profile":
            with self.lock:
                ch = hw.profile_choices()
                cur = hw.get_profile()
                nxt = ch[(ch.index(cur) + 1) % len(ch)] if cur in ch else ch[0]
                self._try("profile", hw.set_profile, nxt)

    def loop(self):
        while True:
            try:
                self.cpuload.sample()
                with self.lock:
                    present = self.kbd.available()
                    if present and not self.kbd_present:
                        # keyboard (re)appeared, e.g. after resume
                        if self.kbd_present is not None:
                            time.sleep(0.5)
                            self.apply_keyboard()
                    self.kbd_present = present
                    ac = hw.battery().get("ac_online", True)
                    if self.last_ac is not None and ac != self.last_ac:
                        log("Power source changed:", "AC" if ac else "battery")
                        if self.config["profile"].get("auto_switch"):
                            self.apply_profile_for_power(ac)
                    if ac != self.on_ac:
                        self.on_ac = ac
                        self.apply_keyboard()  # battery lighting rules
                    self.last_ac = ac
                    self.update_idle()
                    self.fan_curve_tick()
                    self._ticks += 1
                    if self._ticks % 20 == 1:
                        self.maybe_check_update()
                    if self._ticks % 40 == 0:
                        self.refresh_battery_limit()
                    self.state = self.build_state()
                write_state(self.state)
            except Exception:
                traceback.print_exc()
            time.sleep(1.5)

    # ---------------------------------------------------------- Commands
    def handle(self, req):
        cmd = req.get("cmd")
        a = req.get("args", {}) or {}
        with self.lock:
            if cmd == "ping":
                return {"version": __version__}
            if cmd == "state":
                self.state = self.build_state()
                return self.state
            if cmd == "check_update":
                self.update_info, self.update_error = updater.check(), None
                return self.update_info
            if cmd == "install_update":
                info = updater.check()
                self.update_info = info
                return updater.install(info)
            if cmd == "set_updates":
                merge(self.config["updates"], {k: a[k] for k in ("auto_check", "interval_hours") if k in a})
                self.save_config()
                return {}
            if cmd == "set_power":
                p = self.config["power"]
                for key in ("idle_off_ac", "idle_off_battery", "battery_brightness",
                            "battery_lights_off", "battery_lightbar_off"):
                    if key in a:
                        p[key] = a[key]
                p["battery_brightness"] = max(0, min(100, int(p["battery_brightness"])))
                self.save_config()
                self.update_idle()
                self.apply_keyboard()
                return {}
            if cmd == "set_predator_key":
                if a["action"] not in ("open", "cycle_profile", "none"):
                    raise ValueError("unknown action")
                self.config["predator_key"] = a["action"]
                self.save_config()
                return {}
            if cmd == "restore":
                self.restore()
                return {}
            if cmd == "set_profile":
                hw.set_profile(a["profile"])
                # a manual choice becomes the profile for the current power source
                ac = hw.battery().get("ac_online", True)
                self.config["profile"]["ac" if ac else "battery"] = a["profile"]
                self.save_config()
                return {}
            if cmd == "set_profile_policy":
                merge(self.config["profile"], {k: a[k] for k in ("ac", "battery", "auto_switch") if k in a})
                self.save_config()
                if self.config["profile"]["auto_switch"]:
                    self.apply_profile_for_power(hw.battery().get("ac_online", True))
                return {}
            if cmd == "set_fans":
                merge(self.config["fans"], {k: a[k] for k in ("mode", "cpu", "gpu", "curve") if k in a})
                self.save_config()
                self.apply_fans()
                return self._result("fans")
            if cmd == "set_keyboard":
                k = self.config["keyboard"]
                if a.get("effect") == "on":
                    a["effect"] = k.get("last_effect") or "static"
                if a.get("effect") == "off" and k["effect"] != "off":
                    k["last_effect"] = k["effect"]
                for key in ("effect", "color", "brightness", "effect_id", "frame"):
                    if key in a:
                        k[key] = a[key]
                if k["effect"] not in EFFECTS:
                    raise ValueError("unknown effect")
                if a.get("frame") is not None and len(bytes.fromhex(a["frame"])) != NUM_CELLS * 3:
                    raise ValueError("frame must be %d bytes" % (NUM_CELLS * 3))
                if a.get("save", True):
                    self.save_config()
                self.apply_keyboard()
                return self._result("keyboard")
            if cmd == "preview_keyboard":
                # preview without saving (e.g. while painting in the RGB studio)
                k = dict(self.config["keyboard"])
                k.update({x: a[x] for x in ("effect", "color", "brightness", "frame", "effect_id") if x in a})
                self.apply_keyboard(k)
                return self._result("keyboard")
            if cmd == "set_lightbar":
                for key in ("mode", "color", "brightness", "speed", "direction", "sync_keyboard"):
                    if key in a:
                        self.config["lightbar"][key] = a[key]
                if self.config["lightbar"]["mode"] not in hw.LIGHTBAR_MODES:
                    raise ValueError("unknown lightbar mode")
                self.save_config()
                self.apply_lightbar()
                return self._result("lightbar")
            if cmd == "save_preset":
                self.config["keyboard_presets"][a["name"]] = dict(self.config["keyboard"])
                self.save_config()
                return {}
            if cmd == "delete_preset":
                self.config["keyboard_presets"].pop(a["name"], None)
                self.save_config()
                return {}
            if cmd == "load_preset":
                self.config["keyboard"] = dict(self.config["keyboard_presets"][a["name"]])
                self.save_config()
                self.apply_keyboard()
                return self._result("keyboard")
            if cmd == "set_battery_limit":
                hw.set_battery_limit(bool(a["enabled"]))
                self.config["battery"]["limit"] = bool(a["enabled"])
                self.save_config()
                self.refresh_battery_limit()
                return {}
            if cmd == "set_battery_calibration":
                hw.set_battery_calibration(bool(a["enabled"]))
                return {}
            if cmd == "get_firmware":
                out = {}
                for name, (getter, _s, label) in hw.FIRMWARE_OPTIONS.items():
                    try:
                        out[name] = {"label": label, "value": getter()}
                    except Exception as e:
                        out[name] = {"label": label, "value": None, "error": str(e)}
                for name, getter in (("battery_limit", hw.get_battery_limit),
                                     ("battery_calibration", hw.get_battery_calibration)):
                    try:
                        out[name] = {"value": getter()}
                    except Exception as e:
                        out[name] = {"value": None, "error": str(e)}
                return out
            if cmd == "set_firmware":
                name = a["name"]
                if name not in hw.FIRMWARE_OPTIONS:
                    raise ValueError("unknown option")
                hw.FIRMWARE_OPTIONS[name][1](a["value"])
                return {}
        raise ValueError("unknown command: %s" % cmd)

    def _result(self, key):
        if key in self.errors:
            raise RuntimeError(self.errors[key])
        return {}


class CpuLoadProxy:
    def __init__(self):
        self._c = hw.CpuLoad()
        self.value = None

    def sample(self):
        v = self._c.sample()
        if v is not None:
            self.value = v


def write_state(state):
    tmp = STATE_PATH + ".tmp"
    with open(tmp, "w") as f:
        json.dump(state, f)
    os.chmod(tmp, 0o644)
    os.replace(tmp, STATE_PATH)


def peer_allowed(sock, group):
    creds = sock.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i"))
    _pid, uid, gid = struct.unpack("3i", creds)
    if uid == 0:
        return True
    try:
        g = grp.getgrnam(group)
    except KeyError:
        return False
    if gid == g.gr_gid:
        return True
    try:
        user = pwd.getpwuid(uid).pw_name
    except KeyError:
        return False
    return user in g.gr_mem or pwd.getpwuid(uid).pw_gid == g.gr_gid


def make_handler(daemon):
    class Handler(socketserver.StreamRequestHandler):
        def handle(self):
            if not peer_allowed(self.request, daemon.config.get("allowed_group", "wheel")):
                self.wfile.write(b'{"ok": false, "error": "permission denied (user not in allowed group)"}\n')
                return
            for line in self.rfile:
                try:
                    req = json.loads(line)
                    res = {"ok": True, "result": daemon.handle(req)}
                except Exception as e:
                    res = {"ok": False, "error": str(e)}
                self.wfile.write((json.dumps(res) + "\n").encode())
                self.wfile.flush()
    return Handler


class Server(socketserver.ThreadingMixIn, socketserver.UnixStreamServer):
    daemon_threads = True


def main():
    if os.geteuid() != 0:
        sys.exit("predatord must run as root")
    os.makedirs(os.path.dirname(SOCKET_PATH), exist_ok=True)
    os.chmod(os.path.dirname(SOCKET_PATH), 0o755)
    try:
        os.unlink(SOCKET_PATH)
    except FileNotFoundError:
        pass

    d = Daemon()
    d.kbd_present = None
    log("predatord %s started" % __version__)
    d.restore()
    d.kbd_present = d.kbd.available()

    server = Server(SOCKET_PATH, make_handler(d))
    os.chmod(SOCKET_PATH, 0o666)  # access is checked via SO_PEERCRED

    signal.signal(signal.SIGTERM, lambda *_: (server.server_close(), sys.exit(0)))
    signal.signal(signal.SIGUSR1, lambda *_: threading.Thread(target=lambda: (d.lock.acquire(), d.restore(), d.lock.release())).start())
    threading.Thread(target=d.loop, daemon=True).start()
    HotkeyListener(d.predator_key, log).start()
    d.idle.start()
    server.serve_forever()


if __name__ == "__main__":
    main()