import type { PlotLayout } from './types';

/** Collect all widget ids referenced by leaves, in layout order. */
export function collectWidgetIds(layout: PlotLayout): string[] {
  if (layout.type === 'leaf') return [layout.widget];
  return [...collectWidgetIds(layout.a), ...collectWidgetIds(layout.b)];
}

/** Count the number of leaves (panes). */
export function countLeaves(layout: PlotLayout): number {
  if (layout.type === 'leaf') return 1;
  return countLeaves(layout.a) + countLeaves(layout.b);
}

/**
 * Replace the leaf for `widgetId` with a split: the existing leaf becomes child
 * `a` and a new leaf for `newWidgetId` becomes child `b` (ratio 0.5).
 */
export function splitLeaf(
  layout: PlotLayout,
  widgetId: string,
  dir: 'h' | 'v',
  newWidgetId: string,
): PlotLayout {
  if (layout.type === 'leaf') {
    if (layout.widget !== widgetId) return layout;
    return {
      type: 'split',
      dir,
      ratio: 0.5,
      a: { type: 'leaf', widget: widgetId },
      b: { type: 'leaf', widget: newWidgetId },
    };
  }
  return {
    ...layout,
    a: splitLeaf(layout.a, widgetId, dir, newWidgetId),
    b: splitLeaf(layout.b, widgetId, dir, newWidgetId),
  };
}

/**
 * Remove the leaf for `widgetId`, collapsing its parent split into the sibling
 * subtree. Returns the original layout if the leaf is the sole remaining pane.
 */
export function removeLeaf(layout: PlotLayout, widgetId: string): PlotLayout {
  if (layout.type === 'leaf') return layout; // can't remove the root leaf
  // If a direct child is the target leaf, collapse to the sibling.
  if (layout.a.type === 'leaf' && layout.a.widget === widgetId) return layout.b;
  if (layout.b.type === 'leaf' && layout.b.widget === widgetId) return layout.a;
  return {
    ...layout,
    a: removeLeaf(layout.a, widgetId),
    b: removeLeaf(layout.b, widgetId),
  };
}

/**
 * Set the ratio of the split node addressed by `path` (a list of 0/1 = a/b
 * choices from the root).
 */
export function setRatioAtPath(
  layout: PlotLayout,
  path: number[],
  ratio: number,
): PlotLayout {
  if (layout.type === 'leaf') return layout;
  if (path.length === 0) return { ...layout, ratio };
  const [head, ...rest] = path;
  if (head === 0) return { ...layout, a: setRatioAtPath(layout.a, rest, ratio) };
  return { ...layout, b: setRatioAtPath(layout.b, rest, ratio) };
}
