VoidShell/services/sway/workspaces.ts

import GObject, { register, property, signal } from "astal/gobject";
import { readFileAsync, monitorFile } from "astal/file";
import * as types from "./types";
import * as defs from "./definitions";
import Gdk from "gi://Gdk";
import Gio from "gi://Gio";
import GLib from "gi://GLib";
import { Sway } from "./ipc";
import { InvalidContextError } from "./exceptions";
import { delay, dereferenceSymbolicLink } from "@/utils";
import { DBusObject } from "@services/dbus";
import { computeContextScore, computeGroupScore } from "./scoring";
import { Mutex } from "@/synchronization";

@register({
  GTypeName: "SwayWorkspace",
})
export class Workspace extends DBusObject {
  #active: boolean = false;
  #focused: boolean = false;
  #visible: boolean = false;
  #urgent: boolean = false;

  groups: Set<Group> = new Set();
  #currentGroup: Group | null = null;
  #currentOutput: string | null = null;

  constructor(public readonly definition: defs.Workspace) {
    super(
      "dev.ezri.voidshell.workspaces.Workspace",
      `/dev/ezri/VoidShell/workspaces/workspace/${definition.index}`,
    );
  }

  @property(String)
  get indicatorClassName() {
    const result = ["indicator-circle", "sway--ws"];
    if (this.active) {
      result.push("sway--active");
    }
    if (this.visible) {
      result.push("sway--visible");
    }
    if (this.focused) {
      result.push("sway--focused");
    }
    return result.join(" ");
  }

  /**
   * Create an ephemeral workspace from a sync operation that is not defined in the system.
   */
  public static fromSync(update: types.WorkspaceResponse) {
    const [name, index] = (() => {
      if (parseInt(update.name).toString() === update.name) {
        return [`ephemeral: ${update.name}`, parseInt(update.name)];
      } else {
        return [update.name, update.num];
      }
    })();
    const definition: defs.Workspace = {
      index,
      name,
      exec: "console",
      program_name: "console",
    };
    const workspace = new this(definition);
    workspace.sync(update);
    return workspace;
  }

  @property(String)
  get objectPath() {
    return `/dev/ezri/VoidShell/workspaces/workspace/${this.definition.index}`;
  }

  @property(Boolean)
  get active() {
    return this.#active;
  }

  set active(value: boolean) {
    if (this.#active === value) {
      return;
    }
    this.#active = value;
    this.notify("active");
    this.notify("indicator-class-name");
    this.dbusObj?.emit_property_changed(
      "active",
      GLib.Variant.new_boolean(value),
    );
    if (!value) {
      // Whenever a workspace is deactivated, all other flags should be unset.
      this.focused = false;
      this.visible = false;
      this.urgent = false;
    }
  }

  @property(Boolean)
  get focused() {
    return this.#focused;
  }

  set focused(value: boolean) {
    if (this.#focused === value) {
      return;
    }
    this.#focused = value;
    this.notify("focused");
    this.notify("indicator-class-name");
    this.dbusObj?.emit_property_changed(
      "focused",
      GLib.Variant.new_boolean(value),
    );
  }

  @property(Boolean)
  get visible() {
    return this.#visible;
  }

  set visible(value: boolean) {
    if (this.#visible === value) {
      return;
    }
    this.#visible = value;
    this.notify("visible");
    this.notify("indicator-class-name");
    this.dbusObj?.emit_property_changed(
      "visible",
      GLib.Variant.new_boolean(value),
    );
  }

  @property(Boolean)
  get urgent() {
    return this.#urgent;
  }

  set urgent(value: boolean) {
    if (this.#urgent === value) {
      return;
    }
    this.#urgent = value;
    this.notify("urgent");
    this.dbusObj?.emit_property_changed(
      "urgent",
      GLib.Variant.new_boolean(value),
    );
  }

  @property(String)
  get name() {
    return this.definition.name;
  }

  @property(Group)
  /**
   * The group that this workspace is currently assigned to.
   *
   * While a workspace can be in multiple groups, including multiple groups
   * in the same context, it can only be assigned to one group at a time. If
   * activated in a different group in the same context, it will be reassigned
   * to that group.
   *
   * An inactive workspace will not have a current group, even if it only belongs
   * to a single group in the active context.
   */
  get currentGroup() {
    return this.#currentGroup;
  }

  set currentGroup(group: Group | null) {
    if (group !== this.#currentGroup) {
      this.#currentGroup = group;
      this.notify("current-group");
    }
  }

  @property(String)
  get currentOutput() {
    return this.#currentGroup?.outputName ?? null;
  }

  set currentOutput(output) {
    if (output !== this.#currentOutput) {
      this.#currentOutput = output;
      this.notify("current-output");
    }
  }

  deactivate() {
    console.log(`Deactivating workspace ${this.definition.index}`);
    this.active = false;
  }

  sync(update: types.WorkspaceResponse) {
    this.active = true;
    this.visible = update.visible;
    this.focused = update.focused;
    this.urgent = update.urgent;
  }

  async relocate(to: Group | string, retainFocus: boolean = false) {
    if (typeof to === "string") {
      const group = Tree.instance.findGroupByObjectPath(to);
      if (group === undefined) {
        throw new Error("Group not found");
      }
      to = group;
    }
    if (this.#currentGroup === to) {
      return;
    }

    await Sway.instance.command(
      `workspace ${this.definition.index}; move workspace to output ${to.outputName}`,
    );
  }

  async focus(to?: Group) {
    if (this.#currentGroup === to || to === undefined) {
      await Sway.instance.command(`workspace ${this.definition.index}`);
    } else {
      await this.relocate(to);
    }
  }

  async moveContainerTo() {
    await Sway.instance.command(
      `move container to workspace ${this.definition.index}`,
    );
  }
}

@register({
  GTypeName: "SwayWorkspaceGroup",
})
export class Group extends DBusObject {
  #workspaces: Workspace[] = [];
  #context: Context;
  #outputName: string | null = null;

  #monitor: Gdk.Monitor | null = null;

  @property(Context)
  get context() {
    return this.#context;
  }

  @property(Workspace)
  get focusedWorkspace() {
    return this.#workspaces.find((workspace) => workspace.visible);
  }

  @property(Workspace)
  get focused() {
    return this.focusedWorkspace?.focused ?? false;
  }

  @property(Gdk.Monitor)
  get monitor() {
    return this.#monitor;
  }

  @property(String)
  get outputName() {
    return this.#outputName ?? "";
  }

  @property(Object)
  get workspacePaths() {
    return this.#workspaces.map((workspace) => workspace.objectPath);
  }

  @property(String)
  get focusedWorkspacePath() {
    return this.focusedWorkspace?.objectPath ?? "";
  }

  @signal()
  declare gainedFocus: () => void;

  @signal()
  declare lostFocus: () => void;

  @signal(Gdk.Monitor)
  declare activated: (monitor: Gdk.Monitor) => void;

  @signal()
  declare deactivated: () => void;

  constructor(
    public readonly definition: defs.Group,
    public readonly name: string,
    public readonly outputDefinition: defs.Output,
    context: Context,
  ) {
    super(
      "dev.ezri.voidshell.workspaces.Group",
      `${context.objectPath}/group/${name}`,
    );
    this.#context = context;
    this.#workspaces = definition.workspaces
      .map((index) => Tree.instance.workspaces.get(index.toString()))
      .filter((workspace) => workspace !== undefined);

    this.#workspaces.forEach((workspace) => {
      workspace.connect(
        "notify::visible",
        // Bind the function, just to be safe.
        (() => {
          if (workspace.visible) {
            this.notify("focused-workspace");
          }
        }).bind(this),
      );
      workspace.connect(
        "notify::focused",
        (() => {
          if (workspace === this.focusedWorkspace) {
            this.notify("focused");
          }
        }).bind(this),
      );
    });
  }

  /**
   * Compute the score of this group based on the given Sway output.
   */
  score(output: types.OutputResponse): number {
    return computeGroupScore(this.outputDefinition, output);
  }

  async activate(output: types.OutputResponse) {
    // Ensure that this output is compatible with the group.
    const score = this.score(output);
    if (score === 0) {
      console.error(`Output ${output.name} is not compatible with this group`);
      throw new Error(
        `Output ${output.name} is not compatible with this group.`,
      );
    }

    // Configure monitor
    await Sway.instance.command(
      `output ${output.name} pos ${this.outputDefinition.position[0]} ${this.outputDefinition.position[1]} mode ${this.outputDefinition.mode}`,
    );

    // Enable monitor.
    await Sway.instance.command(`output ${output.name} enable`);
    // Store monitor name.
    this.#outputName = output.name;

    // Retrieve GDK monitor.
    const monitor =
      Gdk.Display.get_default()?.get_monitor_at_point(
        this.outputDefinition.position[0],
        this.outputDefinition.position[1],
      ) ?? null;

    // Notify listeners.
    this.#monitor = monitor;
    if (monitor !== null) {
      this.activated(monitor);
    }
    this.notify("monitor");

    // If the output's focused workspace is not in this group, focus the first workspace.
    if (!this.#workspaces.some((workspace) => workspace.visible)) {
      await this.#workspaces[0].focus(this);
    }

    // Add ourselves to each workspace
    this.#workspaces.forEach((workspace) => workspace.groups.add(this));
  }

  async dummyActivate(output: types.OutputResponse) {
    const score = this.score(output);
    if (score === 0) {
      console.error(
        `Output ${output.name} is not compatible with this group. Dummy activation, ignoring...`,
      );
    }

    this.#outputName = output.name;
    const monitor =
      Gdk.Display.get_default()?.get_monitor_at_point(
        this.outputDefinition.position[0],
        this.outputDefinition.position[1],
      ) ?? null;
    this.#monitor = monitor;
    if (monitor !== null) {
      this.activated(monitor);
    }
    this.notify("monitor");

    this.#workspaces.forEach((workspace) => workspace.groups.add(this));
  }

  async focus() {
    if (this.#outputName === null) {
      throw new Error("Group is not activated.");
    }
    await Sway.instance.command(`focus output ${this.#outputName}`);
  }

  getWorkspace(index: number): Workspace | undefined {
    if (index >= this.#workspaces.length) {
      return undefined;
    }
    return this.#workspaces[index];
  }

  async focusWorkspace(index: number): Promise<boolean> {
    const workspace = this.getWorkspace(index);
    if (workspace === undefined) return false;
    try {
      await workspace.focus(this);
    } catch {
      return false;
    }
    return true;
  }

  deactivate() {
    // Remove ourselves from every workspace
    this.#workspaces.forEach((workspace) => workspace.groups.delete(this));
  }

  async focusNextWorkspace() {
    const index =
      (this.#workspaces.findIndex((ws) => ws === this.focusedWorkspace) + 1) %
      this.#workspaces.length;
    return this.#workspaces[index].focus();
  }

  async focusPreviousWorkspace() {
    const index =
      (this.#workspaces.findIndex((ws) => ws === this.focusedWorkspace) - 1) %
      this.#workspaces.length;
    return this.#workspaces[index].focus();
  }
}

@register({
  GTypeName: "SwayWorkspaceContext",
})
export class Context extends DBusObject {
  #groups: Map<string, Group> = new Map();
  #focusedGroup: Group | null = null;
  #groupOrder: string[] = [];
  #active: boolean = false;

  constructor(
    public readonly definition: defs.Context,
    public readonly name: string,
  ) {
    super(
      "dev.ezri.voidshell.workspaces.Context",
      `/dev/ezri/VoidShell/workspaces/context/${name.replace("-", "_")}`,
    );
    definition.outputs.forEach((output) => {
      if (output.group === undefined || !definition.groups[output.group]) {
        throw new Error(`Group ${output.group} not found in context`);
      }
      const group = new Group(
        definition.groups[output.group],
        output.group,
        output,
        this,
      );
      this.#groups.set(output.group, group);
    });
    // Group order is defined by the order in which the groups are defined in the context, not the order in which they are defined in the outputs.
    this.#groupOrder = Object.keys(definition.groups);
  }

  @property(Group)
  get focusedGroup() {
    return this.#focusedGroup;
  }

  get groups(): Omit<Map<string, Group>, "set" | "delete" | "clear"> {
    return this.#groups;
  }

  private set focusedGroup(value) {
    if (this.#focusedGroup !== value) {
      this.#focusedGroup = value;
      this.notify("focused-group");
      this.dbusObj?.emit_property_changed(
        "focusedGroupPath",
        GLib.Variant.new_string(value?.objectPath ?? "/"),
      );
    }
  }

  @property(String)
  get focusedGroupPath() {
    return this.#focusedGroup?.objectPath ?? "/";
  }

  @property(Boolean)
  get active() {
    return this.#active;
  }

  private set active(value: boolean) {
    if (value !== this.#active) {
      this.#active = value;
      this.notify("active");
      this.dbusObj?.emit_property_changed(
        "active",
        GLib.Variant.new_boolean(value),
      );
      if (!value) {
        Tree.instance.deactivating();
      }
    }
  }

  getGroupOnOutput(output: string): Group | undefined {
    return Array.from(this.#groups.values()).find(
      (group) => group.outputName === output,
    );
  }

  /**
   * This function implements the output matching algorithm.
   * @param outputSet The set of outputs to match against. Usually the currently connected outputs, as retrieved from Sway.
   * @returns A map of group names to outputs and scores.
   */
  private matchGroupsWithOutputs(outputSet: types.OutputResponse[]) {
    const [_, match] = computeContextScore(this.definition, outputSet);
    return match;
  }

  /**
   * Compute the score of this context based on the given Sway output set.
   */
  score(outputSet: types.OutputResponse[], includePriority = true) {
    const [score] = computeContextScore(this.definition, outputSet);
    if (includePriority && score !== 0) {
      return score + (this.definition.priority ?? 0);
    }
    return score;
  }

  /**
   * Activate this context.
   */
  async activate() {
    console.log("Context activation requested");
    const outputs = await Sway.instance.getOutputs();
    const match = this.matchGroupsWithOutputs(outputs);
    if (match === undefined) {
      throw new InvalidContextError(
        this,
        "Could not activate context: context is incompatible with current outputs.",
      );
    }
    // Construct a set of output names that this context will use.
    const usedOutputs = new Set(
      Array.from(match.values()).map((output) => output.name),
    );
    if (Tree.instance.currentContext !== null) {
      await Tree.instance.currentContext.deactivate(usedOutputs);
    }
    // Activate groups.
    await Array.from(match.entries()).amap(async ([name, output]) => {
      const group = this.#groups.get(name);
      console.log(`Activating group ${name} on output ${output.name}`);
      if (group === undefined) {
        throw new InvalidContextError(
          this,
          `Could not activate context: group ${name} is not defined.`,
        );
      }
      await group.activate(output);
    });
    // Notify listeners.
    this.active = true;
    Tree.instance.notify("current-context");
  }

  /**
   * Mark the context as activated, without changing outputs
   */
  async dummyActivate() {
    console.log(
      "Context dummy activation requested. This will not produce any changes to the current output configuration!",
    );
    const outputs = await Sway.instance.getOutputs();
    const match = this.matchGroupsWithOutputs(outputs);
    if (match === undefined) {
      throw new InvalidContextError(
        this,
        "Could not activate context: context is incompatible with current outputs.",
      );
    }
    await Tree.instance.currentContext?.dummyDeactivate();

    await Array.from(match.entries()).amap(async ([name, output]) => {
      const group = this.#groups.get(name);
      console.log(`Dummy-activating group ${name} on output ${output.name}`);
      if (group === undefined) {
        throw new InvalidContextError(
          this,
          `Could not activate context: group ${name} is not defined.`,
        );
      }
      await group.dummyActivate(output);
    }),
      (this.active = true);
    Tree.instance.notify("current-context");
  }

  async dummyDeactivate() {
    this.active = false;
    await new Promise((resolve) => setTimeout(resolve, 1));
    Array.from(this.#groups.values()).forEach((group) => group.deactivate());
  }

  /**
   * Deactivate this context.
   * @param outputsToLeave A set of outputs that should not be disabled, selected by output name
   */
  async deactivate(outputsToLeave: Set<string>) {
    // First, notify for deactivation and move ourselves to the lowest priority in the event list
    this.active = false;
    await Promise.resolve(); // This should allow the event loop to run and any event listeners to run before we continue

    // Then, deactivate the outputs and groups.
    const command = Array.from(this.#groups.values())
      .map((group) => {
        group.deactivate();
        return group;
      })
      .filter((group) => !outputsToLeave.has(group.outputName ?? ""))
      .map((group) => `output ${group.outputName} disable`)
      .join("; ");
    await Sway.instance.command(command);
  }

  getWorkspace(
    index: number,
  ): [workspace: Workspace, group: Group] | [null, null] {
    const currentGroupIndex = Math.max(
      this.#groupOrder.indexOf(this.focusedGroup?.name ?? ""),
      0,
    );
    const groupName = this.#groupOrder
      .rotate(currentGroupIndex)
      .find((name) => {
        const group = this.#groups.get(name);
        if (group === undefined) {
          return false;
        }
        const workspace = group.getWorkspace(index);
        if (workspace === this.focusedGroup?.focusedWorkspace) {
          // If the workspace is the focused workspace of the focused group, don't return it; we want to search other outputs.
          return false;
        }
        // If the workspace at this index exists on this output, return it. Otherwise, return the accumulator.
        return !!workspace;
      });
    if (groupName === undefined) {
      return [null, null];
    }
    const group = this.#groups.get(groupName) ?? null;
    const workspace = group?.getWorkspace(index) ?? null;
    if (workspace === null || group === null) {
      return [null, null];
    }
    return [workspace, group];
  }
}

@register({
  GTypeName: "SwayWorkspaceTree",
})
export class Tree extends DBusObject {
  #workspaces: Map<string, Workspace> = new Map();
  #contexts: Map<string, Context> = new Map();

  static _instance: Tree | null = null;

  public get workspaces() {
    return this.#workspaces;
  }

  async #parseFile(path: string) {
    console.log(`reading file ${path}`);
    const contents = await readFileAsync(path);
    const data = JSON.parse(contents) as defs.Config;
    this.#workspaces = new Map();
    this.#contexts = new Map();

    // Create workspaces
    data.workspaces.forEach((workspace) => {
      this.#workspaces.set(
        workspace.index.toString(),
        new Workspace(workspace),
      );
    });

    // Create contexts
    Object.entries(data.contexts).forEach(([name, context]) => {
      this.#contexts.set(name, new Context(context, name));
    });

    // Notify for tree refresh
    this.emit("tree-refresh");

    // If we're connected to Sway, immediately do a full tree sync. Otherwise,
    // wait to be connected and then do a full tree sync.
    if (Sway.instance.isConnected) {
      await this.sync();
      this.onOutputEvent();
    } else {
      const connId = Sway.instance.connect("connected", async () => {
        Sway.instance.disconnect(connId);
        await this.sync();
        this.onOutputEvent();
      });
    }
  }

  /**
   * Synchronize the state of VoidShell's workspace tree with Sway
   */
  async sync() {
    const touched = new Set<Workspace>();
    try {
      const workspaces = await Sway.instance.getWorkspaces();
      workspaces.forEach((workspace) => {
        const ws = this.#workspaces.get(workspace.name);
        if (ws === undefined) {
          this.#workspaces.set(workspace.name, Workspace.fromSync(workspace));
        } else {
          ws.sync(workspace);
          touched.add(ws);
        }
      });
      this.#workspaces.forEach((workspace) => {
        if (!touched.has(workspace)) {
          workspace.active = false;
        }
      });
    } catch (e) {
      console.error("Failed to sync", e);
    }
  }

  static get instance() {
    if (!this._instance) {
      throw new Error("Tree instance not initialized");
    }
    return Tree._instance as Tree;
  }

  constructor(filepath: string) {
    super();
    console.log(`Parsing file ${filepath}`);
    if (Tree._instance !== null) {
      console.error("Tree instance already initialized...");
      throw new Error("Tree instance already initialized");
    }
    Tree._instance = this;
    console.log("Dereferencing symlink");
    dereferenceSymbolicLink(filepath)
      .then(async (resolvedPath) => {
        console.log(`Resolved ${filepath} to ${resolvedPath}`);
        try {
          await this.#parseFile(resolvedPath);
        } catch (e) {
          console.error("Error parsing config file:", e);
        }
        monitorFile(filepath, async (_, event) => {
          if (event === Gio.FileMonitorEvent.CHANGES_DONE_HINT) {
            const oldWorkspaces = this.#workspaces;
            const oldContexts = this.#contexts;
            try {
              await this.#parseFile(filepath);
            } catch (e) {
              console.error("Error parsing config file:", e);
              this.#workspaces = oldWorkspaces;
              this.#contexts = oldContexts;
            }
          }
        });
      })
      .catch((reason) => {
        console.error("Error dereferencing symlink:", reason);
      });
    Sway.instance.connect("workspace", (_, event) => {
      this.onWorkspaceEvent(event);
    });
    Sway.instance.connect("output", (_, event) => this.onOutputEvent(event));
  }

  private onWorkspaceEvent(event: types.WorkspaceEvent) {
    switch (event.change) {
      case types.WorkspaceEventChange.RELOAD:
        break;
      default:
        this.sync();
    }
  }

  private outputDelayLock = new Mutex();

  /**
   * Callback for when we receive output events.
   *
   * @param event - The event to process. Currently, this event is always the same ({ change: "unspecified" }), so it is unused.
   */
  private async onOutputEvent(_?: types.OutputEvent) {
    try {
      // Wait for half a second, since a single hardware change can trigger a flurry of output events
      await this.outputDelayLock.immediate(delay.bind(this, 500));
    } catch {
      return;
    }
    const outputs = await Sway.instance.getOutputs();
    // Score each context
    const scores = Array.from(this.#contexts.entries()).reduce(
      (acc, [name, context]) => {
        acc[name] = context.score(outputs);
        return acc;
      },
      {} as Record<string, number>,
    );
    console.log(scores);

    // Find highest score
    const [highName] = Object.entries(scores).sort(
      ([_, a], [__, b]) => b - a,
    )[0];
    if (highName === undefined) {
      console.warn(
        "No configured contexts are compatible with the current display arrangement!",
      );
      return;
    }
    const newContext = this.#contexts.get(highName);
    if (newContext === undefined) {
      console.error(
        "High score context does not exist in memory. This should not happen.",
      );
      return;
    }
    if (newContext === this.currentContext) {
      console.log("Context is already active, ignoring...");
    }
    console.log("Activating context", highName);
    await newContext.dummyActivate();
  }

  @property(Context)
  get currentContext(): Context | null {
    return (
      Array.from(this.#contexts.values()).find((context) => context.active) ??
      null
    );
  }

  @property(String)
  get currentContextName(): string {
    return (
      Array.from(this.#contexts.entries()).find(
        ([, context]) => context.active,
      )?.[0] ?? ""
    );
  }

  @signal()
  /**
   * Emitted when a context begins deactivating.
   */
  declare deactivating: () => void;

  @signal()
  /**
   * Emitted after a full tree refresh which invalidates all context, group, and workspace objects.
   */
  declare treeRefresh: () => void;

  async findNewContext() {
    const outputs = await Sway.instance.getOutputs();
    const topScore = Array.from(this.#contexts.values())
      .map((context) => [context, context.score(outputs)] as const)
      .sort(([, a], [, b]) => b - a)
      .find(([_, score]) => score > 0)?.[0];
    if (topScore === undefined) {
      // No available context
      return;
    }
    await topScore.activate();
  }

  async requestContext(name: string) {
    const context = this.#contexts.get(name);
    if (context === undefined) {
      throw new Error(`Context ${name} not found`);
    }
    // Get current outputs
    const outputs = await Sway.instance.getOutputs();
    // Score the context
    const score = context.score(outputs, false);
    if (score > 0) {
      await context.activate();
    }
    throw new Error(`Context ${name} is incompatible with current outputs`);
  }

  getWorkspace(
    groupIndex: number,
  ): [workspace: Workspace, group: Group] | [null, null] {
    return this.currentContext?.getWorkspace(groupIndex) ?? [null, null];
  }

  findGroupByObjectPath(path: string): Group | undefined {
    const match = path.match(
      new RegExp(
        "/dev/ezri/VoidShell/workspaces/context/([\\w_]+)/group/([\\w_]+)",
      ),
    );
    if (match === null) {
      console.error(`Group object path ${path} invalid`);
      return undefined;
    }
    const context =
      this.#contexts.get(match[1]) ??
      this.#contexts.get(match[1].replace("_", "-"));
    if (context === undefined) return undefined;
    return (
      context.groups.get(match[2]) ??
      context.groups.get(match[2].replace("_", "-"))
    );
  }
}