import { Box } from "@mui/material";
import { ReactFlow, Background, Controls, MiniMap, SmoothStepEdge } from "@xyflow/react";
import "@xyflow/react/dist/style.css";
import { useMemo, useRef, useState, useCallback } from "react";
import type { IGraph, IAgent } from '@/interfaces/database/agent';
import NodeDrawer from './NodeDrawer';
import ReadonlyNode from './node/ReadonlyNode';
import NoteNode from './node/NoteNode';
import { useAgentCanvas } from '@/hooks/agent-hooks';

const nodeTypes = {
  beginNode: ReadonlyNode,
  agentNode: ReadonlyNode,
  messageNode: ReadonlyNode,
  toolNode: ReadonlyNode,
  noteNode: NoteNode,
  retrievalNode: ReadonlyNode,
  ragNode: ReadonlyNode,
  switchNode: ReadonlyNode,
  categorizeNode: ReadonlyNode,
  iterationNode: ReadonlyNode,
  iterationStartNode: ReadonlyNode,
  keywordNode: ReadonlyNode,
  relevantNode: ReadonlyNode,
  logicNode: ReadonlyNode,
  emailNode: ReadonlyNode,
  tokenizerNode: ReadonlyNode,
  templateNode: ReadonlyNode,
  rewriteNode: ReadonlyNode,
  invokeNode: ReadonlyNode,
  parserNode: ReadonlyNode,
  fileNode: ReadonlyNode,
  extractorNode: ReadonlyNode,
  splitterNode: ReadonlyNode,
  placeholderNode: ReadonlyNode,
  contextNode: ReadonlyNode,
  group: ReadonlyNode,
  generateNode: ReadonlyNode,
} as const;

const edgeTypes = {
  buttonEdge: SmoothStepEdge,
  ButtonEdge: SmoothStepEdge,
} as const;

export interface AgentCanvasProps {
  agent?: IAgent | null;
  graph?: IGraph | null;
  onReady?: (api: { fitView: () => void }) => void;
  onAutoSaved?: (date: Date) => void;
}

function AgentCanvas({ agent, graph, onReady, onAutoSaved }: AgentCanvasProps) {
  const { sanitized, selectedNode, hoverNode, onInit,
    onNodeClick, onPaneClick, onNodeMouseEnter } = useAgentCanvas({ agent, graph, onReady, onAutoSaved });

  return (
    <Box sx={{ flex: 1, position: 'relative', height: '100%' }}>
      {sanitized && (
        <ReactFlow
          nodes={sanitized.nodes || []}
          edges={sanitized.edges || []}
          fitView
          nodesDraggable={false}
          nodesConnectable={false}
          elementsSelectable={false}
          defaultEdgeOptions={{ animated: true, type: 'smoothstep' }}
          nodeTypes={nodeTypes as any}
          edgeTypes={edgeTypes as any}
          onNodeClick={onNodeClick}
          onPaneClick={onPaneClick}
          onNodeMouseEnter={onNodeMouseEnter}
          onInit={onInit}
          proOptions={{ hideAttribution: true }}
          minZoom={0.25}
          maxZoom={2}
          style={{ width: '100%', height: '100%' }}
        >
          <Background />
          <MiniMap pannable zoomable />
          <Controls showInteractive={false} />
        </ReactFlow>
      )}
      <NodeDrawer open={!!selectedNode} node={hoverNode || selectedNode} onClose={onPaneClick} />
    </Box>
  );
}

export default AgentCanvas;