/**
 * Group Store
 *
 * Manages node grouping/nesting: parent-child relationships,
 * collapse/expand state, and group membership operations.
 */

import { atom } from 'jotai';
import { graphAtom, graphUpdateVersionAtom } from './graph-store';
import { nodePositionUpdateCounterAtom } from './graph-position';
import { pushHistoryAtom } from './history-store';
import type { GraphNodeAttributes } from './types';

// =============================================================================
// Collapse State
// =============================================================================

/**
 * Set of group node IDs that are currently collapsed.
 * Collapsed groups hide their children from uiNodesAtom.
 */
export const collapsedGroupsAtom = atom<Set<string>>(new Set<string>());

/**
 * Toggle a group's collapsed state
 */
export const toggleGroupCollapseAtom = atom(null, (get, set, groupId: string) => {
	const current = get(collapsedGroupsAtom);
	const next = new Set(current);
	if (next.has(groupId)) {
		next.delete(groupId);
	} else {
		next.add(groupId);
	}
	set(collapsedGroupsAtom, next);
});

/**
 * Collapse a group
 */
export const collapseGroupAtom = atom(null, (get, set, groupId: string) => {
	const current = get(collapsedGroupsAtom);
	if (!current.has(groupId)) {
		const next = new Set(current);
		next.add(groupId);
		set(collapsedGroupsAtom, next);
	}
});

/**
 * Expand a group
 */
export const expandGroupAtom = atom(null, (get, set, groupId: string) => {
	const current = get(collapsedGroupsAtom);
	if (current.has(groupId)) {
		const next = new Set(current);
		next.delete(groupId);
		set(collapsedGroupsAtom, next);
	}
});

// =============================================================================
// Parent/Child Queries
// =============================================================================

/**
 * Get children of a group node (reads from graph attributes)
 */
export const nodeChildrenAtom = atom((get) => {
	get(graphUpdateVersionAtom);
	const graph = get(graphAtom);

	return (parentId: string): string[] => {
		const children: string[] = [];
		graph.forEachNode((nodeId, attrs) => {
			if (attrs.parentId === parentId) {
				children.push(nodeId);
			}
		});
		return children;
	};
});

/**
 * Get the parent ID of a node
 */
export const nodeParentAtom = atom((get) => {
	get(graphUpdateVersionAtom);
	const graph = get(graphAtom);

	return (nodeId: string): string | undefined => {
		if (!graph.hasNode(nodeId)) return undefined;
		return graph.getNodeAttribute(nodeId, 'parentId') as string | undefined;
	};
});

/**
 * Check if a node is a group (has any children)
 */
export const isGroupNodeAtom = atom((get) => {
	const getChildren = get(nodeChildrenAtom);
	return (nodeId: string): boolean => getChildren(nodeId).length > 0;
});

/**
 * Get count of children for a group node
 */
export const groupChildCountAtom = atom((get) => {
	const getChildren = get(nodeChildrenAtom);
	return (groupId: string): number => getChildren(groupId).length;
});

// =============================================================================
// Group Membership Operations
// =============================================================================

/**
 * Set a node's parent (move into a group)
 */
export const setNodeParentAtom = atom(
	null,
	(get, set, { nodeId, parentId }: { nodeId: string; parentId: string | undefined }) => {
		const graph = get(graphAtom);
		if (!graph.hasNode(nodeId)) return;

		// Prevent circular grouping
		if (parentId) {
			if (parentId === nodeId) return;
			// Check if parentId is a descendant of nodeId
			let current: string | undefined = parentId;
			while (current) {
				if (current === nodeId) return; // Circular!
				if (!graph.hasNode(current)) break;
				current = graph.getNodeAttribute(current, 'parentId') as string | undefined;
			}
		}

		graph.setNodeAttribute(nodeId, 'parentId', parentId);
		set(graphUpdateVersionAtom, (v) => v + 1);
	},
);

/**
 * Move multiple nodes into a group
 */
export const moveNodesToGroupAtom = atom(
	null,
	(get, set, { nodeIds, groupId }: { nodeIds: string[]; groupId: string }) => {
		for (const nodeId of nodeIds) {
			set(setNodeParentAtom, { nodeId, parentId: groupId });
		}
	},
);

/**
 * Remove a node from its group (set parentId to undefined)
 */
export const removeFromGroupAtom = atom(null, (get, set, nodeId: string) => {
	set(setNodeParentAtom, { nodeId, parentId: undefined });
});

/**
 * Create a group from selected nodes:
 * 1. Create a new "group" type node at the bounding box center
 * 2. Set all selected nodes as children
 * Returns the new group node ID (set via callback since node creation is async)
 */
export const groupSelectedNodesAtom = atom(
	null,
	(get, set, { nodeIds, groupNodeId }: { nodeIds: string[]; groupNodeId: string }) => {
		set(pushHistoryAtom, `Group ${nodeIds.length} nodes`);
		const graph = get(graphAtom);

		// Compute bounding box of selected nodes
		let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity;
		for (const nodeId of nodeIds) {
			if (!graph.hasNode(nodeId)) continue;
			const attrs = graph.getNodeAttributes(nodeId) as GraphNodeAttributes;
			minX = Math.min(minX, attrs.x);
			minY = Math.min(minY, attrs.y);
			maxX = Math.max(maxX, attrs.x + (attrs.width || 200));
			maxY = Math.max(maxY, attrs.y + (attrs.height || 100));
		}

		// Position group node at bounding box origin with padding
		const padding = 20;
		if (graph.hasNode(groupNodeId)) {
			graph.setNodeAttribute(groupNodeId, 'x', minX - padding);
			graph.setNodeAttribute(groupNodeId, 'y', minY - padding - 30); // extra for header
			graph.setNodeAttribute(groupNodeId, 'width', maxX - minX + 2 * padding);
			graph.setNodeAttribute(groupNodeId, 'height', maxY - minY + 2 * padding + 30);
		}

		// Set parent for all nodes
		for (const nodeId of nodeIds) {
			if (nodeId !== groupNodeId && graph.hasNode(nodeId)) {
				graph.setNodeAttribute(nodeId, 'parentId', groupNodeId);
			}
		}

		set(graphUpdateVersionAtom, (v) => v + 1);
		set(nodePositionUpdateCounterAtom, (c) => c + 1);
	},
);

/**
 * Ungroup: remove parent from all children of a group node
 */
export const ungroupNodesAtom = atom(null, (get, set, groupId: string) => {
	set(pushHistoryAtom, 'Ungroup nodes');
	const graph = get(graphAtom);
	graph.forEachNode((nodeId, attrs) => {
		if (attrs.parentId === groupId) {
			graph.setNodeAttribute(nodeId, 'parentId', undefined);
		}
	});
	set(graphUpdateVersionAtom, (v) => v + 1);
});

// =============================================================================
// Drag-to-Nest
// =============================================================================

/**
 * Nest dragged nodes into a target group node.
 * Sets parentId for each dragged node, then auto-resizes the target.
 */
export const nestNodesOnDropAtom = atom(
	null,
	(get, set, { nodeIds, targetId }: { nodeIds: string[]; targetId: string }) => {
		set(pushHistoryAtom, 'Nest nodes');
		for (const nodeId of nodeIds) {
			if (nodeId === targetId) continue;
			set(setNodeParentAtom, { nodeId, parentId: targetId });
		}
		set(autoResizeGroupAtom, targetId);
	},
);

// =============================================================================
// Descendants Helper
// =============================================================================

/**
 * Recursively collects all descendant node IDs of a group node.
 * Used by useNodeDrag to move all children when dragging a group.
 */
export function getNodeDescendants(
	graph: { forEachNode: (cb: (nodeId: string, attrs: { parentId?: string }) => void) => void },
	groupId: string,
): string[] {
	const descendants: string[] = [];
	const stack = [groupId];
	while (stack.length > 0) {
		const current = stack.pop()!;
		graph.forEachNode((nodeId, attrs) => {
			if (attrs.parentId === current) {
				descendants.push(nodeId);
				stack.push(nodeId);
			}
		});
	}
	return descendants;
}

// =============================================================================
// Edge Re-routing for Collapsed Groups
// =============================================================================

/**
 * Derived: maps each collapsed node to its outermost collapsed ancestor.
 * Used by visibleEdgeKeysAtom and edgeFamilyAtom to re-route edges
 * from hidden children to their visible group node.
 *
 * Only nodes that are inside a collapsed group appear in this map.
 * The value is always a visible group node (the outermost collapsed ancestor).
 */
export const collapsedEdgeRemapAtom = atom((get) => {
	const collapsed = get(collapsedGroupsAtom);
	if (collapsed.size === 0) return new Map<string, string>();

	get(graphUpdateVersionAtom);
	const graph = get(graphAtom);
	const remap = new Map<string, string>();

	for (const nodeId of graph.nodes()) {
		let current = nodeId;
		let outermost: string | null = null;
		while (true) {
			if (!graph.hasNode(current)) break;
			const parent = graph.getNodeAttribute(current, 'parentId') as string | undefined;
			if (!parent) break;
			if (collapsed.has(parent)) outermost = parent;
			current = parent;
		}
		if (outermost) remap.set(nodeId, outermost);
	}

	return remap;
});

// =============================================================================
// Group Auto-resize
// =============================================================================

/**
 * Write atom: recomputes a group node's position and dimensions
 * to contain all its direct children with padding.
 */
export const autoResizeGroupAtom = atom(null, (get, set, groupId: string) => {
	const graph = get(graphAtom);
	if (!graph.hasNode(groupId)) return;

	const children: string[] = [];
	graph.forEachNode((nodeId, attrs) => {
		if (attrs.parentId === groupId) {
			children.push(nodeId);
		}
	});

	if (children.length === 0) return;

	const padding = 20;
	const headerHeight = 30;
	let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity;

	for (const childId of children) {
		const attrs = graph.getNodeAttributes(childId) as GraphNodeAttributes;
		minX = Math.min(minX, attrs.x);
		minY = Math.min(minY, attrs.y);
		maxX = Math.max(maxX, attrs.x + (attrs.width || 200));
		maxY = Math.max(maxY, attrs.y + (attrs.height || 100));
	}

	graph.setNodeAttribute(groupId, 'x', minX - padding);
	graph.setNodeAttribute(groupId, 'y', minY - padding - headerHeight);
	graph.setNodeAttribute(groupId, 'width', maxX - minX + 2 * padding);
	graph.setNodeAttribute(groupId, 'height', maxY - minY + 2 * padding + headerHeight);

	set(nodePositionUpdateCounterAtom, (c) => c + 1);
});

// =============================================================================
// Visibility Helper
// =============================================================================

/**
 * Check if a node should be hidden because it's inside a collapsed group.
 * Walks up the parent chain — if ANY ancestor is collapsed, the node is hidden.
 */
export function isNodeCollapsed(
	nodeId: string,
	getParentId: (id: string) => string | undefined,
	collapsed: Set<string>,
): boolean {
	let current = nodeId;
	while (true) {
		const parentId = getParentId(current);
		if (!parentId) return false;
		if (collapsed.has(parentId)) return true;
		current = parentId;
	}
}
