import { LinkVertical, LinkVerticalStep, LinkVerticalLine } from "@visx/shape";
import { DEFAULT_STYLES, TreeNode } from "./Interface";


export function collectExpandedNodeIds(
  node: TreeNode,
  expandedIds: Set<number | string> = new Set()
): Set<number | string> {
  const isChildrenAvailable = Array.isArray(node.children) && node.children.length > 0;
  if (isChildrenAvailable) {
    expandedIds.add(node.id);
    node.children!.forEach((childNode) => collectExpandedNodeIds(childNode, expandedIds));
  }
  return expandedIds;
}

export function getLinkComponent(linkType: string) {
  if (linkType === "step") return LinkVerticalStep;
  if (linkType === "line") return LinkVerticalLine;
  return LinkVertical;
}

export function getLevelStyle(
  levelStyles: Record<string, any>,
  depth: number
): any {
  const key = `level${depth}`;
  if (levelStyles[key]) return levelStyles[key];
  return DEFAULT_STYLES;
}

export function extractLevelStyles(
  styleConfig: Record<string, any>
): Record<string, any> {
  const levelStyles: Record<string, any> = {};
  Object.keys(styleConfig).forEach((key) => {
    if (/^level\d+$/.test(key) && styleConfig[key] && typeof styleConfig[key] === "object") {
      levelStyles[key] = styleConfig[key];
    }
  });
  return levelStyles;
}

export function mergeChildren(
  node: TreeNode,
  targetId: number | string,
  children: TreeNode[]
): TreeNode {
  if (node.id === targetId) return { ...node, children };
  if (!Array.isArray(node.children) || node.children.length === 0) return node;
  return {
    ...node,
    children: node.children.map((c) => mergeChildren(c, targetId, children)),
  };
}

export function countVisibleLeaves(
  node: TreeNode,
  expandedIds: Set<number | string>
): number {
  const isLeaf =
    !expandedIds.has(node.id) ||
    !Array.isArray(node.children) ||
    node.children.length === 0;
  if (isLeaf) return 1;
  return node.children!.reduce(
    (sum, child) => sum + countVisibleLeaves(child, expandedIds),
    0
  );
}

export function computeDynamicHeight(
  node: TreeNode,
  expandedIds: Set<number | string>,
  nodeHeight: number,
  nodeHeights: Record<number | string, number> = {},
  levelGap: number = 80,
): number {
  const maxHeightPerDepth: Record<number, number> = {};

  function traverse(currentNode: TreeNode, depth: number) {
    const measured = nodeHeights[currentNode.id];
    const effectiveHeight = Math.max(measured && measured > 0 ? measured : 0, nodeHeight);
    maxHeightPerDepth[depth] = Math.max(maxHeightPerDepth[depth] ?? 0, effectiveHeight);
    if (expandedIds.has(currentNode.id) && Array.isArray(currentNode.children)) {
      currentNode.children.forEach((c) => traverse(c, depth + 1));
    }
  }
  traverse(node, 0);
  const depths = Object.keys(maxHeightPerDepth);
  if (depths.length === 0) return nodeHeight + levelGap;
  return depths.reduce((total, d) => total + maxHeightPerDepth[Number(d)] + levelGap, 0);
}