#!/usr/bin/env python3
import subprocess, json, os, socket
import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk

OUTPUT = None  # set via argv, e.g. "DP-2"

def get_icon(cls: str) -> str | None:
    theme = Gtk.IconTheme.get_default()
    for name in [cls, cls.lower(), cls.lower().rstrip("0123456789")]:
        info = theme.lookup_icon(name, 16, 0)
        if info:
            return info.get_filename()
    return None
# Define which workspace IDs belong to each monitor
MONITOR_WORKSPACES = {
    "DP-2":     [1, 2, 3, 4, 5],
    "HDMI-A-1": [6, 7, 8, 9, 10],
}

def build():
    clients    = json.loads(subprocess.check_output(["hyprctl", "clients",        "-j"]))
    workspaces = json.loads(subprocess.check_output(["hyprctl", "workspaces",     "-j"]))
    active     = json.loads(subprocess.check_output(["hyprctl", "activeworkspace", "-j"]))

    # existing workspace IDs on this monitor (Hyprland only lists non-empty ones)
    existing = {w["id"] for w in workspaces if w["monitor"] == OUTPUT}

    # all IDs we want to show, including empty ones
    all_ids = sorted(set(MONITOR_WORKSPACES.get(OUTPUT, [])) | existing)

    # build icon map
    amount = {}
    icons: dict[int, list[str]] = {}
    for c in clients:
        wid  = c["workspace"]["id"]
        cls  = c.get("class") or c.get("initialClass", "")
        path = get_icon(cls)
        if path and path not in icons.get(wid, []):
            icons.setdefault(wid, []).append(path)
            amount[c["workspace"]["id"]] =+ 1

    result = []
    for wid in all_ids:
        result.append({
            "id":     wid,
            "label":  str(wid),
            "active": wid == active["id"],
            "urgent": any(c.get("urgent") and c["workspace"]["id"] == wid for c in clients),
            "output": OUTPUT,
            "icons":  icons.get(wid, []),
            "pxwidth":  32 + len(icons.get(wid, [])) * 20,
        })

    print(json.dumps(result), flush=True)

def watch():
    xdg = os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}")
    sig = os.environ.get("HYPRLAND_INSTANCE_SIGNATURE", "")
    sock = f"{xdg}/hypr/{sig}/.socket2.sock"

    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
        s.connect(sock)
        buf = ""
        while True:
            buf += s.recv(4096).decode()
            while "\n" in buf:
                line, buf = buf.split("\n", 1)
                ev = line.split(">>")[0]
                if ev in {"workspace", "openwindow", "closewindow",
                          "movewindow", "urgent", "focusedmon", "activelayout"}:
                    build()

if __name__ == "__main__":
    import sys
    OUTPUT = sys.argv[1] if len(sys.argv) > 1 else None
    build()
    watch()