tiletopia/src/lib/layout/tree.ts

//! Splits-tree layout model.
//!
//! The workspace is a binary tree of splits. Internal nodes are HSplit or
//! VSplit with a ratio; leaves are terminal panes. This is the same model
//! tmux / i3 / Zellij use — dragging a gutter mutates one parent ratio,
//! both sibling subtrees reflow automatically.

export type NodeId = string;

/** 'h' = side-by-side (a on left, b on right). 'v' = stacked (a on top, b below). */
export type Orientation = "h" | "v";

export interface LeafNode {
  kind: "leaf";
  id: NodeId;
  /** WSL distro the pane was spawned against. */
  distro?: string;
  /** Working directory the pane was started in. Not currently used at spawn time but preserved for future. */
  cwd?: string;
  /** Optional user label shown in the pane toolbar. */
  label?: string;
  /**
   * If true, keystrokes typed in this pane are mirrored to every other
   * leaf with `broadcast === true`. Toggle via the 📡 button in the
   * pane toolbar.
   */
  broadcast?: boolean;
}

export interface SplitNode {
  kind: "split";
  id: NodeId;
  orientation: Orientation;
  /** Size of `a` divided by total. Kept in (0, 1) — clamped on drag. */
  ratio: number;
  a: TreeNode;
  b: TreeNode;
}

export type TreeNode = LeafNode | SplitNode;

function newId(): NodeId {
  return (
    globalThis.crypto?.randomUUID?.() ??
    Math.random().toString(36).slice(2, 12)
  );
}

export function newLeaf(props: Partial<Omit<LeafNode, "kind" | "id">> = {}): LeafNode {
  return { kind: "leaf", id: newId(), ...props };
}

export function newSplit(
  orientation: Orientation,
  a: TreeNode,
  b: TreeNode,
  ratio = 0.5,
): SplitNode {
  return { kind: "split", id: newId(), orientation, ratio, a, b };
}

/** Walk the tree, replacing the node with the given id via `produce`. Returns a new tree. */
export function replaceById(
  root: TreeNode,
  targetId: NodeId,
  produce: (node: TreeNode) => TreeNode,
): TreeNode {
  if (root.id === targetId) return produce(root);
  if (root.kind === "leaf") return root;
  const a = replaceById(root.a, targetId, produce);
  const b = replaceById(root.b, targetId, produce);
  if (a === root.a && b === root.b) return root;
  return { ...root, a, b };
}

/** Split the leaf with the given id. New pane goes on side `b`. */
export function splitLeaf(
  root: TreeNode,
  leafId: NodeId,
  orientation: Orientation,
  newLeafProps: Partial<Omit<LeafNode, "kind" | "id">> = {},
): TreeNode {
  return replaceById(root, leafId, (node) => {
    if (node.kind !== "leaf") return node;
    return newSplit(orientation, node, newLeaf(newLeafProps));
  });
}

/**
 * Remove the leaf with the given id. The other child of its parent split
 * takes the parent's place in the tree. Returns null if the closed leaf
 * was the entire tree (caller should create a fresh leaf).
 */
export function closeLeaf(root: TreeNode, leafId: NodeId): TreeNode | null {
  if (root.id === leafId) return null;
  if (root.kind === "leaf") return root;
  // If a direct child is the target leaf, collapse this split to the sibling.
  if (root.a.kind === "leaf" && root.a.id === leafId) return root.b;
  if (root.b.kind === "leaf" && root.b.id === leafId) return root.a;
  // Recurse.
  const newA = closeLeaf(root.a, leafId);
  const newB = closeLeaf(root.b, leafId);
  // If either side collapsed to null somehow (target was a split, shouldn't
  // happen with current callers but defensive), fall back to the other side.
  if (newA === null) return newB ?? root;
  if (newB === null) return newA;
  if (newA === root.a && newB === root.b) return root;
  return { ...root, a: newA, b: newB };
}

/** Find a leaf by id. Returns null if not found. */
export function findLeaf(root: TreeNode, leafId: NodeId): LeafNode | null {
  if (root.kind === "leaf") return root.id === leafId ? root : null;
  return findLeaf(root.a, leafId) ?? findLeaf(root.b, leafId);
}

/**
 * Swap the distro on a leaf. The leaf gets a **new id** so the rendering
 * layer's `{#key node.id}` block remounts XtermPane — the old PTY is killed
 * and a fresh one spawns with the new distro.
 */
export function changeDistro(
  root: TreeNode,
  leafId: NodeId,
  distro: string,
): TreeNode {
  return replaceById(root, leafId, (node) => {
    if (node.kind !== "leaf") return node;
    return { ...node, id: newId(), distro };
  });
}

/** Set or clear a leaf's label. Does NOT remount (label is metadata only). */
export function changeLabel(
  root: TreeNode,
  leafId: NodeId,
  label: string | undefined,
): TreeNode {
  return replaceById(root, leafId, (node) => {
    if (node.kind !== "leaf") return node;
    const trimmed = label?.trim();
    return { ...node, label: trimmed ? trimmed : undefined };
  });
}

// ---- preset layouts --------------------------------------------------------

type LeafDefaults = Partial<Omit<LeafNode, "kind" | "id">>;

export function presetSingle(d: LeafDefaults = {}): TreeNode {
  return newLeaf(d);
}

export function presetTwoColumns(d: LeafDefaults = {}): TreeNode {
  return newSplit("h", newLeaf(d), newLeaf(d));
}

export function presetThreeColumns(d: LeafDefaults = {}): TreeNode {
  // Even thirds: outer split at 1/3, inner split at 1/2.
  return newSplit(
    "h",
    newLeaf(d),
    newSplit("h", newLeaf(d), newLeaf(d), 0.5),
    1 / 3,
  );
}

export function presetTwoRows(d: LeafDefaults = {}): TreeNode {
  return newSplit("v", newLeaf(d), newLeaf(d));
}

export function presetTwoByTwo(d: LeafDefaults = {}): TreeNode {
  const row = () => newSplit("h", newLeaf(d), newLeaf(d));
  return newSplit("v", row(), row());
}

/** Number of leaves in the tree. */
export function leafCount(root: TreeNode): number {
  if (root.kind === "leaf") return 1;
  return leafCount(root.a) + leafCount(root.b);
}

/** Iterate all leaves in left-to-right order. */
export function* walkLeaves(root: TreeNode): Generator<LeafNode> {
  if (root.kind === "leaf") {
    yield root;
  } else {
    yield* walkLeaves(root.a);
    yield* walkLeaves(root.b);
  }
}

/** Toggle a leaf's broadcast flag. Metadata-only — does NOT swap the id, so the pane is not respawned. */
export function toggleBroadcast(root: TreeNode, leafId: NodeId): TreeNode {
  return replaceById(root, leafId, (node) => {
    if (node.kind !== "leaf") return node;
    return { ...node, broadcast: !node.broadcast };
  });
}

/**
 * Reshape the tree into the structure produced by `preset`, but PRESERVE
 * existing leaves (and their PTYs) by copying their id/distro/cwd/label/
 * broadcast into the preset's slots, in DFS order.
 *
 * - If the preset has more slots than existing leaves, the extra slots stay
 *   as their freshly-created (empty) leaves — those panes will spawn new
 *   shells.
 * - If the preset has fewer slots than existing leaves, the overflow leaves
 *   are returned in `dropped` so the caller can kill their PTYs.
 *
 * Split ratios reset to the preset's defaults (0.5 — the user's previous
 * resize work is discarded; that's the point of "apply preset").
 */
export function reshapeToPreset(
  existing: TreeNode,
  preset: (d: LeafDefaults) => TreeNode,
  defaults: LeafDefaults,
): { tree: TreeNode; dropped: NodeId[] } {
  const existingLeaves = Array.from(walkLeaves(existing));
  const tree = preset(defaults);
  const slots = Array.from(walkLeaves(tree));
  const dropped: NodeId[] = [];

  for (let i = 0; i < slots.length; i++) {
    const src = existingLeaves[i];
    if (!src) break;
    const slot = slots[i];
    slot.id = src.id;
    if (src.distro !== undefined) slot.distro = src.distro;
    if (src.cwd !== undefined) slot.cwd = src.cwd;
    if (src.label !== undefined) slot.label = src.label;
    if (src.broadcast !== undefined) slot.broadcast = src.broadcast;
  }

  for (let i = slots.length; i < existingLeaves.length; i++) {
    dropped.push(existingLeaves[i].id);
  }

  return { tree, dropped };
}

/** Force every leaf's broadcast flag to `on`. Returns the same root reference
 *  when nothing actually changed, so callers can skip a state update if so. */
export function setAllBroadcast(root: TreeNode, on: boolean): TreeNode {
  if (root.kind === "leaf") {
    if (!!root.broadcast === on) return root;
    return { ...root, broadcast: on };
  }
  const a = setAllBroadcast(root.a, on);
  const b = setAllBroadcast(root.b, on);
  if (a === root.a && b === root.b) return root;
  return { ...root, a, b };
}

// --- flat layout (for absolute-positioned rendering) ------------------------

/** Normalised bounding box: top/left/width/height as fractions [0, 1]. */
export interface Box {
  top: number;
  left: number;
  width: number;
  height: number;
}

/** A leaf rendered as a flat sibling: its current LeafNode plus the box
 *  it occupies in the container. */
export interface LeafSlot {
  leaf: LeafNode;
  box: Box;
}

/** A draggable gutter at a split boundary. `box` is where to render the
 *  draggable strip; `parentBox` is the area the gutter divides (needed to
 *  convert pointer position → ratio). */
export interface GutterInfo {
  splitId: NodeId;
  orientation: Orientation;
  ratio: number;
  box: Box;
  parentBox: Box;
}

/** Walk the tree and produce a flat list of leaf slots + draggable gutters.
 *  Renderer uses these to position all leaves as siblings in the DOM, which
 *  lets React preserve component instances (and thus PTYs) across any tree
 *  reshape — splits, closes, presets, etc. */
export function flattenLayout(
  root: TreeNode,
  box: Box = { top: 0, left: 0, width: 1, height: 1 },
): { leaves: LeafSlot[]; gutters: GutterInfo[] } {
  if (root.kind === "leaf") {
    return { leaves: [{ leaf: root, box }], gutters: [] };
  }
  const isH = root.orientation === "h";
  const r = root.ratio;
  let boxA: Box;
  let boxB: Box;
  let gutter: GutterInfo;
  if (isH) {
    const splitPos = box.width * r;
    boxA = { top: box.top, left: box.left, width: splitPos, height: box.height };
    boxB = {
      top: box.top,
      left: box.left + splitPos,
      width: box.width - splitPos,
      height: box.height,
    };
    gutter = {
      splitId: root.id,
      orientation: "h",
      ratio: r,
      box: {
        top: box.top,
        left: box.left + splitPos,
        width: 0,
        height: box.height,
      },
      parentBox: box,
    };
  } else {
    const splitPos = box.height * r;
    boxA = { top: box.top, left: box.left, width: box.width, height: splitPos };
    boxB = {
      top: box.top + splitPos,
      left: box.left,
      width: box.width,
      height: box.height - splitPos,
    };
    gutter = {
      splitId: root.id,
      orientation: "v",
      ratio: r,
      box: {
        top: box.top + splitPos,
        left: box.left,
        width: box.width,
        height: 0,
      },
      parentBox: box,
    };
  }
  const a = flattenLayout(root.a, boxA);
  const b = flattenLayout(root.b, boxB);
  return {
    leaves: [...a.leaves, ...b.leaves],
    gutters: [gutter, ...a.gutters, ...b.gutters],
  };
}

/** Update a split's ratio by its id. */
export function updateSplitRatio(root: TreeNode, splitId: NodeId, ratio: number): TreeNode {
  return replaceById(root, splitId, (node) => {
    if (node.kind !== "split") return node;
    return { ...node, ratio };
  });
}

export type Direction = "left" | "right" | "up" | "down";

/** Spatial pane navigation: given an active leaf, find the nearest neighbor
 *  in the requested direction. Used for Ctrl+Shift+Arrow shortcuts. */
export function findNeighborInDirection(
  leaves: LeafSlot[],
  fromLeafId: NodeId,
  direction: Direction,
): NodeId | null {
  const from = leaves.find((s) => s.leaf.id === fromLeafId);
  if (!from) return null;
  const fromCenter = {
    x: from.box.left + from.box.width / 2,
    y: from.box.top + from.box.height / 2,
  };
  const EPS = 1e-3;

  let best: { id: NodeId; perpDist: number; primaryDist: number } | null = null;

  for (const slot of leaves) {
    if (slot.leaf.id === fromLeafId) continue;
    const center = {
      x: slot.box.left + slot.box.width / 2,
      y: slot.box.top + slot.box.height / 2,
    };
    let primary: number;
    let perp: number;

    switch (direction) {
      case "right":
        if (slot.box.left < from.box.left + from.box.width - EPS) continue;
        primary = center.x - fromCenter.x;
        perp = Math.abs(center.y - fromCenter.y);
        break;
      case "left":
        if (slot.box.left + slot.box.width > from.box.left + EPS) continue;
        primary = fromCenter.x - center.x;
        perp = Math.abs(center.y - fromCenter.y);
        break;
      case "down":
        if (slot.box.top < from.box.top + from.box.height - EPS) continue;
        primary = center.y - fromCenter.y;
        perp = Math.abs(center.x - fromCenter.x);
        break;
      case "up":
        if (slot.box.top + slot.box.height > from.box.top + EPS) continue;
        primary = fromCenter.y - center.y;
        perp = Math.abs(center.x - fromCenter.x);
        break;
    }

    if (
      best === null ||
      perp < best.perpDist ||
      (perp === best.perpDist && primary < best.primaryDist)
    ) {
      best = { id: slot.leaf.id, perpDist: perp, primaryDist: primary };
    }
  }

  return best?.id ?? null;
}

/** Swap two leaves' tree positions. Each leaf carries its own data
 *  (id, distro, cwd, label, broadcast) into the other's slot. PTYs stay
 *  alive because React keys on leaf.id and our renderer is flat. */
export function swapLeaves(root: TreeNode, idA: NodeId, idB: NodeId): TreeNode {
  if (idA === idB) return root;
  const a = findLeaf(root, idA);
  const b = findLeaf(root, idB);
  if (!a || !b) return root;
  function walk(n: TreeNode): TreeNode {
    if (n.kind === "leaf") {
      if (n.id === idA) return b!;
      if (n.id === idB) return a!;
      return n;
    }
    const na = walk(n.a);
    const nb = walk(n.b);
    if (na === n.a && nb === n.b) return n;
    return { ...n, a: na, b: nb };
  }
  return walk(root);
}

export function serialize(root: TreeNode): string {
  return JSON.stringify(root);
}

/** Parse JSON back to a tree. Returns null on invalid input. */
export function deserialize(json: string): TreeNode | null {
  try {
    const parsed = JSON.parse(json);
    if (!isTreeNode(parsed)) return null;
    return parsed;
  } catch {
    return null;
  }
}

function isTreeNode(x: unknown): x is TreeNode {
  if (typeof x !== "object" || x === null) return false;
  const o = x as Record<string, unknown>;
  if (typeof o.id !== "string") return false;
  if (o.kind === "leaf") return true;
  if (o.kind === "split") {
    return (
      (o.orientation === "h" || o.orientation === "v") &&
      typeof o.ratio === "number" &&
      isTreeNode(o.a) &&
      isTreeNode(o.b)
    );
  }
  return false;
}