#! /usr/bin/python3

import json
import re
import subprocess
from collections.abc import Iterable, Iterator
from typing import Any


def urgent_descendents(node: dict[str, Any]) -> Iterator[dict[str, Any]]:
    for subnode in node["nodes"]:
        yield from find_urgent(subnode)


def find_urgent(node: dict[str, Any]) -> Iterator[dict[str, Any]]:
    if node["type"] in ("root", "output"):
        yield from urgent_descendents(node)
    elif node["type"] == "workspace":
        for unode in urgent_descendents(node):
            yield unode | {"workspace": node["name"]}
    elif node["type"] == "con":
        if node["focus"]:
            yield from urgent_descendents(node)
        elif node["urgent"]:
            yield {"id": node["id"], "name": node["name"]}


def display_option(unode: dict[str, Any]) -> str:
    return f"[#{unode['id']} @ WS {unode['workspace']}] {unode['name']}"


def select_option(unodes: Iterable[dict[str, Any]]) -> str:
    joined = "".join(display_option(unode) + "\n" for unode in unodes)
    proc = subprocess.run(
        ["fuzzel", "--dmenu"],
        input=joined,
        capture_output=True,
        text=True,
    )

    return proc.stdout


def parse_id_from_option(opt: str) -> int:
    m = re.match("^\\[#(?P<id>[0-9]+) ", opt)
    if m is None:
        raise ValueError()

    return int(m["id"])


def main() -> None:
    proc = subprocess.run(["swaymsg", "-t", "get_tree"], capture_output=True, text=True)
    sway_tree = json.loads(proc.stdout)
    opt = select_option(find_urgent(sway_tree))
    con_id = parse_id_from_option(opt)
    subprocess.run(["swaymsg", "--raw", f"[con_id={con_id}]", "focus"])


if __name__ == "__main__":
    main()