import type { Edge, Node } from "reactflow";
import { Position } from "reactflow";
import type { GraphData, GraphNodeData, MindmapNode } from "../types/mindmap";

const LEVEL_GAP = 320;
const SIBLING_GAP = 120;
const PADDING_X = 80;
const PADDING_Y = 100;

type PositionMap = Map<string, { x: number; y: number }>;

type TreeToGraphOptions = {
  expandedNodeIds?: Set<string>;
  onToggleNode?: (nodeId: string) => void;
  onOpenNodeChat?: (nodeLabel: string) => void;
};

function getVisibleChildren(
  node: MindmapNode,
  expandedNodeIds: Set<string>,
): MindmapNode[] {
  if (node.children.length === 0) {
    return [];
  }

  return expandedNodeIds.has(node.id) ? node.children : [];
}

function assignPositions(
  node: MindmapNode,
  positions: PositionMap,
  leafIndexRef: { value: number },
  expandedNodeIds: Set<string>,
): number {
  const visibleChildren = getVisibleChildren(node, expandedNodeIds);

  if (visibleChildren.length === 0) {
    const y = leafIndexRef.value * SIBLING_GAP;
    leafIndexRef.value += 1;
    positions.set(node.id, { x: node.level * LEVEL_GAP, y });
    return y;
  }

  const childYs = visibleChildren.map((child) =>
    assignPositions(child, positions, leafIndexRef, expandedNodeIds),
  );
  const y =
    childYs.reduce((total, current) => total + current, 0) / childYs.length;

  positions.set(node.id, { x: node.level * LEVEL_GAP, y });
  return y;
}

function buildGraph(
  node: MindmapNode,
  positions: PositionMap,
  nodes: Node<GraphNodeData>[],
  edges: Edge[],
  expandedNodeIds: Set<string>,
  options: TreeToGraphOptions,
): void {
  const position = positions.get(node.id);

  if (!position) {
    return;
  }

  const hasChildren = node.children.length > 0;
  const isExpanded = expandedNodeIds.has(node.id);
  const onToggleNode = options.onToggleNode;
  const onOpenNodeChat = options.onOpenNodeChat;

  nodes.push({
    id: node.id,
    type: "mindmapNode",
    position: {
      x: position.x + PADDING_X,
      y: position.y + PADDING_Y,
    },
    style: {
      cursor: "default",
    },
    data: {
      label: node.label,
      level: node.level,
      isLeaf: node.is_leaf,
      hasChildren,
      isExpanded,
      onToggle: hasChildren && onToggleNode ? () => onToggleNode(node.id) : undefined,
      onOpenChat: onOpenNodeChat ? () => onOpenNodeChat(node.label) : undefined,
    },
    sourcePosition: Position.Right,
    targetPosition: Position.Left,
    draggable: false,
    selectable: false,
  });

  getVisibleChildren(node, expandedNodeIds).forEach((child) => {
    edges.push({
      id: `${node.id}-${child.id}`,
      source: node.id,
      target: child.id,
      type: "smoothstep",
      animated: false,
      style: {
        stroke: "#60a5fa",
        strokeWidth: 2,
      },
    });

    buildGraph(child, positions, nodes, edges, expandedNodeIds, options);
  });
}

export function treeToGraph(
  tree: MindmapNode,
  options: TreeToGraphOptions = {},
): GraphData {
  const positions: PositionMap = new Map();
  const leafIndexRef = { value: 0 };
  const nodes: Node<GraphNodeData>[] = [];
  const edges: Edge[] = [];
  const expandedNodeIds = options.expandedNodeIds ?? new Set([tree.id]);

  assignPositions(tree, positions, leafIndexRef, expandedNodeIds);
  buildGraph(tree, positions, nodes, edges, expandedNodeIds, options);

  return { nodes, edges };
}