windmill/frontend/src/lib/components/flows/flowDiff.ts

import type { FlowModule, FlowValue } from '$lib/gen'
import { deepEqual } from 'fast-equals'
import {
	collectAllFlowModuleIds,
	collectFlowNodes,
	ensureModuleArrayByLocation,
	findModuleParent,
	getModuleArrayByLocation,
	type ModuleParentLocation
} from './flowTree'

export { findModuleInFlow, findModuleParent, type ModuleParentLocation } from './flowTree'

/** Prefix added to module IDs when the original module coexists with a replacement */
export const DUPLICATE_MODULE_PREFIX = 'old__'

/** Prefix added to new module IDs when restoring original during type change rejection */
export const NEW_MODULE_PREFIX = 'new__'

/**
 * Action types for flow module changes during diff tracking
 * - added: Module was added to the flow
 * - modified: Module content was changed
 * - removed: Module was deleted from the flow
 * - shadowed: Module is shown as removed (visualization mode)
 */
export type AIModuleAction = 'added' | 'modified' | 'removed' | 'shadowed' | undefined

/**
 * Tracks the action performed on a module and whether it requires user approval
 */
export type ModuleActionInfo = {
	action: AIModuleAction
	/** Whether this change is pending user approval (accept/reject) */
	pending: boolean
}

/**
 * Normalizes a FlowModule for comparison by removing properties that
 * should be ignored when determining if a module has changed.
 * Specifically, removes empty `assets` arrays since their presence/absence
 * is not a meaningful difference. Recursively normalizes nested modules
 * within container types (branchone, branchall, forloopflow, whileloopflow, aiagent).
 */
function normalizeModuleForComparison(module: FlowModule): FlowModule {
	// Deep clone to avoid mutating the original and to handle nested structures
	const normalized = JSON.parse(JSON.stringify(module)) as FlowModule

	// Helper to remove empty assets from a value object
	function removeEmptyAssets(value: Record<string, unknown>): void {
		if (Array.isArray(value.assets) && value.assets.length === 0) {
			delete value.assets
		}
	}

	if ('value' in normalized && normalized.value && typeof normalized.value === 'object') {
		removeEmptyAssets(normalized.value as Record<string, unknown>)

		// Recursively normalize nested modules based on type
		const value = normalized.value
		if (value.type === 'forloopflow' || value.type === 'whileloopflow') {
			value.modules = value.modules.map((m) => normalizeModuleForComparison(m))
		} else if (value.type === 'branchone') {
			value.default = value.default.map((m) => normalizeModuleForComparison(m))
			value.branches = value.branches.map((branch) => ({
				...branch,
				modules: branch.modules.map((m) => normalizeModuleForComparison(m))
			}))
		} else if (value.type === 'branchall') {
			value.branches = value.branches.map((branch) => ({
				...branch,
				modules: branch.modules.map((m) => normalizeModuleForComparison(m))
			}))
		}
	}

	return normalized
}

/**
 * The complete diff result with action maps and merged flow
 */
export type FlowTimeline = {
	/** Actions for modules in the before flow */
	beforeActions: Record<string, ModuleActionInfo>

	/** Actions for modules in the after flow (adjusted based on display mode) */
	afterActions: Record<string, ModuleActionInfo>

	/** The merged flow containing both after modules and removed modules properly nested */
	mergedFlow: FlowValue
}

/**
 * Computes the difference between two flow versions and returns a map of module IDs to their actions.
 *
 * When a module exists in both flows but has a different type, it's treated as removed + added
 * rather than modified, since it's effectively a completely different module.
 *
 * @param beforeFlow - The original flow value
 * @param afterFlow - The modified flow value
 * @returns A record mapping module IDs to their diff result with separate before/after actions
 */
export function computeFlowModuleDiff(
	beforeFlow: FlowValue,
	afterFlow: FlowValue,
	options: { markAsPending: boolean } = { markAsPending: false }
): {
	beforeActions: Record<string, ModuleActionInfo>
	afterActions: Record<string, ModuleActionInfo>
} {
	const beforeActions: Record<string, ModuleActionInfo> = {}
	const afterActions: Record<string, ModuleActionInfo> = {}

	// Get all modules with their locations from both flows
	const beforeModulesWithLoc = getAllModulesWithLocation(beforeFlow)
	const afterModulesWithLoc = getAllModulesWithLocation(afterFlow)

	// Find all module IDs
	const allModuleIds = new Set([...beforeModulesWithLoc.keys(), ...afterModulesWithLoc.keys()])

	for (const moduleId of allModuleIds) {
		const beforeEntry = beforeModulesWithLoc.get(moduleId)
		const afterEntry = afterModulesWithLoc.get(moduleId)

		if (!beforeEntry && afterEntry) {
			// Module exists in after but not before -> added
			afterActions[moduleId] = { action: 'added', pending: options.markAsPending }
		} else if (beforeEntry && !afterEntry) {
			// Module exists in before but not after -> removed
			beforeActions[moduleId] = { action: 'removed', pending: options.markAsPending }
			afterActions[moduleId] = { action: 'shadowed', pending: options.markAsPending }
		} else if (beforeEntry && afterEntry) {
			// Module exists in both -> check location first, then type and content
			if (!locationsEqual(beforeEntry.location, afterEntry.location)) {
				// Location changed -> treat as removed from old + added at new
				beforeActions[moduleId] = { action: 'removed', pending: options.markAsPending }
				afterActions[moduleId] = { action: 'added', pending: options.markAsPending }
			} else {
				// Same location -> check type and content
				const typeChanged = beforeEntry.module.value.type !== afterEntry.module.value.type
				if (typeChanged) {
					// Type changed -> treat as removed + added
					beforeActions[moduleId] = { action: 'removed', pending: options.markAsPending }
					afterActions[moduleId] = { action: 'added', pending: options.markAsPending }
				} else if (
					!deepEqual(
						normalizeModuleForComparison(beforeEntry.module),
						normalizeModuleForComparison(afterEntry.module)
					)
				) {
					// Same type but different content -> modified
					beforeActions[moduleId] = { action: 'modified', pending: options.markAsPending }
					afterActions[moduleId] = { action: 'modified', pending: options.markAsPending }
				}
			}
		}
	}

	return { beforeActions, afterActions }
}

/**
 * Helper function to get all modules from a flow as a Map
 */
function getAllModulesMap(flow: FlowValue): Map<string, FlowModule> {
	const moduleMap = new Map<string, FlowModule>()

	for (const { module } of collectFlowNodes(flow)) {
		if (module?.id) {
			moduleMap.set(module.id, module)
		}
	}

	return moduleMap
}

/**
 * Represents a module along with its location in the flow
 */
type ModuleWithLocation = {
	module: FlowModule
	location: ModuleParentLocation
}

/**
 * Helper function to get all modules from a flow as a Map with their locations
 */
function getAllModulesWithLocation(flow: FlowValue): Map<string, ModuleWithLocation> {
	const result = new Map<string, ModuleWithLocation>()

	for (const entry of collectFlowNodes(flow)) {
		if (entry.module?.id) {
			result.set(entry.module.id, {
				module: entry.module,
				location: entry.location
			})
		}
	}

	return result
}

/**
 * Compares two module locations for equality.
 * Two locations are equal if they refer to the same parent container.
 * Index within the container is not considered (modules can be reordered).
 */
export function locationsEqual(
	a: ModuleParentLocation | null,
	b: ModuleParentLocation | null
): boolean {
	if (!a || !b) return a === b
	if (a.type !== b.type) return false

	switch (a.type) {
		case 'root':
		case 'failure':
		case 'preprocessor':
			return true // Same type is enough (index doesn't matter for location equality)
		case 'forloop':
		case 'whileloop':
		case 'aiagent':
			return a.parentId === (b as typeof a).parentId
		case 'branchone-default':
			return a.parentId === (b as typeof a).parentId
		case 'branchone-branch':
		case 'branchall-branch':
			return (
				a.parentId === (b as typeof a).parentId && a.branchIndex === (b as typeof a).branchIndex
			)
		default:
			return false
	}
}

/**
 * Deep clones a module to avoid mutation
 */
function cloneModule(module: FlowModule): FlowModule {
	return JSON.parse(JSON.stringify(module))
}

/**
 * Prepends a prefix to a module's ID and all nested child module IDs to avoid collisions
 */
function prependModuleId(module: FlowModule, prefix: string): FlowModule {
	const newModule = cloneModule(module)
	newModule.id = prefix + newModule.id

	// Recursively prefix nested module IDs
	if (newModule.value.type === 'forloopflow' || newModule.value.type === 'whileloopflow') {
		newModule.value.modules = newModule.value.modules.map((m) => prependModuleId(m, prefix))
	} else if (newModule.value.type === 'branchone') {
		newModule.value.default = newModule.value.default.map((m) => prependModuleId(m, prefix))
		newModule.value.branches = newModule.value.branches.map((branch) => ({
			...branch,
			modules: branch.modules.map((m) => prependModuleId(m, prefix))
		}))
	} else if (newModule.value.type === 'branchall') {
		newModule.value.branches = newModule.value.branches.map((branch) => ({
			...branch,
			modules: branch.modules.map((m) => prependModuleId(m, prefix))
		}))
	} else if (newModule.value.type === 'aiagent' && newModule.value.tools) {
		// Handle aiagent tools - only prefix FlowModule tools, not MCP tools
		newModule.value.tools = newModule.value.tools.map((tool) => {
			// MCP tools have tool_type: 'mcp', FlowModule tools have tool_type: 'flowmodule' or undefined
			if (tool.value.tool_type === 'mcp') {
				return tool // MCP tools don't have nested module IDs
			}
			// For FlowModule tools, prefix the ID and recurse
			const prefixedTool = {
				...tool,
				id: prefix + tool.id
			}
			// If the tool has nested modules (it's a container type), recurse
			const innerValue = tool.value as FlowModule['value']
			if (innerValue.type === 'forloopflow' || innerValue.type === 'whileloopflow') {
				;(prefixedTool.value as any).modules = (innerValue as any).modules.map((m: FlowModule) =>
					prependModuleId(m, prefix)
				)
			} else if (innerValue.type === 'branchone') {
				;(prefixedTool.value as any).default = (innerValue as any).default.map((m: FlowModule) =>
					prependModuleId(m, prefix)
				)
				;(prefixedTool.value as any).branches = (innerValue as any).branches.map((branch: any) => ({
					...branch,
					modules: branch.modules.map((m: FlowModule) => prependModuleId(m, prefix))
				}))
			} else if (innerValue.type === 'branchall') {
				;(prefixedTool.value as any).branches = (innerValue as any).branches.map((branch: any) => ({
					...branch,
					modules: branch.modules.map((m: FlowModule) => prependModuleId(m, prefix))
				}))
			}
			return prefixedTool
		})
	}

	return newModule
}

function getAllModuleIds(flow: FlowValue): Set<string> {
	return new Set(collectAllFlowModuleIds(flow))
}

/**
 * Scans the merged flow for duplicate IDs and prefixes duplicates with 'old__'.
 * This handles the case where a module is moved from one location to another -
 * both the old and new versions end up in the merged flow with the same ID.
 */
function fixDuplicateIds(merged: FlowValue, beforeFlow: FlowValue): void {
	const seenIds = new Set<string>()
	const beforeModulesMap = getAllModulesMap(beforeFlow)

	// Process a single module - returns the (possibly prefixed) module
	function processModule(module: FlowModule): FlowModule {
		let result = module

		if (seenIds.has(module.id)) {
			// Duplicate found! Check if this one exists in beforeFlow
			const beforeModule = beforeModulesMap.get(module.id)
			if (beforeModule) {
				// This is the "old" version - prefix it and all its children
				result = prependModuleId(module, DUPLICATE_MODULE_PREFIX)
			}
		} else {
			seenIds.add(module.id)
		}

		// Recurse into nested modules (use result which may be prefixed)
		processNestedModules(result)

		return result
	}

	// Process nested modules in-place
	function processNestedModules(module: FlowModule): void {
		if (module.value.type === 'forloopflow' || module.value.type === 'whileloopflow') {
			module.value.modules = module.value.modules.map((m) => processModule(m))
		} else if (module.value.type === 'branchone') {
			module.value.default = module.value.default.map((m) => processModule(m))
			for (const branch of module.value.branches) {
				branch.modules = branch.modules.map((m) => processModule(m))
			}
		} else if (module.value.type === 'branchall') {
			for (const branch of module.value.branches) {
				branch.modules = branch.modules.map((m) => processModule(m))
			}
		} else if (module.value.type === 'aiagent' && module.value.tools) {
			// For aiagent tools, we need to track IDs of FlowModule tools
			for (const tool of module.value.tools) {
				if (tool.value.tool_type !== 'mcp') {
					if (seenIds.has(tool.id)) {
						// Can't easily prefix in-place here, but aiagent tools rarely move
						// The main use case is regular modules moving in/out of loops/branches
					} else {
						seenIds.add(tool.id)
					}
				}
			}
		}
	}

	// Process root modules
	if (merged.modules) {
		merged.modules = merged.modules.map((m) => processModule(m))
	}

	// Process special modules
	if (merged.failure_module) {
		if (seenIds.has(merged.failure_module.id)) {
			const beforeModule = beforeModulesMap.get(merged.failure_module.id)
			if (beforeModule) {
				merged.failure_module = prependModuleId(merged.failure_module, DUPLICATE_MODULE_PREFIX)
			}
		} else {
			seenIds.add(merged.failure_module.id)
		}
	}

	if (merged.preprocessor_module) {
		if (seenIds.has(merged.preprocessor_module.id)) {
			const beforeModule = beforeModulesMap.get(merged.preprocessor_module.id)
			if (beforeModule) {
				merged.preprocessor_module = prependModuleId(
					merged.preprocessor_module,
					DUPLICATE_MODULE_PREFIX
				)
			}
		} else {
			seenIds.add(merged.preprocessor_module.id)
		}
	}
}

/**
 * Reconstructs the merged flow with removed modules properly nested
 */
function reconstructMergedFlow(
	afterFlow: FlowValue,
	beforeFlow: FlowValue,
	beforeActions: Record<string, ModuleActionInfo>
): FlowValue {
	// Deep clone afterFlow to avoid mutation
	const merged: FlowValue = JSON.parse(JSON.stringify(afterFlow))

	// Get all removed/shadowed modules from beforeFlow
	const removedModules = Object.entries(beforeActions)
		.filter(([_, action]) => action.action === 'removed' || action.action === 'shadowed')
		.map(([id]) => id)

	// Create a Set for faster lookup
	const removedModulesSet = new Set(removedModules)

	// Cache beforeFlow modules map and merged IDs to avoid recomputing in the loop
	const beforeModulesMap = getAllModulesMap(beforeFlow)
	const mergedIds = getAllModuleIds(merged)

	// For each removed module, find its parent and insert it
	for (const removedId of removedModules) {
		const beforeModule = beforeModulesMap.get(removedId)
		if (!beforeModule) continue

		const parentLocation = findModuleParent(beforeFlow, removedId)
		if (!parentLocation) continue

		// Skip if parent is also removed - the module will be inserted as part of its parent
		// This prevents duplicates when removing container modules with nested children
		if (
			parentLocation.type !== 'root' &&
			parentLocation.type !== 'failure' &&
			parentLocation.type !== 'preprocessor'
		) {
			if (removedModulesSet.has(parentLocation.parentId)) {
				// Parent is also removed, skip this module
				continue
			}
		}

		let clonedModule = cloneModule(beforeModule)

		// Check for ID collision - this happens when a module type changed
		// In this case, the new module is already in the merged flow as 'added'
		// We prepend the duplicate prefix to the removed module's ID so both can coexist
		if (mergedIds.has(clonedModule.id)) {
			clonedModule = prependModuleId(clonedModule, DUPLICATE_MODULE_PREFIX)
		}

		// Track the newly added module ID
		mergedIds.add(clonedModule.id)

		insertModuleIntoFlow(merged, clonedModule, beforeFlow, removedId)
	}

	// Post-process: fix any duplicate IDs that may have been created
	// This handles the case where a module moved from one location to another
	fixDuplicateIds(merged, beforeFlow)

	return merged
}

/**
 * Finds the best position to insert a removed module in a module array
 */
function findBestInsertPosition(
	targetModules: FlowModule[],
	beforeModules: FlowModule[],
	originalIndex: number,
	removedId: string
): number {
	// Look for anchors (modules that exist in both before and after)
	// Try to find previous anchor
	for (let i = originalIndex - 1; i >= 0; i--) {
		const anchorId = beforeModules[i]?.id
		const anchorIndex = targetModules.findIndex((m) => m.id === anchorId)
		if (anchorIndex >= 0) {
			return anchorIndex + 1
		}
	}

	// Try to find next anchor
	for (let i = originalIndex + 1; i < beforeModules.length; i++) {
		const anchorId = beforeModules[i]?.id
		const anchorIndex = targetModules.findIndex((m) => m.id === anchorId)
		if (anchorIndex >= 0) {
			return anchorIndex
		}
	}

	// No anchors found, append to end
	return targetModules.length
}

/**
 * Inserts a removed module into its nested parent in the merged flow
 */
/**
 * Adjusts the after actions based on display mode and adds entries for prefixed IDs
 */
function adjustActionsForDisplay(
	afterActions: Record<string, ModuleActionInfo>,
	beforeActions: Record<string, ModuleActionInfo>,
	markRemovedAsShadowed: boolean,
	mergedFlow: FlowValue
): Record<string, ModuleActionInfo> {
	const adjusted: Record<string, ModuleActionInfo> = {}

	// Copy all existing actions
	for (const [id, action] of Object.entries(afterActions)) {
		if (!markRemovedAsShadowed && action.action === 'shadowed') {
			// In unified mode, change 'shadowed' to 'removed' for proper coloring
			adjusted[id] = { action: 'removed', pending: action.pending }
		} else {
			adjusted[id] = action
		}
	}

	// Add entries for prefixed IDs (modules that had type changes or were removed)
	// These are the old versions that got the duplicate prefix prepended to their ID
	const allMergedIds = getAllModuleIds(mergedFlow)
	for (const id of allMergedIds) {
		if (id.startsWith(DUPLICATE_MODULE_PREFIX) && !adjusted[id]) {
			// This is a prefixed ID for a module that was removed
			const originalId = id.substring(DUPLICATE_MODULE_PREFIX.length)
			// Check beforeActions to see if this module was removed
			if (beforeActions[originalId]?.action === 'removed') {
				adjusted[id] = {
					action: markRemovedAsShadowed ? 'shadowed' : 'removed',
					pending: beforeActions[originalId].pending
				}
			}
		}
	}

	return adjusted
}

/**
 * Builds the complete flow diff result with action maps and merged flow.
 * The merged flow contains all modules from afterFlow plus removed modules from
 * beforeFlow properly nested in their original locations.
 *
 * @param beforeFlow - The original flow value
 * @param afterFlow - The modified flow value
 * @param options - Display options
 * @returns Complete diff result with beforeActions, afterActions, and mergedFlow
 */
export function buildFlowTimeline(
	beforeFlow: FlowValue,
	afterFlow: FlowValue,
	options: { markRemovedAsShadowed: boolean; markAsPending: boolean } = {
		markRemovedAsShadowed: false,
		markAsPending: false
	}
): FlowTimeline {
	// Compute the diff between the two flows
	const { beforeActions, afterActions } = computeFlowModuleDiff(beforeFlow, afterFlow, {
		markAsPending: options.markAsPending
	})

	// Reconstruct merged flow with removed modules properly nested
	const mergedFlow = reconstructMergedFlow(afterFlow, beforeFlow, beforeActions)

	// Adjust after actions based on display mode and add entries for prefixed IDs
	const adjustedAfterActions = adjustActionsForDisplay(
		afterActions,
		beforeActions,
		options.markRemovedAsShadowed,
		mergedFlow
	)

	return {
		beforeActions,
		afterActions: adjustedAfterActions,
		mergedFlow
	}
}

/**
 * Inserts a module into a flow at its correct position based on where it was located in the source flow.
 * This is useful when restoring a removed module - it finds the correct parent and position.
 *
 * @param targetFlow - The flow to insert the module into
 * @param moduleToInsert - The module to insert
 * @param sourceFlow - The flow where the module originally existed (to find parent location and ordering)
 * @param moduleId - The ID of the module being inserted
 */
export function insertModuleIntoFlow(
	targetFlow: FlowValue,
	moduleToInsert: FlowModule,
	sourceFlow: FlowValue,
	moduleId: string
): void {
	const parentLocation = findModuleParent(sourceFlow, moduleId)
	if (!parentLocation) return

	// Handle special modules
	if (parentLocation.type === 'failure') {
		targetFlow.failure_module = moduleToInsert
		return
	}
	if (parentLocation.type === 'preprocessor') {
		targetFlow.preprocessor_module = moduleToInsert
		return
	}

	const targetModules = ensureModuleArrayByLocation(targetFlow, parentLocation, sourceFlow)
	const sourceModules = getModuleArrayByLocation(sourceFlow, parentLocation)
	if (!targetModules || !sourceModules) {
		return
	}

	const insertIndex = findBestInsertPosition(
		targetModules,
		sourceModules,
		parentLocation.index,
		moduleId
	)
	targetModules.splice(insertIndex, 0, moduleToInsert)
}

/**
 * Checks if the input schema has changed between two flow versions.
 * The input schema always exists (even if empty), so we only check for modifications.
 *
 * @param beforeFlow - The original flow (can be OpenFlow or just have schema property)
 * @param afterFlow - The modified flow (can be OpenFlow or just have schema property)
 * @returns true if the schemas are different, false if identical
 */
export function hasInputSchemaChanged(
	beforeFlow: { schema?: { [key: string]: unknown } } | undefined,
	afterFlow: { schema?: { [key: string]: unknown } } | undefined
): boolean {
	if (!beforeFlow || !afterFlow) {
		return false
	}

	return !deepEqual(beforeFlow.schema, afterFlow.schema)
}