#!/usr/bin/env python3
import math
import os
import select
import socket
import struct
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import asrock_hid as h

SOCK_PATH = os.environ.get(
    "ASROCK_SOCK", os.path.join(os.environ.get("XDG_RUNTIME_DIR", "/tmp"), "asrock-led.sock")
)
FRAME = 0.05

NATIVE_TAKES_COLOR = {
    "static": (h.A_STATIC, True),
    "breath": (h.A_BREATH, True),
    "strobe": (h.A_STROBE, True),
    "direct": (h.A_DIRECT, True),
    "cycle": (h.A_CYCLE, False),
    "random": (h.A_RANDOM, False),
    "music": (h.A_MUSIC, False),
    "wave": (h.A_WAVE, False),
    "spring": (h.A_SPRING, False),
    "stack": (h.A_STACK, False),
    "cram": (h.A_CRAM, False),
    "scan": (h.A_SCAN, False),
    "neon": (h.A_NEON, False),
    "water": (h.A_WATER, False),
    "rainbow": (h.A_RAINBOW, False),
}

COLORS = {
    "black": "000000", "off": "000000", "white": "ffffff",
    "red": "ff0000", "green": "00ff00", "blue": "0000ff",
    "yellow": "ffff00", "cyan": "00ffff", "magenta": "ff00ff",
    "orange": "ff8000", "purple": "8000ff", "pink": "ff0080",
    "amber": "ffbf00",
}


def parse_color(s):
    if s in COLORS:
        s = COLORS[s]
    s = s.lstrip("#")
    if len(s) == 6 and all(c in "0123456789abcdefABCDEF" for c in s):
        return struct.unpack("BBB", bytes.fromhex(s))
    parts = s.replace(",", " ").split()
    if len(parts) == 3:
        vals = [int(p, 0) for p in parts]
        if all(0 <= v <= 255 for v in vals):
            return tuple(vals)
    raise ValueError("bad color: %s" % s)


def pulse_frames(rgb, count, period):
    n = max(2, int(round(period / FRAME)))
    frames = []
    for _ in range(count):
        for i in range(n):
            br = math.sin(math.pi * i / (n - 1))
            frames.append((int(rgb[0] * br), int(rgb[1] * br), int(rgb[2] * br)))
    return frames


def blink_frames(rgb, count, period):
    half = max(1, int(round(period / 2 / FRAME)))
    frames = []
    for _ in range(count):
        frames.extend([rgb] * half)
        frames.extend([(0, 0, 0)] * half)
    return frames


class Daemon:
    def __init__(self, dev=h.DEFAULT_DEV):
        self.fd = h.open_dev(dev)
        self.zone = 0
        self.allzone = 1
        self.idle = {"kind": "native", "mode": h.A_STATIC, "rgb": (0, 0, 255), "speed": 0xE0, "written": True}
        self.current = {"kind": "native", "mode": h.A_STATIC, "rgb": (0, 0, 255), "speed": 0xE0, "written": False}

    def set_zone(self, zone):
        if zone >= 0:
            self.zone = zone
            self.allzone = 0
        else:
            self.zone = 0
            self.allzone = 1

    def write(self, mode, rgb, speed):
        r, g, b = rgb
        return h.write_set(self.fd, self.zone, mode, r, g, b, speed, self.allzone)

    def native(self, mode, rgb, speed):
        self.current = {"kind": "native", "mode": mode, "rgb": rgb, "speed": speed, "written": False}
        return "ok"

    def frames(self, frames):
        self.current = {"kind": "frames", "frames": frames, "idx": 0}
        return "ok"

    def apply(self, line):
        toks = line.split()
        if not toks:
            return "err empty"
        if toks[0] == "zone":
            self.set_zone(-1 if toks[1] == "all" else int(toks[1], 0))
            toks = toks[2:]
        if not toks:
            return "ok"
        effect = toks[0]
        args = toks[1:]

        if effect == "idle":
            self.idle = {"kind": "native", "mode": h.A_STATIC, "rgb": parse_color(args[0]),
                         "speed": int(args[1], 0) if len(args) > 1 else 0xE0, "written": False}
            self.current = dict(self.idle, written=False)
            return "ok idle set"

        if effect == "off":
            return self.native(h.A_OFF, (0, 0, 0), 0)

        if effect == "static":
            rgb = parse_color(args[0])
            speed = int(args[1], 0) if len(args) > 1 else 0xE0
            return self.native(h.A_STATIC, rgb, speed)

        if effect == "pulse":
            rgb = parse_color(args[0])
            count = int(args[1], 0) if len(args) > 1 else 1
            period = float(args[2]) if len(args) > 2 else 0.8
            return self.frames(pulse_frames(rgb, count, period))

        if effect == "blink" or effect == "flash":
            rgb = parse_color(args[0])
            count = int(args[1], 0) if len(args) > 1 else 1
            period = float(args[2]) if len(args) > 2 else (0.2 if effect == "flash" else 0.4)
            return self.frames(blink_frames(rgb, count, period))

        if effect in NATIVE_TAKES_COLOR:
            mode, takes_color = NATIVE_TAKES_COLOR[effect]
            speed = 0xE0
            rgb = (255, 255, 255)
            n = 0
            if args:
                try:
                    speed = int(args[0], 0)
                    n = 1
                except ValueError:
                    pass
            if n < len(args) and takes_color:
                rgb = parse_color(args[n])
                n += 1
            elif takes_color and len(args) == 1:
                try:
                    rgb = parse_color(args[0])
                except ValueError:
                    pass
            return self.native(mode, rgb, speed)

        if effect == "query":
            st = h.read_zone(self.fd, self.zone)
            if not st:
                return "err no reply"
            return "ok z%d mode=0x%02x rgb=(%d,%d,%d) speed=%d all=%d" % (
                st["zone"], st["mode"], st["r"], st["g"], st["b"], st["speed"], st["all"])

        if effect == "state":
            c = self.current
            if c["kind"] == "native":
                desc = "native mode=0x%02x rgb=%s speed=%d" % (c["mode"], c["rgb"], c["speed"])
            else:
                nxt = c["frames"][c["idx"]] if c["idx"] < len(c["frames"]) else None
                desc = "frames %d/%d next=%s" % (c["idx"], len(c["frames"]), nxt)
            return "ok zone=%d all=%d | current=%s | idle=%s" % (
                self.zone, self.allzone, desc, self.idle)

        if effect == "commit":
            h.commit(self.fd)
            return "ok committed"

        if effect == "help":
            return ("ok static|off|idle <c> | pulse <c> [n] [s] | blink <c> [n] [s] | flash <c> | "
                    + "|".join(NATIVE_TAKES_COLOR) + " [speed] [color] | query | commit | zone <n> <cmd>")

        raise ValueError("unknown effect: %s (try 'help')" % effect)

    def render(self):
        cur = self.current
        if cur is None:
            return
        if cur["kind"] == "native":
            if not cur["written"]:
                self.write(cur["mode"], cur["rgb"], cur["speed"])
                cur["written"] = True
        elif cur["kind"] == "frames":
            if cur["idx"] < len(cur["frames"]):
                self.write(h.A_STATIC, cur["frames"][cur["idx"]], 0xE0)
                cur["idx"] += 1
            else:
                self.current = dict(self.idle, written=False)

    def run(self):
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        try:
            os.unlink(SOCK_PATH)
        except FileNotFoundError:
            pass
        sock.bind(SOCK_PATH)
        os.chmod(SOCK_PATH, 0o600)
        sock.listen(16)
        sock.settimeout(FRAME)
        next_tick = time.monotonic()
        while True:
            try:
                conn, _ = sock.accept()
            except socket.timeout:
                conn = None
            if conn is not None:
                with conn:
                    data = b""
                    while b"\n" not in data:
                        chunk = conn.recv(1024)
                        if not chunk:
                            break
                        data += chunk
                    try:
                        ack = self.apply(data.decode().strip())
                    except Exception as e:
                        ack = "err %s" % e
                    conn.sendall(ack.encode() + b"\n")
            now = time.monotonic()
            if now >= next_tick:
                self.render()
                next_tick = now + FRAME


def ctl(args):
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.settimeout(1.5)
    sock.connect(SOCK_PATH)
    sock.sendall(" ".join(args).encode() + b"\n")
    data = b""
    while b"\n" not in data:
        chunk = sock.recv(1024)
        if not chunk:
            break
        data += chunk
    print(data.decode().strip())


def main():
    args = sys.argv[1:]
    if not args or args[0] == "daemon":
        Daemon().run()
    elif args[0] == "ctl":
        ctl(args[1:])
    else:
        sys.exit("usage: asrock-indicator.py daemon | ctl <command>")


if __name__ == "__main__":
    main()