import { Node, Edge } from "@xyflow/react";

interface LayoutOptions {
    nodes: Node[];
    edges: Edge[];
    levelSeparation?: number;
    nodeSeparation?: number;
}

interface NodeWithLayout extends Node {
    width?: number;
    height?: number;
    children?: NodeWithLayout[];
    parent?: NodeWithLayout;
    subtreeHeight?: number;
    isRight?: boolean;
}

export function getMindMapLayout(options: LayoutOptions) {
    const {
        nodes,
        edges,
        levelSeparation = 200,
        nodeSeparation = 60
    } = options;

    // 构建树形结构
    const nodeMap = new Map<string, NodeWithLayout>();
    nodes.forEach(node => {
        nodeMap.set(node.id, { ...node, children: [], width: 150, height: 40 });
    });

    let rootNode: NodeWithLayout | undefined;
    edges.forEach(edge => {
        const source = nodeMap.get(edge.source);
        const target = nodeMap.get(edge.target);
        if (source && target) {
            source.children?.push(target);
            target.parent = source;
        }
    });

    // 找到根节点
    rootNode = Array.from(nodeMap.values()).find(node => !node.parent);
    if (!rootNode) return { nodes, edges };

    // 分配节点到左右两侧
    function assignSides(node: NodeWithLayout, isRight: boolean = true) {
        if (!node.children?.length) return;

        const len = node.children.length;
        const midIndex = Math.floor(len / 2);

        // 如果是根节点，将子节点分为左右两部分
        if (!node.parent) {
            for (let i = 0; i < len; i++) {
                const child = node.children[i];
                assignSides(child, i < midIndex);
                child.isRight = i < midIndex;
            }
        }
        // 如果不是根节点，所有子节点继承父节点的方向
        else {
            node.children.forEach(child => {
                assignSides(child, isRight);
                child.isRight = isRight;
            });
        }
    }

    // 计算子树高度
    function calculateSubtreeHeight(node: NodeWithLayout): number {
        if (!node.children?.length) {
            node.subtreeHeight = node.height || 40;
            return node.subtreeHeight;
        }

        const childrenHeight = node.children.reduce((sum, child) => {
            return sum + calculateSubtreeHeight(child);
        }, 0);

        const totalGaps = (node.children.length - 1) * nodeSeparation;
        node.subtreeHeight = Math.max(node.height || 40, childrenHeight + totalGaps);
        return node.subtreeHeight;
    }

    // 布局计算
    function calculateLayout(node: NodeWithLayout, x: number, y: number) {
        node.position = { x, y };
        if (!node.children?.length) return;

        let currentY = y - (node.subtreeHeight || 0) / 2;

        node.children.forEach(child => {
            const direction = child.isRight ? 1 : -1;
            const childX = x + (levelSeparation * direction);
            const childY = currentY + (child.subtreeHeight || 0) / 2;

            calculateLayout(child, childX, childY);
            currentY += (child.subtreeHeight || 0) + nodeSeparation;
        });
    }

    // 执行布局流程
    if (rootNode) {
        // 1. 分配节点到左右两侧
        assignSides(rootNode);
        // 2. 计算子树高度
        calculateSubtreeHeight(rootNode);
        // 3. 执行布局计算
        calculateLayout(rootNode, 0, 0);
    }

    // 转换回原始格式
    const layoutedNodes = Array.from(nodeMap.values()).map(node => ({
        ...node,
        position: node.position,
    }));

    return {
        nodes: layoutedNodes,
        edges,
    };
}