# modules/darwin/yabai/default.nix
{ config, lib, pkgs, ... }:

let
  cfg = config.local.yabai;

  # symbolic hotkey IDs 118..126 == "Switch to Desktop 1".."Switch to Desktop 9"
  spaceKeycodes = [ 18 19 20 21 22 23 25 26 28 ];

  mkSpaceHotkey = i: {
    name = toString (117 + i);
    value = {
      enabled = true;
      value = {
        type = "standard";
        parameters = [
          (48 + i)                                   # ASCII '1'..'9'
          (builtins.elemAt spaceKeycodes (i - 1))    # virtual keycode
          cfg.spaceHotkeys.modifierMask
        ];
      };
    };
  };

  spaceHotkeys = builtins.listToAttrs
    (map mkSpaceHotkey (lib.range 1 cfg.spaceHotkeys.count));

  padArms = lib.concatStringsSep "\n" (lib.mapAttrsToList
    (uuid: pad: "        ${uuid}) pad=${toString pad} ;;")
    cfg.displayTopPadding);

  displayPaddingScript = pkgs.writeShellScript "yabai-display-padding" ''
    set -eu
    yabai="${cfg.package}/bin/yabai"

    "$yabai" -m query --displays \
      | ${lib.getExe pkgs.jq} -r '.[] | "\(.uuid) \(.spaces | map(tostring) | join(" "))"' \
    | while read -r uuid spaces; do
        case "$uuid" in
${padArms}
          *) pad=${toString cfg.defaultTopPadding} ;;
        esac
        for s in $spaces; do
          for k in top_padding; do
            "$yabai" -m config --space "$s" "$k" "$pad" || true
          done
        done
      done
  '';

  yabaiBin = "${cfg.package}/bin/yabai";

  directionBinds = lib.concatStringsSep "\n" (map (d: ''
    alt - ${d.key} : ${yabaiBin} -m window --focus ${d.dir} || ${yabaiBin} -m display --focus ${d.dir}
    shift + alt - ${d.key} : ${yabaiBin} -m window --swap ${d.dir} || ${yabaiBin} -m window --display ${d.dir}
  '') [
    { key = "up";    dir = "north"; }
    { key = "down";  dir = "south"; }
    { key = "right"; dir = "east";  }
    { key = "left";  dir = "west";  }
  ]);

  spaceMoveBinds = lib.concatStringsSep "\n"
    (map (i: "shift + alt - ${toString i} : ${yabaiBin} -m window --space ${toString i}")
      (lib.range 1 cfg.spaceHotkeys.count));

in {
  options.local.yabai = {
    enable = lib.mkEnableOption "yabai tiling window manager setup";

    package = lib.mkOption {
      type = lib.types.package;
      default = pkgs.yabai;
      description = "yabai package. Referenced by keybindings and the padding script.";
    };

    enableScriptingAddition = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = ''
        Only gates space creation/destruction and moving windows between
        spaces. Window geometry and padding work without it.
      '';
    };

    defaultPadding = lib.mkOption {
      type = lib.types.int;
      default = 10;
      description = ''
        Padding and window_gap for any display not listed in displayTopPadding.
        Sized for Retina (2x) panels.
      '';
    };

    defaultTopPadding = lib.mkOption {
      type = lib.types.int;
      default = 10;
      description = ''
        Padding for top edge of screen
      '';
    };

    displayTopPadding = lib.mkOption {
      type = lib.types.attrsOf lib.types.int;
      default = { };
      example = lib.literalExpression ''
        { "37D8832A-2D66-02CA-B9F7-8F30A301B230" = 40; }
      '';
      description = ''
        Display UUID -> padding value. The value is applied to all four
        paddings and to window_gap for every space on that display.

        Get UUIDs with:
          yabai -m query --displays | jq -r '.[] | "\(.uuid)  \(.frame)"'

        Do not key off frame dimensions: yabai reports points, not pixels,
        so a scaled 1080p panel will not report 1920x1080.
      '';
    };

    manageKeybindings = lib.mkOption {
      type = lib.types.bool;
      default = true;
      description = "Emit yabai skhd bindings into services.skhd.skhdConfig.";
    };

    spaceHotkeys = {
      enable = lib.mkEnableOption "opt+<n> macOS space-switching hotkeys" // {
        default = true;
      };

      count = lib.mkOption {
        type = lib.types.ints.between 1 9;
        default = 9;
        description = "Number of spaces to bind. macOS exposes 9 symbolic hotkeys.";
      };

      modifierMask = lib.mkOption {
        type = lib.types.int;
        default = 524288;  # option
        description = "Carbon modifier mask. 524288 = option, 1048576 = command.";
      };
    };

    extraSymbolicHotkeys = lib.mkOption {
      type = lib.types.attrs;
      default = { };
      description = ''
        Merged into com.apple.symbolichotkeys alongside the space hotkeys.
        Put non-yabai symbolic hotkeys here rather than defining
        CustomUserPreferences."com.apple.symbolichotkeys" in the host config --
        CustomUserPreferences is types.attrs and merges shallowly, so a second
        definition of that domain silently clobbers this one.
      '';
    };

    extraSettings = lib.mkOption {
      type = lib.types.attrs;
      default = { };
      description = "Extra keys merged into services.yabai.config.";
    };

    extraConfig = lib.mkOption {
      type = lib.types.lines;
      default = "";
      description = "Extra shell appended to yabairc (rules, signals).";
    };
  };

  config = lib.mkIf cfg.enable {
    services.yabai = {
      enable = true;
      inherit (cfg) package enableScriptingAddition;

      config = {
        mouse_follows_focus = "on";
        layout = "bsp";
        window_placement = "second_child";
        top_padding = toString cfg.defaultTopPadding;
        bottom_padding = toString cfg.defaultPadding;
        left_padding = toString cfg.defaultPadding;
        right_padding = toString cfg.defaultPadding;
        window_gap = toString cfg.defaultPadding;
        mouse_modifier = "alt";
        mouse_drop_action = "swap";
        mouse_action1 = "move";
        mouse_action2 = "resize";
        focus_follows_mouse = "autofocus";
      } // cfg.extraSettings;

      extraConfig = ''
        yabai -m rule --add app='System Settings' manage=off
        yabai -m rule --add app="^Zen$" title="^Picture-in-Picture$" manage=off
      ''
      + lib.optionalString (cfg.displayTopPadding != { }) ''

        # Per-display padding. space_created is required, not redundant:
        # display_added fires when yabai sees the display, which can precede
        # macOS finishing space creation on it, so the first pass may find
        # zero or one space. display_removed is required because yabai
        # reindexes spaces on disconnect and the overrides are keyed by index.
        yabai -m signal --add event=display_added   action="${displayPaddingScript}" label="per-display padding (added)"
        yabai -m signal --add event=display_removed action="${displayPaddingScript}" label="per-display padding (removed)"
        yabai -m signal --add event=space_created   action="${displayPaddingScript}" label="per-display padding (space)"
        ${displayPaddingScript}
      ''
      + cfg.extraConfig;
    };

    services.skhd = lib.mkIf cfg.manageKeybindings {
      enable = true;
      skhdConfig = ''
        shift + alt - q : ${yabaiBin} -m window --close

        ${directionBinds}

        shift + alt - space : ${yabaiBin} -m window --toggle float
        shift + alt - return : ${yabaiBin} -m window --toggle sticky
        alt - f : ${yabaiBin} -m window --toggle zoom-fullscreen
        shift + alt - f : ${yabaiBin} -m window --toggle native-fullscreen

        ${spaceMoveBinds}
      '';
    };

    # Prerequisites, not preferences: yabai requires per-display spaces, and
    # macOS's own tiling fights it for window placement.
    system.defaults = {
      spaces."spans-displays" = false;
      WindowManager.EnableTilingOptionAccelerator = false;

      CustomUserPreferences."com.apple.symbolichotkeys".AppleSymbolicHotKeys =
        cfg.extraSymbolicHotkeys
        // (lib.optionalAttrs cfg.spaceHotkeys.enable spaceHotkeys);
    };
  };
}