import { useEffect, useRef, useImperativeHandle, forwardRef, useCallback } from 'react';
import * as d3 from 'd3';
import type { AttributeNode } from '../types';
import '../styles/mindmap.css';

interface VisualSettings {
  nodeSpacing: number;
  fontSize: number;
}

interface MindmapD3Props {
  data: AttributeNode;
  isDark: boolean;
  visualSettings: VisualSettings;
}

interface TreeNode extends d3.HierarchyPointNode<AttributeNode> {
  _children?: TreeNode[];
  x0?: number;
  y0?: number;
}

export interface MindmapD3Ref {
  expandAll: () => void;
  collapseAll: () => void;
}

export const MindmapD3 = forwardRef<MindmapD3Ref, MindmapD3Props>(
  ({ data, isDark, visualSettings }, ref) => {
    const svgRef = useRef<SVGSVGElement>(null);
    const containerRef = useRef<HTMLDivElement>(null);
    const rootRef = useRef<TreeNode | null>(null);
    const updateFnRef = useRef<((source: TreeNode) => void) | null>(null);

    const expandAllNodes = useCallback((node: TreeNode) => {
      if (node._children) {
        node.children = node._children;
        node._children = undefined;
      }
      if (node.children) {
        node.children.forEach(expandAllNodes);
      }
    }, []);

    const collapseAllNodes = useCallback((node: TreeNode) => {
      if (node.children) {
        node.children.forEach(collapseAllNodes);
        if (node.depth > 0) {
          node._children = node.children;
          node.children = undefined;
        }
      }
    }, []);

    useImperativeHandle(ref, () => ({
      expandAll: () => {
        if (rootRef.current && updateFnRef.current) {
          expandAllNodes(rootRef.current);
          updateFnRef.current(rootRef.current);
        }
      },
      collapseAll: () => {
        if (rootRef.current && updateFnRef.current) {
          collapseAllNodes(rootRef.current);
          updateFnRef.current(rootRef.current);
        }
      },
    }), [expandAllNodes, collapseAllNodes]);

    useEffect(() => {
      if (!svgRef.current || !containerRef.current || !data) return;

      const container = containerRef.current;
      const width = container.clientWidth;
      const height = container.clientHeight;
      const margin = { top: 40, right: 20, bottom: 40, left: 20 };
      const { nodeSpacing, fontSize } = visualSettings;

      // Clear previous content
      d3.select(svgRef.current).selectAll('*').remove();

      const svg = d3
        .select(svgRef.current)
        .attr('width', width)
        .attr('height', height);

      const g = svg
        .append('g')
        .attr('transform', `translate(${margin.left}, ${margin.top})`);

      // Add zoom behavior
      const zoom = d3.zoom<SVGSVGElement, unknown>()
        .scaleExtent([0.3, 3])
        .on('zoom', (event) => {
          g.attr('transform', event.transform);
        });

      svg.call(zoom);
      svg.call(zoom.transform, d3.zoomIdentity.translate(margin.left, margin.top));

      // Create hierarchy first to calculate dynamic spacing
      const root = d3.hierarchy(data) as TreeNode;
      // For horizontal layout: x0 is vertical center, y0 is left edge
      root.x0 = height / 2;
      root.y0 = 0;
      rootRef.current = root;

      // Create horizontal tree layout (left to right)
      const verticalNodeSpacing = Math.max(nodeSpacing, 35);
      const horizontalLevelSpacing = 180; // space between levels (left to right)
      const treeLayout = d3.tree<AttributeNode>()
        .nodeSize([verticalNodeSpacing, horizontalLevelSpacing])
        .separation((a, b) => (a.parent === b.parent ? 1 : 1.2));

      // Initialize all nodes as expanded
      root.descendants().forEach((d: TreeNode) => {
        d._children = undefined;
      });

      // Dynamically extract category labels from the tree based on depth
      // Each depth level corresponds to a category
      const categoryByDepth: Record<number, string> = {};
      root.descendants().forEach((d: TreeNode) => {
        if (d.depth > 0 && d.data.category && !categoryByDepth[d.depth]) {
          categoryByDepth[d.depth] = d.data.category;
        }
      });
      const maxDepthWithCategory = Math.max(...Object.keys(categoryByDepth).map(Number), 0);
      const categoryLabels = [''];
      for (let i = 1; i <= maxDepthWithCategory; i++) {
        categoryLabels.push(categoryByDepth[i] || '');
      }
      const headerHeight = 40;

      function update(source: TreeNode) {
        const duration = 300;
        const nodes = treeLayout(root);
        const descendants = nodes.descendants() as TreeNode[];
        const links = nodes.links();

        // For horizontal layout: swap x and y
        // d.x becomes vertical position, d.y becomes horizontal position
        // Center the tree vertically (leave room for header)
        let minX = Infinity;
        let maxX = -Infinity;
        descendants.forEach((d) => {
          if (d.x < minX) minX = d.x;
          if (d.x > maxX) maxX = d.x;
        });
        const treeHeight = maxX - minX;
        const offsetY = (height - margin.top - margin.bottom - treeHeight - headerHeight) / 2 - minX + headerHeight;
        const offsetX = 50; // left margin for root node

        // Draw category headers with background
        g.selectAll('.category-header-group').remove();
        const maxDepth = Math.max(...descendants.map(d => d.depth));

        // Dynamic color palette for categories
        const colorPalette = [
          { 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
        ];

        // Generate colors dynamically based on category position
        const categoryColors: Record<string, string> = {};
        categoryLabels.forEach((label, index) => {
          if (label && index > 0) {
            const colorIndex = (index - 1) % colorPalette.length;
            categoryColors[label] = isDark ? colorPalette[colorIndex].dark : colorPalette[colorIndex].light;
          }
        });

        for (let depth = 1; depth <= Math.min(maxDepth, categoryLabels.length - 1); depth++) {
          const label = categoryLabels[depth];
          if (label) {
            const headerX = depth * horizontalLevelSpacing + offsetX;
            const headerY = offsetY + minX - 45; // Position above the first node with more space

            const headerGroup = g.append('g')
              .attr('class', 'category-header-group');

            // Background rectangle
            const textWidth = label.length * 16 + 16;
            headerGroup.append('rect')
              .attr('x', headerX - textWidth / 2)
              .attr('y', headerY - 14)
              .attr('width', textWidth)
              .attr('height', 24)
              .attr('rx', 4)
              .attr('fill', categoryColors[label] || (isDark ? '#444' : '#ddd'))
              .attr('opacity', 0.9);

            // Text label
            headerGroup.append('text')
              .attr('x', headerX)
              .attr('y', headerY)
              .attr('text-anchor', 'middle')
              .attr('dominant-baseline', 'middle')
              .attr('font-size', '13px')
              .attr('font-weight', 'bold')
              .attr('fill', '#fff')
              .text(label);
          }
        }

        // Update nodes
        const node = g
          .selectAll<SVGGElement, TreeNode>('g.node')
          .data(descendants, (d) => d.data.name + d.depth);

        // Enter new nodes - swap x,y for horizontal layout
        const nodeEnter = node
          .enter()
          .append('g')
          .attr('class', 'node')
          .attr('transform', () => `translate(${(source.y0 || 0) + offsetX},${(source.x0 || 0) + offsetY})`)
          .on('click', (_event, d: TreeNode) => {
            if (d.children) {
              d._children = d.children;
              d.children = undefined;
            } else if (d._children) {
              d.children = d._children;
              d._children = undefined;
            }
            update(d);
          });

        // Add rounded rectangle for nodes
        nodeEnter
          .append('rect')
          .attr('width', (d) => Math.max(d.data.name.length * fontSize, 60))
          .attr('height', 28)
          .attr('x', (d) => -Math.max(d.data.name.length * fontSize, 60) / 2)
          .attr('y', -14)
          .attr('rx', 6)
          .attr('ry', 6)
          .attr('class', (d) => {
            const hasChildren = d.children || d._children;
            const isCollapsed = d._children && !d.children;
            const category = d.data.category || (d.depth === 0 ? 'root' : '');
            const categoryClass = category ? `category-${category}` : '';
            return `node-rect depth-${d.depth} ${categoryClass} ${hasChildren ? 'has-children' : ''} ${isCollapsed ? 'collapsed' : ''}`;
          })
          .style('opacity', 0);

        // Add text labels inside rect
        nodeEnter
          .append('text')
          .attr('dy', '0.35em')
          .attr('text-anchor', 'middle')
          .attr('class', (d) => {
            const category = d.data.category || (d.depth === 0 ? 'root' : '');
            const categoryClass = category ? `category-${category}` : '';
            return `node-text depth-${d.depth} ${categoryClass}`;
          })
          .style('font-size', `${fontSize}px`)
          .text((d) => d.data.name)
          .style('opacity', 0);

        // Merge enter and update
        const nodeUpdate = nodeEnter.merge(node);

        // Horizontal layout: y is horizontal (depth), x is vertical (spread)
        nodeUpdate
          .transition()
          .duration(duration)
          .attr('transform', (d) => `translate(${d.y + offsetX},${d.x + offsetY})`);

        nodeUpdate
          .select('rect')
          .transition()
          .duration(duration)
          .attr('width', (d) => Math.max(d.data.name.length * fontSize, 60))
          .attr('x', (d) => -Math.max(d.data.name.length * fontSize, 60) / 2)
          .style('opacity', 1)
          .attr('class', (d) => {
            const hasChildren = d.children || d._children;
            const isCollapsed = d._children && !d.children;
            const category = d.data.category || (d.depth === 0 ? 'root' : '');
            const categoryClass = category ? `category-${category}` : '';
            return `node-rect depth-${d.depth} ${categoryClass} ${hasChildren ? 'has-children' : ''} ${isCollapsed ? 'collapsed' : ''}`;
          });

        nodeUpdate
          .select('text')
          .transition()
          .duration(duration)
          .style('font-size', `${fontSize}px`)
          .style('opacity', 1);

        // Exit nodes
        const nodeExit = node
          .exit<TreeNode>()
          .transition()
          .duration(duration)
          .attr('transform', () => `translate(${source.y + offsetX},${source.x + offsetY})`)
          .remove();

        nodeExit.select('rect').style('opacity', 0);
        nodeExit.select('text').style('opacity', 0);

        // Update links - horizontal curves (left to right)
        const link = g
          .selectAll<SVGPathElement, d3.HierarchyPointLink<AttributeNode>>('path.link')
          .data(links, (d) => d.target.data.name + (d.target as TreeNode).depth);

        const linkEnter = link
          .enter()
          .insert('path', 'g')
          .attr('class', 'link')
          .attr('d', () => {
            const o = { x: (source.y0 || 0) + offsetX, y: (source.x0 || 0) + offsetY };
            return `M${o.x},${o.y} L${o.x},${o.y}`;
          });

        linkEnter
          .merge(link)
          .transition()
          .duration(duration)
          .attr('d', (d) => {
            // Horizontal layout: source.y is horizontal, source.x is vertical
            const sx = d.source.y + offsetX;
            const sy = d.source.x + offsetY;
            const tx = d.target.y + offsetX;
            const ty = d.target.x + offsetY;
            const midX = (sx + tx) / 2;
            return `M${sx},${sy} C${midX},${sy} ${midX},${ty} ${tx},${ty}`;
          });

        link
          .exit()
          .transition()
          .duration(duration)
          .attr('d', () => {
            const o = { x: source.y + offsetX, y: source.x + offsetY };
            return `M${o.x},${o.y} L${o.x},${o.y}`;
          })
          .remove();

        // Store positions for next transition (keep original x,y for d3 tree)
        descendants.forEach((d) => {
          d.x0 = d.x;
          d.y0 = d.y;
        });
      }

      updateFnRef.current = update;
      update(root);

      // Handle resize
      const resizeObserver = new ResizeObserver(() => {
        const newWidth = container.clientWidth;
        const newHeight = container.clientHeight;
        svg.attr('width', newWidth).attr('height', newHeight);
      });

      resizeObserver.observe(container);

      return () => {
        resizeObserver.disconnect();
      };
    }, [data, visualSettings, isDark]);

    return (
      <div
        ref={containerRef}
        className={`mindmap-container ${isDark ? 'mindmap-dark' : 'mindmap-light'}`}
      >
        <svg ref={svgRef} className="mindmap-svg" />
      </div>
    );
  }
);

MindmapD3.displayName = 'MindmapD3';
