import { type IArtifactGraphEntity } from '@/interfaces/database/dataset';
import { cn } from '@/lib/utils';
import { memo, useCallback, useEffect, useMemo, useRef } from 'react';
import ForceGraph2D, { type ForceGraphMethods } from 'react-force-graph-2d';
import {
  getNodeColor as defaultGetNodeColor,
  getNodeRadius as defaultGetNodeRadius,
  MinNodeRadius,
} from './node-style';
import {
  type ArtifactForceGraphProps,
  type ArtifactGraphLink,
  type ArtifactGraphNode,
} from './types';
import { useArtifactGraphData } from './use-artifact-graph-data';
import { useCenterGravity } from './use-center-gravity';
import { useContainerDimensions } from './use-container-dimensions';
import { useGraphHighlight } from './use-graph-highlight';
import { defaultMapNodeToValue } from './utils';

const defaultGetNodeId = (node: IArtifactGraphEntity) => node.slug;

const nodeCanvasObjectMode = () => 'after' as const;

function ArtifactForceGraph<TNodeValue = IArtifactGraphEntity>({
  data,
  show = true,
  onNodeClick,
  mapNodeToValue = defaultMapNodeToValue as (
    node: IArtifactGraphEntity,
  ) => TNodeValue,
  getNodeId = defaultGetNodeId,
  getNodeColor = defaultGetNodeColor,
  getNodeRadius = defaultGetNodeRadius,
  highlightNodeId,
}: ArtifactForceGraphProps<TNodeValue>) {
  const containerRef = useRef<HTMLDivElement>(null);
  const fgRef = useRef<ForceGraphMethods<ArtifactGraphNode> | undefined>(
    undefined,
  );
  const hasFittedRef = useRef(false);
  const dimensions = useContainerDimensions(containerRef, show);
  const hasDimensions = dimensions.width > 0 && dimensions.height > 0;

  const graphData = useArtifactGraphData({
    data,
    getNodeId,
    getNodeColor,
    getNodeRadius,
  });

  // Resolve the controlled id back to a node object reference (highlighting relies on the node's __neighbors/__links)
  const pinnedNode = useMemo(
    () =>
      highlightNodeId
        ? ((graphData.nodes.find((node) => node.id === highlightNodeId) ??
            null) as ArtifactGraphNode | null)
        : null,
    [graphData, highlightNodeId],
  );

  const {
    handleNodeHover,
    getNodeColor: nodeColor,
    getLinkColor,
    getLinkWidth,
    paintNode,
  } = useGraphHighlight(containerRef, pinnedNode);

  useEffect(() => {
    hasFittedRef.current = false;
  }, [graphData]);

  useCenterGravity(fgRef, hasDimensions);

  const handleEngineStop = useCallback(() => {
    if (!hasFittedRef.current && fgRef.current) {
      fgRef.current.zoomToFit(400);
      hasFittedRef.current = true;
    }
  }, []);

  const handleNodeClick = useCallback(
    (node: IArtifactGraphEntity) => {
      onNodeClick?.(mapNodeToValue(node));
    },
    [onNodeClick, mapNodeToValue],
  );

  const nodeVal = useCallback(
    (node: ArtifactGraphNode) => node.__radius ?? MinNodeRadius,
    [],
  );

  // Hover tooltip shows the entity description; empty string hides it
  const getNodeLabel = useCallback(
    (node: ArtifactGraphNode) => node.description ?? '',
    [],
  );

  // Empty label hides the tooltip, so relations without a type show nothing
  const getLinkLabel = useCallback(
    (link: ArtifactGraphLink) => link.type ?? '',
    [],
  );

  return (
    <div
      ref={containerRef}
      className={cn('flex-1 min-h-0 h-full', !show && 'hidden')}
    >
      {hasDimensions && (
        <ForceGraph2D
          ref={fgRef}
          width={dimensions.width}
          height={dimensions.height}
          graphData={graphData}
          nodeRelSize={1}
          nodeColor={nodeColor}
          nodeVal={nodeVal}
          cooldownTicks={100}
          nodeLabel={getNodeLabel}
          autoPauseRedraw={false}
          onEngineStop={handleEngineStop}
          onNodeClick={handleNodeClick}
          onNodeHover={handleNodeHover}
          nodeCanvasObject={paintNode}
          nodeCanvasObjectMode={nodeCanvasObjectMode}
          linkColor={getLinkColor}
          linkWidth={getLinkWidth}
          linkLabel={getLinkLabel}
        />
      )}
    </div>
  );
}

const MemoizedArtifactForceGraph = memo(ArtifactForceGraph) as <
  TNodeValue = IArtifactGraphEntity,
>(
  props: ArtifactForceGraphProps<TNodeValue>,
) => React.ReactElement;

export { MemoizedArtifactForceGraph as ArtifactForceGraph };
export default MemoizedArtifactForceGraph;
