import { useMemo } from 'react';
import { Table, Tag, Checkbox, Typography, Tooltip } from 'antd';
import type { AttributeDAG, CrossoverPair, DAGNode } from '../../types';

const { Text } = Typography;

interface CrossoverMatrixProps {
  dagA: AttributeDAG;
  dagB: AttributeDAG;
  pairs: CrossoverPair[];
  onTogglePair: (pairId: string) => void;
  isDark: boolean;
}

interface MatrixCell {
  nodeA: DAGNode;
  nodeB: DAGNode;
  pair: CrossoverPair | null;
}

export function CrossoverMatrix({
  dagA,
  dagB,
  pairs,
  onTogglePair,
  isDark,
}: CrossoverMatrixProps) {
  // Group nodes by category
  const groupByCategory = (nodes: DAGNode[]) => {
    return nodes.reduce((acc, node) => {
      if (!acc[node.category]) acc[node.category] = [];
      acc[node.category].push(node);
      return acc;
    }, {} as Record<string, DAGNode[]>);
  };

  const nodesByCategoryA = useMemo(() => groupByCategory(dagA.nodes), [dagA.nodes]);
  const nodesByCategoryB = useMemo(() => groupByCategory(dagB.nodes), [dagB.nodes]);

  // Create a lookup map for pairs
  const pairLookup = useMemo(() => {
    const lookup = new Map<string, CrossoverPair>();
    for (const pair of pairs) {
      // Key format: sourceNodeId-targetNodeId
      lookup.set(`${pair.sourceNode.id}-${pair.targetNode.id}`, pair);
    }
    return lookup;
  }, [pairs]);

  // Get all categories from both DAGs
  const categoriesA = Object.keys(nodesByCategoryA);
  const categoriesB = Object.keys(nodesByCategoryB);

  // Render a matrix for each category pair
  const renderCategoryMatrix = (categoryA: string, categoryB: string) => {
    const nodesA = nodesByCategoryA[categoryA] || [];
    const nodesB = nodesByCategoryB[categoryB] || [];

    if (nodesA.length === 0 || nodesB.length === 0) return null;

    // Create matrix data
    const matrixData = nodesA.map((nodeA) => {
      const row: Record<string, MatrixCell | DAGNode> = { nodeA };
      for (const nodeB of nodesB) {
        const pair = pairLookup.get(`${nodeA.id}-${nodeB.id}`) ||
                     pairLookup.get(`${nodeB.id}-${nodeA.id}`) ||
                     null;
        row[nodeB.id] = { nodeA, nodeB, pair };
      }
      return row;
    });

    // Create columns
    const columns = [
      {
        title: (
          <div style={{ display: 'flex', alignItems: 'center', gap: 4 }}>
            <Tag color="blue">A</Tag>
            <Text strong>{categoryA}</Text>
          </div>
        ),
        dataIndex: 'nodeA',
        key: 'nodeA',
        width: 120,
        fixed: 'left' as const,
        render: (node: DAGNode) => (
          <Text ellipsis style={{ maxWidth: 100 }}>{node.name}</Text>
        ),
      },
      ...nodesB.map((nodeB) => ({
        title: (
          <Tooltip title={nodeB.name}>
            <Text ellipsis style={{ maxWidth: 80 }}>{nodeB.name}</Text>
          </Tooltip>
        ),
        dataIndex: nodeB.id,
        key: nodeB.id,
        width: 60,
        align: 'center' as const,
        render: (cell: MatrixCell) => {
          if (!cell || !cell.pair) {
            return <div style={{ color: isDark ? '#555' : '#ddd' }}>-</div>;
          }
          return (
            <Checkbox
              checked={cell.pair.selected}
              onChange={() => onTogglePair(cell.pair!.id)}
            />
          );
        },
      })),
    ];

    return (
      <div key={`${categoryA}-${categoryB}`} style={{ marginBottom: 16 }}>
        <div style={{
          display: 'flex',
          alignItems: 'center',
          gap: 8,
          marginBottom: 8,
          padding: '4px 8px',
          background: isDark ? '#1f1f1f' : '#fafafa',
          borderRadius: 4,
        }}>
          <Tag color="blue">A: {categoryA}</Tag>
          <Text type="secondary">×</Text>
          <Tag color="green">B: {categoryB}</Tag>
        </div>
        <Table
          dataSource={matrixData}
          columns={columns}
          pagination={false}
          size="small"
          scroll={{ x: 'max-content' }}
          rowKey={(row) => (row.nodeA as DAGNode).id}
          bordered
        />
      </div>
    );
  };

  // Render matrices for matching categories (same category in both paths)
  const matchingCategories = categoriesA.filter(cat => categoriesB.includes(cat));

  return (
    <div style={{ padding: 16 }}>
      {matchingCategories.length > 0 ? (
        <>
          <Text type="secondary" style={{ display: 'block', marginBottom: 16 }}>
            Select attribute pairs to cross between Path A and Path B
          </Text>
          {matchingCategories.map(category =>
            renderCategoryMatrix(category, category)
          )}
        </>
      ) : (
        <Text type="secondary">
          No matching categories found between the two paths
        </Text>
      )}
    </div>
  );
}
