novelty-seeking/frontend/src/components/dag/useDAGLayout.ts

import { useMemo } from 'react';
import type { Node } from '@xyflow/react';
import type { AttributeDAG, DAGNode } from '../../types';

interface LayoutConfig {
  nodeHeight: number;
  headerGap: number;
  headerHeight: number;
  nodeSpacing: number;
  fontSize: number;
  isDark: boolean;
}

const COLOR_PALETTE = [
  { dark: '#177ddc', light: '#1890ff' }, // blue
  { dark: '#854eca', light: '#722ed1' }, // purple
  { dark: '#13a8a8', light: '#13c2c2' }, // cyan
  { dark: '#d87a16', light: '#fa8c16' }, // orange
  { dark: '#49aa19', light: '#52c41a' }, // green
  { dark: '#1677ff', light: '#1890ff' }, // blue
  { dark: '#eb2f96', light: '#f759ab' }, // magenta
  { dark: '#faad14', light: '#ffc53d' }, // gold
  { dark: '#a0d911', light: '#bae637' }, // lime
];

function darken(hex: string, amount: number): string {
  const num = parseInt(hex.slice(1), 16);
  const r = Math.max(0, ((num >> 16) & 0xff) - Math.floor(255 * amount));
  const g = Math.max(0, ((num >> 8) & 0xff) - Math.floor(255 * amount));
  const b = Math.max(0, (num & 0xff) - Math.floor(255 * amount));
  return `#${((r << 16) | (g << 8) | b).toString(16).padStart(6, '0')}`;
}

export function useDAGLayout(
  data: AttributeDAG | null,
  config: LayoutConfig
): Node[] {
  return useMemo(() => {
    if (!data) return [];

    const {
      nodeHeight,
      headerGap,
      headerHeight,
      nodeSpacing,
      fontSize,
      isDark,
    } = config;

    const nodes: Node[] = [];
    const sortedCategories = [...data.categories].sort((a, b) => a.order - b.order);

    // Build category color map
    const categoryColors: Record<string, { fill: string; stroke: string }> = {};
    sortedCategories.forEach((cat, index) => {
      const paletteIndex = index % COLOR_PALETTE.length;
      const color = isDark
        ? COLOR_PALETTE[paletteIndex].dark
        : COLOR_PALETTE[paletteIndex].light;
      categoryColors[cat.name] = {
        fill: color,
        stroke: darken(color, 0.15),
      };
    });

    // Group nodes by category
    const nodesByCategory: Record<string, DAGNode[]> = {};
    for (const node of data.nodes) {
      if (!nodesByCategory[node.category]) {
        nodesByCategory[node.category] = [];
      }
      nodesByCategory[node.category].push(node);
    }

    // Sort nodes within each category by order
    for (const cat of Object.keys(nodesByCategory)) {
      nodesByCategory[cat].sort((a, b) => a.order - b.order);
    }

    // Calculate max column height for centering
    const maxNodesInColumn = Math.max(
      ...sortedCategories.map((cat) => (nodesByCategory[cat.name] || []).length),
      1
    );
    const maxTotalHeight = maxNodesInColumn * (nodeHeight + nodeSpacing) - nodeSpacing;
    const contentStartY = headerHeight + headerGap + 20; // Added extra top padding

    // Layout constants - INCREASED spacing
    const colStep = 340; // Reduced from 400
    const nodeWidth = 240;
    const headerWidth = 260;
    const groupWidth = 300; // Reduced from 360 to be less "too wide"

    // Helper function for column X position (Center of the column)
    const getColumnX = (colIndex: number) => colIndex * colStep;

    // 0. Add Group Nodes (Backgrounds) - Rendered first
    sortedCategories.forEach((cat, colIndex) => {
      const categoryNodes = nodesByCategory[cat.name] || [];
      const columnX = getColumnX(colIndex + 1); // +1 because column 0 is query

      // Calculate group height to include header and attributes
      const nodesCount = categoryNodes.length;
      const nodesHeight = Math.max(
        nodesCount * (nodeHeight + nodeSpacing) - nodeSpacing,
        0
      );

      // Start group above the header (header is at y=0)
      const groupStartY = -20;
      const groupPaddingBottom = 20;

      // Total height = distance from startY to contentStart + nodes height + bottom padding
      const groupHeight = (contentStartY - groupStartY) + nodesHeight + groupPaddingBottom;

      nodes.push({
        id: `group-${cat.name}`,
        type: 'group',
        position: {
          x: columnX - groupWidth / 2, // Centered
          y: groupStartY
        },
        data: {
          label: '', // Label is handled by header node
          color: categoryColors[cat.name]?.fill || '#666',
          width: groupWidth,
          height: Math.max(groupHeight, 100), // Ensure min height
          isDark,
        },
        draggable: false,
        selectable: false,
        zIndex: -1,
      });
    });

    // 1. Add Query Node (column 0, centered vertically)
    const queryY = contentStartY + (maxTotalHeight - nodeHeight) / 2;
    nodes.push({
      id: 'query-node',
      type: 'query',
      position: { x: getColumnX(0) - 60, y: queryY }, // Assuming query node width ~120
      data: {
        label: data.query,
        isDark,
        fontSize: fontSize + 2,
      },
      draggable: false,
      selectable: false,
    });

    // 2. Add Category Headers (starting from column 1)
    sortedCategories.forEach((cat, colIndex) => {
      const columnX = getColumnX(colIndex + 1);
      nodes.push({
        id: `header-${cat.name}`,
        type: 'categoryHeader',
        position: { x: columnX - headerWidth / 2, y: 0 }, // Centered
        data: {
          label: cat.name,
          color: categoryColors[cat.name]?.fill || '#666',
          isFixed: cat.is_fixed,
          isDark,
        },
        draggable: false,
        selectable: false,
        zIndex: 10,
      });
    });

    // 3. Add Attribute Nodes
    sortedCategories.forEach((cat, colIndex) => {
      const categoryNodes = nodesByCategory[cat.name] || [];
      const columnX = getColumnX(colIndex + 1);

      categoryNodes.forEach((node, rowIndex) => {
        const y = contentStartY + rowIndex * (nodeHeight + nodeSpacing);
        nodes.push({
          id: node.id,
          type: 'attribute',
          position: { x: columnX - nodeWidth / 2, y }, // Centered
          data: {
            label: node.name,
            fillColor: categoryColors[node.category]?.fill || '#666',
            strokeColor: categoryColors[node.category]?.stroke || '#444',
            fontSize,
          },
          draggable: false,
          selectable: false,
          zIndex: 5,
        });
      });
    });

    return nodes;
  }, [data, config]);
}