'use client';

import { useMemo } from 'react';
import { motion } from 'framer-motion';
import { cn } from '@/lib/utils';
import type { UnionFindState, UnionFindNodeState } from '@/lib/visualizations/types';
import { UnionFindNode } from '../primitives/union-find-node';

interface UnionFindViewProps {
  unionFind: UnionFindState;
  className?: string;
}

const NODE_RADIUS = 24;
const LEVEL_HEIGHT = 72;
const MIN_NODE_SPACING = 64;
const SVG_PADDING = 48;
const TREE_GAP = 40;

interface NodePosition {
  node: UnionFindNodeState;
  x: number;
  y: number;
  parentX?: number;
  parentY?: number;
}

interface TreeInfo {
  rootId: string;
  nodes: Set<string>;
  depth: number;
  width: number;
}

function buildNodeMap(nodes: UnionFindNodeState[]): Map<string, UnionFindNodeState> {
  const map = new Map<string, UnionFindNodeState>();
  for (const node of nodes) {
    map.set(node.id, node);
  }
  return map;
}

/**
 * Build children map: for each node, which nodes have it as parent?
 */
function buildChildrenMap(nodes: UnionFindNodeState[]): Map<string, string[]> {
  const childrenMap = new Map<string, string[]>();
  for (const node of nodes) {
    childrenMap.set(node.id, []);
  }
  for (const node of nodes) {
    if (node.parentId !== null) {
      const children = childrenMap.get(node.parentId);
      if (children) {
        children.push(node.id);
      }
    }
  }
  return childrenMap;
}

/**
 * Identify roots (nodes where parentId is null).
 */
function findRoots(nodes: UnionFindNodeState[]): string[] {
  return nodes.filter(n => n.parentId === null).map(n => n.id);
}

/**
 * Calculate the depth of a tree rooted at the given node.
 */
function calculateTreeDepth(
  nodeId: string,
  childrenMap: Map<string, string[]>
): number {
  const children = childrenMap.get(nodeId) ?? [];
  if (children.length === 0) {
    return 1;
  }
  let maxChildDepth = 0;
  for (const childId of children) {
    const childDepth = calculateTreeDepth(childId, childrenMap);
    maxChildDepth = Math.max(maxChildDepth, childDepth);
  }
  return 1 + maxChildDepth;
}

/**
 * Calculate the width needed for a subtree rooted at the given node.
 */
function calculateSubtreeWidth(
  nodeId: string,
  childrenMap: Map<string, string[]>
): number {
  const children = childrenMap.get(nodeId) ?? [];
  if (children.length === 0) {
    return MIN_NODE_SPACING;
  }
  let totalWidth = 0;
  for (const childId of children) {
    totalWidth += calculateSubtreeWidth(childId, childrenMap);
  }
  return Math.max(totalWidth, MIN_NODE_SPACING);
}

/**
 * Collect all node IDs in a tree rooted at rootId.
 */
function collectTreeNodes(
  rootId: string,
  childrenMap: Map<string, string[]>,
  collected: Set<string>
): void {
  collected.add(rootId);
  const children = childrenMap.get(rootId) ?? [];
  for (const childId of children) {
    collectTreeNodes(childId, childrenMap, collected);
  }
}

/**
 * Build tree info for each root in the forest.
 */
function buildTreeInfos(
  roots: string[],
  childrenMap: Map<string, string[]>
): TreeInfo[] {
  return roots.map(rootId => {
    const nodes = new Set<string>();
    collectTreeNodes(rootId, childrenMap, nodes);
    const depth = calculateTreeDepth(rootId, childrenMap);
    const width = calculateSubtreeWidth(rootId, childrenMap);
    return { rootId, nodes, depth, width };
  });
}

/**
 * Calculate positions for all nodes in a forest layout.
 * Roots are at the TOP, children are BELOW.
 * Parent pointers go UPWARD from child to parent.
 */
function calculatePositions(
  unionFind: UnionFindState
): { positions: NodePosition[]; svgWidth: number; svgHeight: number } {
  const nodeMap = buildNodeMap(unionFind.nodes);
  const childrenMap = buildChildrenMap(unionFind.nodes);
  const roots = findRoots(unionFind.nodes);
  const treeInfos = buildTreeInfos(roots, childrenMap);

  // Calculate total width needed
  const totalWidth = treeInfos.reduce((sum, t) => sum + t.width, 0) +
    (treeInfos.length - 1) * TREE_GAP +
    SVG_PADDING * 2;

  // Calculate max depth for height
  const maxDepth = Math.max(1, ...treeInfos.map(t => t.depth));
  const svgHeight = maxDepth * LEVEL_HEIGHT + SVG_PADDING * 2;

  // Pre-calculate subtree widths
  const subtreeWidths = new Map<string, number>();
  for (const node of unionFind.nodes) {
    subtreeWidths.set(node.id, calculateSubtreeWidth(node.id, childrenMap));
  }

  const positions: NodePosition[] = [];
  let currentX = SVG_PADDING;

  for (const treeInfo of treeInfos) {
    const treeWidth = treeInfo.width;

    // Traverse tree and position nodes
    function traverse(
      nodeId: string,
      level: number,
      left: number,
      right: number,
      parentX?: number,
      parentY?: number
    ) {
      const node = nodeMap.get(nodeId);
      if (!node) return;

      const x = (left + right) / 2;
      const y = SVG_PADDING + level * LEVEL_HEIGHT + NODE_RADIUS;

      positions.push({ node, x, y, parentX, parentY });

      const children = childrenMap.get(nodeId) ?? [];
      if (children.length > 0) {
        const childrenWidth = children.reduce(
          (sum, childId) => sum + (subtreeWidths.get(childId) ?? MIN_NODE_SPACING),
          0
        );

        let childLeft = x - childrenWidth / 2;

        for (const childId of children) {
          const childWidth = subtreeWidths.get(childId) ?? MIN_NODE_SPACING;
          traverse(
            childId,
            level + 1,
            childLeft,
            childLeft + childWidth,
            x,
            y
          );
          childLeft += childWidth;
        }
      }
    }

    traverse(treeInfo.rootId, 0, currentX, currentX + treeWidth);
    currentX += treeWidth + TREE_GAP;
  }

  return { positions, svgWidth: Math.max(totalWidth, 200), svgHeight };
}

export function UnionFindView({ unionFind, className }: UnionFindViewProps) {
  const { positions, svgWidth, svgHeight } = useMemo(
    () => calculatePositions(unionFind),
    [unionFind]
  );

  // Build a set of highlighted path nodes for edge styling
  const pathNodeIds = new Set(unionFind.findPath ?? []);

  return (
    <div className={cn('flex flex-col items-center', className)}>
      {unionFind.label && (
        <span className="mb-2 text-sm font-medium text-[var(--muted-foreground)]">
          {unionFind.label}
        </span>
      )}
      <svg
        width={svgWidth}
        height={svgHeight}
        viewBox={`0 0 ${svgWidth} ${svgHeight}`}
        className="overflow-visible"
      >
        {/* Draw edges (from child UP to parent) */}
        {positions.map(({ node, x, y, parentX, parentY }) => {
          if (parentX === undefined || parentY === undefined) return null;

          // Determine if this edge is on the current find path
          const isOnPath = pathNodeIds.has(node.id);

          // Draw line from child to parent (upward)
          // Arrow points toward parent (at the top)
          const startY = y - NODE_RADIUS;
          const endY = parentY + NODE_RADIUS;

          return (
            <g key={`edge-${node.id}`}>
              <motion.line
                x1={x}
                y1={startY}
                x2={parentX}
                y2={endY}
                className={cn(
                  'stroke-[var(--border)]',
                  isOnPath && 'stroke-[var(--primary)]',
                  node.state === 'cycle' && 'stroke-[var(--error)]'
                )}
                strokeWidth={isOnPath ? 2.5 : 2}
                markerEnd="url(#arrowhead)"
                initial={false}
                animate={{
                  opacity: 1,
                }}
                transition={{ duration: 0.2 }}
              />
            </g>
          );
        })}

        {/* Arrow marker definition */}
        <defs>
          <marker
            id="arrowhead"
            markerWidth="10"
            markerHeight="7"
            refX="9"
            refY="3.5"
            orient="auto"
          >
            <polygon
              points="0 0, 10 3.5, 0 7"
              className="fill-[var(--border)]"
            />
          </marker>
          <marker
            id="arrowhead-primary"
            markerWidth="10"
            markerHeight="7"
            refX="9"
            refY="3.5"
            orient="auto"
          >
            <polygon
              points="0 0, 10 3.5, 0 7"
              className="fill-[var(--primary)]"
            />
          </marker>
        </defs>

        {/* Draw nodes */}
        {positions.map(({ node, x, y }) => (
          <UnionFindNode
            key={node.id}
            node={node}
            x={x}
            y={y}
            radius={NODE_RADIUS}
          />
        ))}

        {/* Current operation indicator */}
        {unionFind.message && (
          <text
            x={svgWidth / 2}
            y={svgHeight - 10}
            textAnchor="middle"
            className="fill-[var(--muted-foreground)] text-xs"
          >
            {unionFind.message}
          </text>
        )}

        {/* Current edge indicator */}
        {unionFind.currentEdge && (
          <text
            x={svgWidth / 2}
            y={20}
            textAnchor="middle"
            className="fill-[var(--foreground)] text-sm font-mono"
          >
            Processing edge: [{unionFind.currentEdge[0]}, {unionFind.currentEdge[1]}]
          </text>
        )}
      </svg>
    </div>
  );
}
