#!/usr/bin/env python3

# pyright: strict, reportAny=false, reportExplicitAny=false, reportUnusedCallResult=false

from json import loads
from logging import NOTSET, basicConfig, debug, info
from os import getenv
from subprocess import CalledProcessError, check_output
from sys import argv
from typing import Any, Literal, TypeAlias, cast, override

HOME = getenv("HOME")
XDG_CONFIG_HOME = getenv("XDG_CONFIG_HOME")
TARGET_DISPLAY_WIDTH = 1920
TARGET_DISPLAY_HEIGHT = 1080


class Window:
    def __init__(self, data: dict[str, Any]):
        self.id: int = data["id"]
        self.app: str = data["app"]
        self.space: str = data["space"]
        self.title: str = data["title"]
        self.has_focus: bool = data["has-focus"]
        self.frame_w: int = data["frame"]["w"]
        self.frame_h: int = data["frame"]["h"]

    @property
    def size(self) -> int:
        return self.frame_w * self.frame_h

    @override
    def __repr__(self) -> str:
        return (
            f"Window({self.title} | {self.app}"
            f", {self.id}"
            f"{', Focused' if self.has_focus else ''})"
        )

    def __gt__(self, other: "Window"):
        return self.id.__gt__(other.id)


class Space:
    def __init__(self, data: dict[str, Any]):
        self.id: int = data["id"]
        self.index: int = data["index"]
        self.label: str = data["label"]
        self.display: int = data["display"]
        self.windows: list[int] = data["label"]
        self.is_visible: str = data["is-visible"]
        self.has_focus: bool = data["has-focus"]
        self.is_native_fullscreen: bool = data["is-native-fullscreen"]

    @override
    def __repr__(self) -> str:
        return (
            f"Space({self.label if self.label and len(self.label) > 0 else '<NoLabel>'}"
            f", {self.index}"
            f"{', Fullscreen' if self.is_native_fullscreen else ''}"
            f"{', Focused' if self.has_focus else ''})"
        )

    def __gt__(self, other: "Space"):
        return self.index.__gt__(other.index)


class YabaiDisplay:
    def __init__(self, data: dict[str, Any]):
        self.id: int = data["id"]
        self.uuid: str = data["uuid"]
        self.index: int = data["index"]
        self.label: str = data["label"]
        self.has_focus: bool = data["has-focus"]
        self.spaces: list[str] = data["spaces"]
        self.frame: dict[str, int] = data["frame"]

    @override
    def __repr__(self) -> str:
        return (
            f"Display({self.label if self.label and len(self.label) > 0 else '<NoLabel>'}"
            f", {self.index}"
            f", Frame ({self.frame})"
            f"{', Focused' if self.has_focus else ''})"
        )

    def __gt__(self, other: "YabaiDisplay"):
        return self.frame["w"].__gt__(other.frame["w"])


SpaceSel: TypeAlias = int | str


class CLIWrapper:
    _base_args: list[str] = []

    def message(self, args: list[str | int]) -> str:
        return self.execute(self._base_args + [str(arg) for arg in args]).decode("utf8")

    def execute(self, *args: Any, **kwargs: Any) -> Any:
        debug(f"Executing: ({args}), ({kwargs})")
        return cast(str, check_output(*args, **kwargs))


class Yabai(CLIWrapper):
    _base_args: list[str] = ["yabai", "-m"]

    spaces: list[tuple[str, bool, list[str], bool]] = [
        ("Desktop", False, ["*"], True),
        ("Finder", False, ["Finder"], False),
        ("Terminal", True, ["Alacritty"], False),
        ("Browser", True, ["Firefox"], False),
    ]
    _initial_window: Window | None = None
    _dual_display: None | Literal[True] | Literal[False] = None
    _exit_with_rule_apply: bool = False

    def __init__(self):
        self._dual_display = len(self.get_displays()) > 1
        if self._dual_display:
            self.spaces.append(
                (
                    "Communication",
                    False,
                    ["Slack", "Signal", "Spotify"],
                    True,
                )
            )
            for application in [
                "Google Meet",
                "Obsidian",
                "Asana",
                "Notion",
                "Netflix",
                "YouTube",
            ]:
                self.spaces.append(
                    (
                        application,
                        True,
                        [application],
                        True,
                    )
                )
        else:
            self.spaces.append(
                (
                    "Communication",
                    False,
                    ["Slack", "Signal", "Spotify", "Notion"],
                    False,
                )
            )
            self.spaces.append(("Notetaking", True, ["Obsidian", "Asana"], False))

    def __enter__(self):
        self._initial_window = self.get_focused_window()
        return self

    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
        if exc_type is not None:
            debug(f"Exited with {exc_type} {exc_value}")
        if self._exit_with_rule_apply:
            self.message(["rule", "--apply"])
            if self._initial_window is not None:
                self.message(["window", "--focus", self._initial_window.id])
        if exc_type is None:
            debug(f"Executed successfully")

    def get_windows(self) -> set[Window]:
        return {
            Window(window) for window in loads(self.message(["query", "--windows"]))
        }

    def get_spaces(self) -> set[Space]:
        return {Space(space) for space in loads(self.message(["query", "--spaces"]))}

    def get_displays(self) -> set[YabaiDisplay]:
        return {
            YabaiDisplay(display)
            for display in loads(self.message(["query", "--displays"]))
        }

    def get_main_display(
        self, displays: set[YabaiDisplay] | None = None
    ) -> YabaiDisplay:
        return sorted(list(displays if displays != None else self.get_displays()))[-1]

    def is_blank_space(self, space: Space) -> bool:
        return (
            space.label not in {s[0] for s in self.spaces}
            and not space.is_native_fullscreen
        )

    def manage_displays(self):
        displays = self.get_displays()
        main_display = self.get_main_display(displays)
        for display in displays:
            if display.index == main_display.index:
                self.message(["display", display.index, "--label", "Main"])
            elif display.index == main_display.index + 1:
                self.message(["display", display.index, "--label", f"Secondary"])
            else:
                self.message(
                    [
                        "display",
                        display.index,
                        "--label",
                        f"YabaiDisplay {display.index}",
                    ]
                )

    def manage_spaces(self):
        initial_window = self.get_focused_window()
        # Start by making sure that the expected number of spaces are present
        spaces = self.get_spaces()
        occupied_spaces = {
            space
            for space in spaces
            if not self.is_blank_space(space) and not space.is_native_fullscreen
        }
        blank_spaces = {space for space in spaces if self.is_blank_space(space)}
        for _ in range(
            max(0, len(self.spaces) - (len(occupied_spaces) + len(blank_spaces)))
        ):
            self.message(["space", "--create"])
        # Use blank spaces to create occupied spaces as necessary
        spaces = self.get_spaces()
        blank_spaces = {space for space in spaces if self.is_blank_space(space)}
        created_space_labels: set[tuple[str, bool]] = set()
        for space_label, space_fullscreen, _, _ in self.spaces:
            if any(space.label == space_label for space in spaces):
                continue
            space = blank_spaces.pop()
            self.message(["space", space.index, "--label", space_label])
            created_space_labels.add((space_label, space_fullscreen))
        # Remove unnecessary spaces
        spaces = self.get_spaces()
        blank_spaces = [space for space in spaces if self.is_blank_space(space)]
        blank_spaces.sort(key=lambda s: s.index, reverse=True)
        # Make sure that the focused space isn't a blank one
        if any(space.has_focus for space in blank_spaces):
            self.message(["space", "--focus", self.spaces[0][0]])
        for space in blank_spaces:
            self.message(["space", "--destroy", space.index])
        # Configure the new spaces
        main_display = self.get_main_display()
        for label, fullscreen in created_space_labels:
            self.set_space_background(label)
        for label, fullscreen, _, _ in self.spaces:
            self.set_config(
                label,
                fullscreen=fullscreen,
                horizontal_padding=int(
                    (main_display.frame["w"] - TARGET_DISPLAY_WIDTH) / 2
                ),
                vertical_padding=int(
                    (main_display.frame["h"] - TARGET_DISPLAY_HEIGHT) / 2
                ),
            )
        if len(created_space_labels) > 0 and initial_window is not None:
            self.message(["window", "--focus", initial_window.id])
        # Return focus
        if initial_window is not None:
            self.message(["window", "--focus", initial_window.id])
        info(f"Spaces configured: {sorted(self.get_spaces())}")

    def set_space_background(self, space: SpaceSel):
        try:
            self.message(["space", "--focus", space])
        except CalledProcessError:
            # Almost certainly thrown because space is already focused, so no problem
            pass
        self.execute(
            [
                "osascript",
                "-e",
                'tell application "System Events" to tell every desktop to set picture to "/System/Library/Desktop Pictures/Solid Colors/Black.png"',
            ]
        )

    def set_config(
        self,
        space: SpaceSel,
        fullscreen: bool = False,
        gap: int = 10,
        vertical_padding: int = 0,
        horizontal_padding: int = 0,
    ):
        for config in [
            ["window_shadow", "float"],
            ["window_opacity", "on"],
            ["layout", "bsp"],
            ["top_padding", int(max(0 if fullscreen else gap, vertical_padding))],
            ["bottom_padding", int(max(0 if fullscreen else gap, vertical_padding))],
            ["left_padding", int(max(0 if fullscreen else gap, horizontal_padding))],
            ["right_padding", int(max(0 if fullscreen else gap, horizontal_padding))],
            ["window_gap", gap],
        ]:
            self.message(["config", "--space", space] + config)

    def set_rules_and_signals(self):
        # Reset rules and signals
        for domain in ["rule", "signal"]:
            for _ in range(len(loads(self.message([domain, "--list"])))):
                self.message([domain, "--remove", 0])
        # Load the system agent on dock restart
        self.message(
            [
                "signal",
                "--add",
                "label=SystemAgentReloadSignal",
                "event=dock_did_restart",
                "action=sudo yabai --load-sa",
            ]
        )
        # Reload spaces when displays are reset
        for reload_event in ["display_added", "display_removed"]:
            self.message(
                [
                    "signal",
                    "--add",
                    f"label={reload_event}RestartSignal",
                    f"event={reload_event}",
                    f"action=/bin/zsh {XDG_CONFIG_HOME}/yabai/yabairc",
                ]
            )
        # Normal windows should be put on the desktop
        self.message(
            [
                "rule",
                "--add",
                "label=DefaultDesktopRule",
                "subrole=AXStandardWindow",
                "space=^Desktop",
            ]
        )
        # Rules for applications that get their own spaces
        for label, _, apps, _ in self.spaces:
            for app in apps:
                if app == "*":
                    continue
                self.message(
                    [
                        "rule",
                        "--add",
                        f"label={app}{label}Rule",
                        f"app={app}",
                        f"space=^{label}",
                    ]
                )

        # Compile SurfingKeys configuration when Firefox is launched
        self.message(
            [
                "signal",
                "--add",
                "event=application_launched",
                "app=Firefox",
                "label=FirefoxCompileExtensionsSignal",
                f"action=/bin/zsh {XDG_CONFIG_HOME}/surfingkeys/compile.sh",
            ]
        )
        # Run Slack client script (populating cached values) when Slack is launched
        self.message(
            [
                "signal",
                "--add",
                "event=application_launched",
                "app=Slack",
                "label=SlackRunSlackClientScript",
                f"action=/usr/bin/env python3 {HOME}/.scripts/slack_client.py",
            ]
        )
        # "Attach" OBS to Google Meet
        self.message(
            [
                "signal",
                "--add",
                "event=application_launched",
                "app=Google Meet",
                "label=GoogleMeetOBSJoiner",
                f"action=/usr/bin/env python3 {HOME}/.scripts/obs_client.py Activate Default StartVirtualCam",
            ]
        )
        self.message(
            [
                "signal",
                "--add",
                "event=application_terminated",
                "app=Google Meet",
                "label=GoogleMeetOBSJoinerEnd",
                f"action=/usr/bin/env python3 {HOME}/.scripts/obs_client.py StopVirtualCam Activate Disabled",
            ]
        )
        # Check if dark mode settings have been updated when focusing terminal
        self.message(
            [
                "signal",
                "--add",
                "event=window_focused",
                "app=Alacritty",
                "label=AlacrittyCheckDarkMode",
                f"action=/bin/zsh {HOME}/.scripts/lightmode.zsh",
            ]
        )

        for trigger in [
            "display_added",
            "display_removed",
            "display_resized",
            "system_woke",
        ]:
            self.message(
                [
                    "signal",
                    "--add",
                    f"event={trigger}",
                    f"label=DisplayChange{trigger}",
                    f"action=/usr/bin/env python3 {HOME}/.config/yabai.py manage",
                ]
            )

        # Rules that differ for one or multiple displays
        if self._dual_display:
            # space
            self.message(
                [
                    "signal",
                    "--add",
                    "event=space_changed",
                    "label=SpaceChanged",
                    f"action=/usr/bin/env python3 {HOME}/.config/yabai.py move",
                ]
            )
        else:
            # Google Meet and Slack Huddles should be "sticky"
            for app, title in (("Google Meet", ".*"), ("Slack", "Huddle.*")):
                self.message(
                    [
                        "rule",
                        "--add",
                        f"label={app}VideoCallFloatingWindowRule",
                        f"app={app}",
                        f"title={title}",
                        "sticky=on",
                        "manage=on",
                        "opacity=0.9",
                        "grid=10:10:6:6:4:4",
                    ]
                )
            # Tiny streaming player
            self.message(
                [
                    "rule",
                    "--add",
                    f"label=NetflixFloatingWindowRule",
                    f"app=Netflix",
                    "sticky=on",
                    "manage=on",
                ]
            )

    def move_spaces(self):
        initial_window = self.get_focused_window()
        final_window = self.get_focused_window()

        displays = self.get_displays()
        main_display = self.get_main_display(displays)
        secondary_display = next(
            (d for d in displays if not d.id == main_display.id), None
        )

        current_spaces = self.get_spaces()

        incorrect_main_spaces: set[Space] = set()
        incorrect_secondary_spaces: set[Space] = set()
        for space_label, _, _, is_secondary_space in self.spaces:
            mapped_space = next(
                (s for s in current_spaces if s.label == space_label), None
            )
            if mapped_space is None:
                continue

            if (
                mapped_space.display == main_display.index
                and is_secondary_space
                and secondary_display
            ):
                incorrect_secondary_spaces.add(mapped_space)
            if mapped_space.display != main_display.index and not is_secondary_space:
                incorrect_main_spaces.add(mapped_space)
        while len(incorrect_main_spaces) > 0 and len(incorrect_secondary_spaces) > 0:
            self.message(
                [
                    "space",
                    incorrect_main_spaces.pop().label,
                    "--switch",
                    incorrect_secondary_spaces.pop().label,
                ]
            )
        for spaces, secondary in (
            (incorrect_main_spaces, False),
            (incorrect_secondary_spaces, True),
        ):
            for space in spaces:
                self.message(
                    [
                        "space",
                        space.label,
                        "--display",
                        secondary_display.label
                        if secondary and secondary_display
                        else main_display.label,
                    ]
                )

        if (
            final_window is not None
            and initial_window is not None
            and final_window.id != initial_window.id
        ):
            self.message(["window", "--focus", initial_window.id])

    def get_focused_window(self) -> Window | None:
        windows = [
            window
            for window in [
                Window(window) for window in loads(self.message(["query", "--windows"]))
            ]
            if window.has_focus
        ]
        if len(windows) > 0:
            return windows.pop()
        return None

    def enable_exit_with_rule_apply(self):
        self._exit_with_rule_apply = True


if __name__ == "__main__":
    basicConfig(level=NOTSET)
    debug(f"Called with parameters {argv}")
    with Yabai() as yabai:
        if argv[1] == "manage" or argv[1] == "initialize":
            yabai.enable_exit_with_rule_apply()
            yabai.manage_displays()
            yabai.manage_spaces()
        if argv[1] == "initialize":
            yabai.set_rules_and_signals()
        if argv[1] == "move" or argv[1] == "manage" or argv[1] == "initialize":
            yabai.move_spaces()