# A script to generate keybindings for i3 Window Manager modes with consistent behavior and help text.
from __future__ import annotations

import io
from contextlib import redirect_stdout
from dataclasses import dataclass, field
from pathlib import Path
from textwrap import indent
from typing import List, Optional, Dict, Tuple


@dataclass
class Binding:
    """An abstract superclass for both mode and command bindings."""

    name: str
    key: str = field(default="")

    # Execute the binding only on key release
    release: bool = field(default=False)

    # A hidden binding won't show up in mode descriptions. Useful if you want to have
    # a second version of a binding with a modifier key, but don't need it to show up separately."""
    hidden: bool = field(default=False)
    # Override the binding's default hint text.
    hint: str = field(init=False, default="")

    def __post_init__(self):
        if not self.key:
            self.key = self.name.lower()[0]

    def binding(self, in_mode: Optional[str] = None) -> str:
        pass

    def _bind_prefix(self, key: Optional[str] = None) -> str:
        prefix = "bindsym "
        if self.release:
            prefix += "--release "
        prefix += key or self.key
        return prefix

    def mode_hint(self) -> str:
        """How to display this binding in the help text of a Mode."""
        if self.hint:
            return self.hint

        if len(self.key) == 1 and self.key in self.name:
            return self.name.replace(self.key, f"<b>{self.key}</b>", 1)

        return f"[{self.key}]{self.name}"


@dataclass
class Mode(Binding):
    """A mode which can contain a list of sub-bindings, which may be a mix of commands and modes.
    exit_bindings contains a list of binds that will cancel the mode and return to default."""

    bindings: List[Binding] = field(default_factory=list)

    exit_bindings = [
        "Return",
        "Escape",
        "control+g",
        "control+bracketleft",
    ]

    def __post_init__(self):
        super().__post_init__()

        used_keys: Dict[str, Binding] = {}
        for binding in self.bindings:
            # Deconflicting directional bindings is complicated, so let's punt on it for now
            if isinstance(binding, Directional):
                continue
            err_base = (
                f"Binding {binding.key} for {self.name}.{binding.name} conflicts with "
            )
            if binding.key in self.exit_bindings:
                raise ValueError(err_base + "mode escape key.")
            if binding.key in used_keys:
                raise ValueError(err_base + used_keys[binding.key].name)
            used_keys[binding.key] = binding

    def mode_name(self, parent: Optional[str]) -> str:
        """The full name of the mode including its parents names, for disambiguation."""
        if parent:
            return f"{parent}_{self.name}"
        else:
            return self.name

    def mode_var(self, parent: Optional[str]) -> str:
        """The name of the variable representing the mode."""
        return f"$mode_{self.mode_name(parent)}"

    def mode_label(self, parent: Optional[str]) -> str:
        """The mode label to show at the start of the mode help text."""
        lineage = self.mode_name(parent).split("_")

        # Remove the common [space] prefix on nested modes.
        if len(lineage) > 1:
            lineage.remove("space")
        return "".join(f"[{part}]" for part in lineage)

    def mode_def(self, parent: Optional[str]) -> str:
        name = self.mode_name(parent)
        var = self.mode_var(parent)

        binding_names = ", ".join(
            [binding.mode_hint() for binding in self.bindings if not binding.hidden]
        )
        help_text = f"{self.mode_label(parent)}: {binding_names}"

        mode = f"set {var} {help_text}\n" f'mode --pango_markup "{var}" {{\n'

        for binding in self.bindings:
            mode += indent(binding.binding(name) + "\n", "  ")

        for binding in self.exit_bindings:
            mode += indent(f'{self._bind_prefix(binding)} mode "default"\n', "  ")

        mode += "}\n"

        submodes = [binding for binding in self.bindings if isinstance(binding, Mode)]

        for sub in submodes:
            mode += "\n"
            mode += sub.mode_def(name)

        return mode

    def binding(self, parent: Optional[str] = None) -> str:
        return f'{self._bind_prefix()} mode "{self.mode_var(parent)}"'


@dataclass
class Command(Binding):
    command: Optional[str] = field(default=None)
    exit_mode: bool = field(default=True)

    def __post_init__(self):
        super().__post_init__()
        if not self.command:
            self.command = self.name.lower()

    def binding(self, parent: Optional[str] = None) -> str:
        return self._binding(self.key, self.command, parent)

    def _binding(self, key: str, command: str, parent: Optional[str] = None) -> str:
        bind = f"{self._bind_prefix(key)} {command}"
        if parent and self.exit_mode:
            bind += ', mode "default"'
        return bind


@dataclass
class Exec(Command):
    startup_id: bool = field(default=False)

    def binding(self, parent: Optional[str] = None) -> str:
        return self._binding(
            self.key,
            f"exec {'--no-startup-id ' if not self.startup_id else ''}{self.command}",
            parent,
        )


@dataclass
class Directional(Command):
    """A convenience class for directional commands that can be executed upleft/down/up/right with h/j/k/l or the
    arrow keys. Optional parameters include a subset of directions to use, and a modifier for the command.
    The command should include the template variable {direction}."""

    subset: Optional[List[str]] = None
    modifier: Optional[str] = None
    flip: bool = False

    flips = {"left": "right", "right": "left", "down": "up", "up": "down"}

    directions = {
        "left": ["h", "Left"],
        "down": ["j", "Down"],
        "up": ["k", "Up"],
        "right": ["l", "Right"],
    }

    def prefix(self) -> str:
        return self.modifier + "+" if self.modifier else ""

    def used_keys(self) -> Tuple[List[str], List[str]]:
        letter_keys = []
        arrow_keys = []

        for direction in self.directions:
            if not self.subset or direction in self.subset:
                letter, arrow = self.directions[direction]
                letter_keys.append(letter)
                arrow_keys.append(arrow)

        return letter_keys, arrow_keys

    def mode_hint(self) -> str:
        letter_keys, arrow_keys = self.used_keys()

        hint = self.prefix() + "<b>"

        hint += "/".join(letter_keys)
        hint += "/"
        hint += "/".join(arrow_keys)

        hint += f"</b> {self.name}"
        return hint

    def binding(self, parent: Optional[str] = None):
        bind = ""
        for direction, keys in self.directions.items():
            if not self.subset or direction in self.subset:
                if self.flip:
                    direction = self.flips[direction]
                for key in keys:
                    bind += self._binding(self.prefix() + key, self.command).format(
                        direction=direction
                    )
                    bind += "\n"
        return bind


@dataclass
class App:
    """
    Required args:
        name: The name of the app.
    Accepted args:
        key: The key to use. Defaults to the first lower case letter of the name.
        Regular press will try to switch to the app if it's already open.
        Shift press will guarantee opening a new instance. Shift bindings are hidden from mode help text.

        path: The binary to execute. Defaults to the lowercase version of the name.
        args: Args to pass to the app on launch
        class: The X11 Window Class to search for. Defaults to the app name. Case insensitive.
        title: The X11 Window Title to search for.
        switch: If True, switch to the app instead of launching a new instance unless shift is held. Default True.
    """

    name: str
    key: str = field(default="")
    path: str = field(default="")
    args: str = field(default="")
    window_class: str = field(default="")
    window_title: str = field(default="")
    switch: bool = field(default=True)

    def __post_init__(self):
        if not self.key:
            self.key = self.name.lower()[0]
        if not self.path:
            self.path = self.name.lower()
        if not self.window_class:
            self.window_class = self.name

    @property
    def class_query(self) -> str:
        return f'class="(?i){self.window_class}"'

    @property
    def title_query(self) -> str:
        if self.window_title:
            return f' title="(?i){self.window_title}"'
        else:
            return ""

    @property
    def query(self) -> str:
        return f"[{self.class_query}{self.title_query}]"

    def commands(self) -> List[Command]:
        cmd = f"{self.path} {self.args}"
        switch_cmd = f"pgrep {self.path} && i3-msg '{self.query} focus' || {cmd}"

        commands = [
            Exec(self.name, key=self.key, command=switch_cmd if self.switch else cmd),
            Exec(self.name, key=f"Shift+{self.key}", command=cmd, hidden=True),
        ]

        return commands

    @staticmethod
    def find(name: str) -> Optional[App]:
        for app in apps:
            if app.name == name:
                return app
        return None


apps = [
    App("firefox"),
    App("kitty", key="t", args="--single-instance", switch=False),
    App("emacs"),
    App("pycharm"),
    App("idea", key="j"),
    App("discord"),
    App("slack"),
    App("zoom"),
    App(
        "windows",
        path="virt-viewer",
        args="--connect qemu:///system -w Windows10",
        window_class="virt-viewer",
    ),
]

bindings = [
    Exec("Terminal", key="$mod+Return", command="kitty --single-instance"),
    Command("next workspace", key="$mod+Tab", command="workspace next_on_output"),
    Command("prev workspace", key="$mod+Shift+Tab", command="workspace prev_on_output"),
    # Volume control
    Exec(
        "Volume Up",
        key="XF86AudioRaiseVolume",
        command="pactl set-sink-volume @DEFAULT_SINK@ +10%",
        startup_id=False,
    ),
    Exec(
        "Volume Down",
        key="XF86AudioLowerVolume",
        command="pactl set-sink-volume @DEFAULT_SINK@ -10%",
        startup_id=False,
    ),
    Exec(
        "Mute",
        key="XF86AudioMute",
        command="pactl set-sink-mute @DEFAULT_SINK@ toggle",
        startup_id=False,
    ),
    Directional("Focus", command="focus {direction}", modifier="$mod"),
    Directional(
        "Move", command="move {direction}", modifier="$mod+Shift", release=True
    ),
    Mode(
        "space",
        key="$mod+space",
        bindings=[
            Exec("Run", key="space", command="rofi -show run"),
            Mode(
                "workspace",
                key="p",
                bindings=[
                    Mode(
                        "move",
                        bindings=[
                            Directional(
                                "Move",
                                command="exec --no-startup-id i3-msg move workspace to output {direction}",
                                subset=["left", "right"],
                                exit_mode=False,
                            ),
                        ],
                    ),
                    Exec("switch to", key="p", command="i3_switch_workspace.sh"),
                ],
            ),
            Mode(
                "window",
                bindings=[
                    Command("delete", command="kill"),
                    Command("fullscreen", command="fullscreen toggle"),
                    Command("float", key="o", command="floating toggle"),
                    Command("split", command="split h"),
                    Command("vertical split", key="v", command="split v"),
                    Command("parent", command="focus parent", exit_mode=False),
                    Mode(
                        "move",
                        bindings=[
                            Directional(
                                "Move",
                                command="move {direction}",
                                exit_mode=False,
                                release=True,
                            ),
                            Exec(
                                "to workspace",
                                key="p",
                                command="i3_send_to_workspace.sh",
                            ),
                        ],
                    ),
                    Mode(
                        "resize",
                        bindings=[
                            Directional(
                                "Grow",
                                command="resize grow {direction} 10 px or 10 ppt",
                                exit_mode=False,
                            ),
                            Directional(
                                "Shrink",
                                command="resize shrink {direction} 10 px or 10 ppt",
                                modifier="Shift",
                                exit_mode=False,
                                flip=True,
                            ),
                        ],
                    ),
                    Mode(
                        "layout",
                        bindings=[
                            Command("split", command="toggle split"),
                            Command("tabbed", command="layout tabbed"),
                            Command("stacking", key="k", command="layout stacking"),
                        ],
                    ),
                ],
            ),
            Mode("open", bindings=[cmd for app in apps for cmd in app.commands()]),
            Mode(
                "goto",
                bindings=[
                    Exec("server", command="rofi -show ssh"),
                ],
            ),
            Mode(
                "quit",
                bindings=[
                    Exec("logout", key="q", command="i3-msg exit"),
                    Exec("reload", command="i3_reconfigure", key="r"),
                    Command("restart", key="Shift+r"),
                    Mode(
                        "system",
                        bindings=[
                            Exec(
                                "suspend", startup_id=False, command="systemctl suspend"
                            ),
                            Exec(
                                "reboot", startup_id=False, command="systemctl reboot"
                            ),
                            Exec(
                                "power off",
                                startup_id=False,
                                command="systemctl poweroff",
                            ),
                        ],
                    ),
                ],
            ),
        ],
    ),
]


@dataclass
class Output:
    name: str
    pos: str
    mode: str = field(default="2560x1440")
    rate: int = field(default=60)
    primary: bool = field(default=False)

    def xrandr_flags(self) -> str:
        flags = f"--output {self.name} --pos {self.pos} --mode {self.mode} --rate {self.rate}"
        if self.primary:
            flags += " --primary"
        return flags


outputs = [
    Output("DP-0", "0x0"),
    Output("DP-4", "2560x0", rate=144, primary=True),
    Output("HDMI-0", "5120x0"),
]


def xrandr_command() -> str:
    cmd = "exec --no-startup-id xrandr\\\n"
    for output in outputs:
        cmd += indent(output.xrandr_flags() + " \\\n", "    ")
    cmd = cmd[:-3]
    return cmd


@dataclass
class Workspace:
    name: str
    output: str = field(default="primary")
    assigned_apps: List[str] = field(default_factory=list)

    def config(self, index: int) -> str:
        ws_var = f"$ws_{self.name.lower()}"
        cfg = f"set {ws_var} {index}: {self.name}\n"
        cfg += f'workspace "{ws_var}" output {self.output}\n'
        for a in self.assigned_apps:
            app = App.find(a)
            if not app:
                raise ValueError(f"Couldn't find app {a}")
            cfg += f'assign {app.query} "{ws_var}"\n'

        return cfg


workspaces = [
    Workspace("Main", output="DP-4"),
    Workspace("Tasks", output="DP-0", assigned_apps=["windows"]),
    Workspace("Comms", output="HDMI-0", assigned_apps=["discord", "slack"]),
    Workspace("Python", output="primary", assigned_apps=["pycharm"]),
    Workspace("Java", output="primary", assigned_apps=["idea"]),
]


def workspace_config() -> str:
    f = io.StringIO()

    with redirect_stdout(f):
        for idx, workspace in enumerate(workspaces):
            print(workspace.config(idx + 1))

    return f.getvalue()


def binds_config() -> str:
    f = io.StringIO()

    preamble = (Path(__file__).parent / "preamble").read_text()
    postscript = (Path(__file__).parent / "postscript").read_text()

    with redirect_stdout(f):
        print(preamble)
        print()
        print(xrandr_command())
        print()
        for binding in bindings:
            if isinstance(binding, Mode):
                print(binding.mode_def(None))
            print(binding.binding())

        print()
        print(workspace_config(), end="")
        print(postscript)

    return f.getvalue()


def main():
    print(binds_config())


if __name__ == "__main__":
    main()