src/utils/tree.js

/**
 * Tree manipulation utilities for the conversation graph.
 * Nodes are stored as a flat map: { [nodeId]: ConversationNode }
 */

/**
 * Get the path from root to a specific node (inclusive).
 * Returns an array of nodeIds ordered root → target.
 */
export function getAncestryPath(nodes, nodeId) {
  const path = [];
  let current = nodeId;
  while (current) {
    path.unshift(current);
    const node = nodes[current];
    if (!node) break;
    current = node.parentId;
  }
  return path;
}

/**
 * Get direct children of a node.
 */
export function getChildren(nodes, nodeId) {
  return Object.values(nodes).filter((n) => n.parentId === nodeId);
}

/**
 * Get all descendant nodeIds (BFS order).
 */
export function getDescendants(nodes, nodeId) {
  const descendants = [];
  const queue = [nodeId];
  while (queue.length > 0) {
    const current = queue.shift();
    const children = getChildren(nodes, current);
    for (const child of children) {
      descendants.push(child.id);
      queue.push(child.id);
    }
  }
  return descendants;
}

/**
 * Get the deepest leaf node in a subtree, following the first child at each level.
 * Used for finding the "default" branch path.
 */
export function getDeepestLeaf(nodes, nodeId) {
  let current = nodeId;
  while (true) {
    const children = getChildren(nodes, current);
    if (children.length === 0) return current;
    // Follow the first (oldest) child
    children.sort((a, b) => a.timestamp - b.timestamp);
    current = children[0].id;
  }
}

/**
 * Build the messages array for LLM context from root to a specific node.
 * Returns an array of { role, content } messages.
 */
export function buildConversationHistory(nodes, nodeId) {
  const path = getAncestryPath(nodes, nodeId);
  const messages = [
    {
      role: 'system',
      content:
        'You are a helpful, concise assistant. Keep responses focused and well-structured.',
    },
  ];

  for (const id of path) {
    const node = nodes[id];
    if (!node) continue;
    if (node.userMessage) {
      messages.push({ role: 'user', content: node.userMessage });
    }
    if (node.assistantMessage) {
      messages.push({ role: 'assistant', content: node.assistantMessage });
    }
  }
  return messages;
}

/**
 * Build messages for generating a response at a specific node.
 * Includes all parent context but only the user message of the target node.
 */
export function buildGenerationContext(nodes, nodeId) {
  const path = getAncestryPath(nodes, nodeId);
  const messages = [
    {
      role: 'system',
      content:
        'You are a helpful, concise assistant. Keep responses focused and well-structured.',
    },
  ];

  for (let i = 0; i < path.length; i++) {
    const node = nodes[path[i]];
    if (!node) continue;
    if (node.userMessage) {
      messages.push({ role: 'user', content: node.userMessage });
    }
    // Only include assistant message for nodes before the target
    if (i < path.length - 1 && node.assistantMessage) {
      messages.push({ role: 'assistant', content: node.assistantMessage });
    }
  }
  return messages;
}

/**
 * Convert the conversation tree into react-force-graph data format.
 */
export function toGraphData(nodes, mainBranchPath = [], activeBranchPath = []) {
  const mainSet = new Set(mainBranchPath);
  const activeSet = new Set(activeBranchPath);

  const graphNodes = Object.values(nodes).map((node) => ({
    id: node.id,
    name: node.userMessage
      ? node.userMessage.slice(0, 60) + (node.userMessage.length > 60 ? '…' : '')
      : 'Start',
    userMessage: node.userMessage || '',
    assistantMessage: node.assistantMessage || '',
    status: node.status,
    isMain: mainSet.has(node.id),
    isActive: activeSet.has(node.id),
    isRoot: !node.parentId,
    childCount: getChildren(nodes, node.id).length,
    depth: getAncestryPath(nodes, node.id).length - 1,
  }));

  const graphLinks = Object.values(nodes)
    .filter((n) => n.parentId)
    .map((n) => ({
      source: n.parentId,
      target: n.id,
      isMain: mainSet.has(n.parentId) && mainSet.has(n.id),
      isActive: activeSet.has(n.parentId) && activeSet.has(n.id),
    }));

  return { nodes: graphNodes, links: graphLinks };
}