#!/usr/bin/env python

"""
Utility to generate the workspace buttons for eww bar widget.

Currently it fetches information from monitor 0. If all monitors have
synchronised workspaces this should not be a problem.
"""
import os, sys, socket, json, subprocess
from typing import Optional

WORKSPACE_ICONS = [
    (1, "☱"),
    (2, "☲"),
    (3, "☳"),
    (4, "☴"),
    (5, "☵"),
    (6, "☶"),
    (7, "☷"),
    (8, "☰"),
]
EXTRA_WORKSPACE_ICONS = [
    (0, "⚌"),
    (1, "⚍"),
    (2, "⚎"),
    (3, "⚏"),
]

Workspaces = dict[int, bool]

def get_widgets(workspaces: Workspaces) -> str:
    """
    Create widget sexp from workspace information
    """
    def get_class(k):
        key = workspaces.get(k, None)
        if key is None:
            return "inactive"
        elif key:
            return "focused"
        else:
            return "active"

    buttons = [
        f'"id": "{k}", "text": "{icon}", "class": "{get_class(k)}"'
        for k, icon in WORKSPACE_ICONS
    ]
    buttons = ", ".join(["{" + b + "}" for b in buttons])
    return "[" + buttons + "]"


# Compositor specific information

def hypr_get_workspace_info() -> Workspaces:
    workspaces = json.loads(os.popen("hyprctl workspaces -j").read())
    monitors = json.loads(os.popen("hyprctl monitors -j").read())
    if not monitors:
        print("No monitor found!", file=sys.stderr)
        return []

    monitor, active = [
        (m["name"], m["activeWorkspace"]["id"]) for m in monitors if m["id"] == 0
    ][0]
    workspaces = {
        w["id"]: w["id"] == active
        for w in workspaces if w["monitor"] == monitor
    }
    return workspaces


def hypr_listen(socket_addr: str):
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        sock.connect(socket_addr)
    except:
        print(f"Could not connect to {socket_addr}", file=sys.stderr)
        raise

    # The flush here is very important since python by default buffers!
    workspace_info = hypr_get_workspace_info()
    print(get_widgets(workspace_info), flush=True)
    while sock:
        data = sock.recv(1024)
        if b"workspace" in data:
            workspace_info = hypr_get_workspace_info()
            print(get_widgets(workspace_info), flush=True)


def sway_get_workspace_info() -> Workspaces:
    workspaces = json.loads(os.popen("swaymsg --raw -t get_workspaces").read())
    # REVIEW: This has not been tested on a multi-monitor setup.
    workspaces = {
        int(w["num"]): w["focused"]
        for w in workspaces
    }
    return workspaces

def sway_listen():
    proc = subprocess.Popen(
        ["swaymsg", "-t", "subscribe", "-m", '["workspace"]'], stdout=subprocess.PIPE
    )
    workspace_info = sway_get_workspace_info()
    print(get_widgets(workspace_info), flush=True)
    while line := proc.stdout.readline():
        # Not needed
        info = json.loads(line)
        workspace_info = sway_get_workspace_info()
        print(get_widgets(workspace_info), flush=True)


def detect():
    """
    Main entry point, detects the type of compositor used
    """
    if signature := os.environ.get("HYPRLAND_INSTANCE_SIGNATURE", None):
        socket_addr = f"/tmp/hypr/{signature}/.socket2.sock"
        return hypr_listen(socket_addr)
    if signature := os.environ.get("SWAYSOCK", None):
        return sway_listen()

    print("No compositor found.", file=sys.stderr)
    return None


if __name__ == "__main__":
    detect()