backend/llama.cpp/tools/ui/src/lib/utils/branching.ts

/**
 * Message branching utilities for conversation tree navigation.
 *
 * Conversation branching allows users to edit messages and create alternate paths
 * while preserving the original conversation flow. Each message has parent/children
 * relationships forming a tree structure.
 *
 * Example tree:
 * root
 *  ├── message 1 (user)
 *  │      └── message 2 (assistant)
 *  │             ├── message 3 (user)
 *  │             └── message 6 (user) ← new branch
 *  └── message 4 (user)
 *        └── message 5 (assistant)
 */

import { MessageRole } from '$lib/enums';

/**
 * Finds a message by its ID in the given messages array.
 */
export function findMessageById(
	messages: readonly DatabaseMessage[],
	id: string | null | undefined
): DatabaseMessage | undefined {
	if (!id) return undefined;
	return messages.find((m) => m.id === id);
}

/**
 * Filters messages to get the conversation path from root to a specific leaf node.
 * If the leafNodeId doesn't exist, returns the path with the latest timestamp.
 *
 * @param messages - All messages in the conversation
 * @param leafNodeId - The target leaf node ID to trace back from
 * @param includeRoot - Whether to include root messages in the result
 * @returns Array of messages from root to leaf, sorted by timestamp
 */
export function filterByLeafNodeId(
	messages: readonly DatabaseMessage[],
	leafNodeId: string,
	includeRoot: boolean = false
): readonly DatabaseMessage[] {
	const result: DatabaseMessage[] = [];
	const nodeMap = new Map<string, DatabaseMessage>();

	// Build node map for quick lookups
	for (const msg of messages) {
		nodeMap.set(msg.id, msg);
	}

	// Find the starting node (leaf node or latest if not found)
	let startNode: DatabaseMessage | undefined = nodeMap.get(leafNodeId);
	if (!startNode) {
		// If leaf node not found, use the message with latest timestamp
		let latestTime = -1;
		for (const msg of messages) {
			if (msg.timestamp > latestTime) {
				startNode = msg;
				latestTime = msg.timestamp;
			}
		}
	}

	// Traverse from leaf to root, collecting messages
	let currentNode: DatabaseMessage | undefined = startNode;
	while (currentNode) {
		// Include message if it's not root, or if we want to include root
		if (currentNode.type !== 'root' || includeRoot) {
			result.push(currentNode);
		}

		// Stop traversal if parent is null (reached root)
		if (currentNode.parent === null) {
			break;
		}
		currentNode = nodeMap.get(currentNode.parent);
	}

	// Sort: system messages first, then by timestamp
	result.sort((a, b) => {
		if (a.role === MessageRole.SYSTEM && b.role !== MessageRole.SYSTEM) return -1;
		if (a.role !== MessageRole.SYSTEM && b.role === MessageRole.SYSTEM) return 1;

		return a.timestamp - b.timestamp;
	});
	return result;
}

/**
 * Finds the leaf node (message with no children) for a given message branch.
 * Traverses down the tree following the last child until reaching a leaf.
 *
 * @param nodeMap - Map of messages keyed by ID
 * @param messageId - Starting message ID to find leaf for
 * @returns The leaf node ID, or the original messageId if no children
 */
function findLeafNodeInMap(
	nodeMap: ReadonlyMap<string, DatabaseMessage>,
	messageId: string
): string {
	let currentNode: DatabaseMessage | undefined = nodeMap.get(messageId);
	while (currentNode && currentNode.children.length > 0) {
		// Follow the last child (most recent branch)
		const lastChildId = currentNode.children[currentNode.children.length - 1];
		currentNode = nodeMap.get(lastChildId);
	}

	return currentNode?.id ?? messageId;
}

/**
 * Convenience wrapper around {@link findLeafNodeInMap} for callers that have a flat message array.
 */
export function findLeafNode(messages: readonly DatabaseMessage[], messageId: string): string {
	const nodeMap = new Map(messages.map((msg) => [msg.id, msg] as const));
	return findLeafNodeInMap(nodeMap, messageId);
}

/**
 * Finds all descendant messages (children, grandchildren, etc.) of a given message.
 * This is used for cascading deletion to remove all messages in a branch.
 *
 * @param messages - All messages in the conversation
 * @param messageId - The root message ID to find descendants for
 * @returns Array of all descendant message IDs
 */
export function findDescendantMessages(
	messages: readonly DatabaseMessage[],
	messageId: string
): string[] {
	const nodeMap = new Map<string, DatabaseMessage>();

	// Build node map for quick lookups
	for (const msg of messages) {
		nodeMap.set(msg.id, msg);
	}

	const descendants: string[] = [];
	const queue: string[] = [messageId];

	while (queue.length > 0) {
		const currentId = queue.shift()!;
		const currentNode = nodeMap.get(currentId);

		if (currentNode) {
			// Add all children to the queue and descendants list
			for (const childId of currentNode.children) {
				descendants.push(childId);
				queue.push(childId);
			}
		}
	}

	return descendants;
}

/**
 * Gets sibling information for a message, including all sibling IDs and current position.
 * Siblings are messages that share the same parent.
 *
 * @param nodeMap - Map of messages keyed by ID
 * @param messageId - The message to get sibling info for
 * @returns Sibling information including leaf node IDs for navigation
 */
export function getMessageSiblings(
	nodeMap: ReadonlyMap<string, DatabaseMessage>,
	messageId: string
): ChatMessageSiblingInfo | null {
	const message = nodeMap.get(messageId);
	if (!message) {
		return null;
	}

	// Handle null parent (root message) case
	if (message.parent === null) {
		// No parent means this is likely a root node with no siblings
		return {
			message,
			siblingIds: [messageId],
			currentIndex: 0,
			totalSiblings: 1
		};
	}

	const parentNode = nodeMap.get(message.parent);
	if (!parentNode) {
		// Parent not found - treat as single message
		return {
			message,
			siblingIds: [messageId],
			currentIndex: 0,
			totalSiblings: 1
		};
	}

	// Get all sibling IDs (including self)
	const siblingIds = parentNode.children;

	// Convert sibling message IDs to their corresponding leaf node IDs
	// This allows navigation between different conversation branches
	const siblingLeafIds = siblingIds.map((siblingId: string) =>
		findLeafNodeInMap(nodeMap, siblingId)
	);

	// Find current message's position among siblings
	const currentIndex = siblingIds.indexOf(messageId);

	return {
		message,
		siblingIds: siblingLeafIds,
		currentIndex,
		totalSiblings: siblingIds.length
	};
}

/**
 * Builds sibling information for every message in a conversation.
 *
 * @param messages - All messages in the conversation
 * @returns Map of message ID to its sibling information
 */
export function buildSiblingInfoMap(
	messages: readonly DatabaseMessage[]
): Map<string, ChatMessageSiblingInfo> {
	const nodeMap = new Map(messages.map((msg) => [msg.id, msg] as const));
	const siblingMap = new Map<string, ChatMessageSiblingInfo>();
	for (const msg of messages) {
		const info = getMessageSiblings(nodeMap, msg.id);
		if (info) {
			siblingMap.set(msg.id, info);
		}
	}
	return siblingMap;
}