import { forwardRef, useImperativeHandle } from 'react';
import {
  ReactFlow,
  useReactFlow,
  ReactFlowProvider,
  Background,
} from '@xyflow/react';
import '@xyflow/react/dist/style.css';

import type {
  TransformationDAGResult,
  ExpertTransformationDAGResult,
  CategoryDefinition
} from '../types';
import {
  transformationNodeTypes,
  useTransformationLayout,
  useExpertTransformationLayout
} from './transformation';
import '../styles/mindmap.css';

interface TransformationDAGProps {
  data: TransformationDAGResult | ExpertTransformationDAGResult;
  categories: CategoryDefinition[];
  isDark: boolean;
  crossoverSource?: string; // If set, marks all attributes as crossed over from this source
}

export interface TransformationDAGRef {
  resetView: () => void;
}

const TransformationDAGInner = forwardRef<TransformationDAGRef, TransformationDAGProps>(
  ({ data, categories, isDark, crossoverSource }, ref) => {
    const { setViewport } = useReactFlow();

    // Check if data is ExpertTransformationDAGResult by checking for 'experts' property
    const isExpertTransformation = 'experts' in data;

    // Use appropriate layout hook based on data type
    const regularLayout = useTransformationLayout(
      !isExpertTransformation ? (data as TransformationDAGResult) : null,
      categories,
      { isDark, fontSize: 13 }
    );

    const expertLayout = useExpertTransformationLayout(
      isExpertTransformation ? (data as ExpertTransformationDAGResult) : null,
      categories,
      { isDark, fontSize: 13, crossoverSource }
    );

    const { nodes, edges } = isExpertTransformation ? expertLayout : regularLayout;

    useImperativeHandle(
      ref,
      () => ({
        resetView: () => {
          setViewport({ x: 50, y: 50, zoom: 1 }, { duration: 300 });
        },
      }),
      [setViewport]
    );

    return (
      <ReactFlow
        nodes={nodes}
        edges={edges}
        nodeTypes={transformationNodeTypes}
        fitView
        fitViewOptions={{ padding: 0.3 }}
        minZoom={0.3}
        maxZoom={2}
        defaultViewport={{ x: 50, y: 50, zoom: 1 }}
        proOptions={{ hideAttribution: true }}
        nodesDraggable={false}
        nodesConnectable={false}
        elementsSelectable={false}
        panOnDrag
        zoomOnScroll
        zoomOnPinch
      >
        <Background
          color={isDark ? '#333' : '#e0e0e0'}
          gap={20}
          size={1}
        />
      </ReactFlow>
    );
  }
);

TransformationDAGInner.displayName = 'TransformationDAGInner';

export const TransformationDAG = forwardRef<TransformationDAGRef, TransformationDAGProps>(
  (props, ref) => (
    <div
      className={`mindmap-container ${props.isDark ? 'mindmap-dark' : 'mindmap-light'}`}
    >
      <ReactFlowProvider>
        <TransformationDAGInner {...props} ref={ref} />
      </ReactFlowProvider>
    </div>
  )
);

TransformationDAG.displayName = 'TransformationDAG';
