import React, {useCallback} from 'react';
import ReactFlow, {addEdge, ConnectionLineType, useNodesState, useEdgesState} from 'react-flow-renderer';
import dagre from 'dagre';
import DataConstructor from './DataConstructor'
import {Box} from "@mui/material";


const dagreGraph = new dagre.graphlib.Graph();
dagreGraph.setDefaultEdgeLabel(() => ({}));

const nodeWidth = 172;
const nodeHeight = 36;

const getLayoutedElements = (nodes, edges, direction = 'TB') => {
    const isHorizontal = direction === 'LR';
    dagreGraph.setGraph({rankdir: direction});

    nodes.forEach((node) => {
        dagreGraph.setNode(node.id, {width: nodeWidth, height: nodeHeight});
    });

    edges.forEach((edge) => {
        dagreGraph.setEdge(edge.source, edge.target);
    });

    dagre.layout(dagreGraph);

    nodes.forEach((node) => {
        const nodeWithPosition = dagreGraph.node(node.id);
        node.targetPosition = isHorizontal ? 'left' : 'top';
        node.sourcePosition = isHorizontal ? 'right' : 'bottom';

        // We are shifting the dagre node position (anchor=center center) to the top left
        // so it matches the React Flow node anchor point (top left).
        node.position = {
            x: nodeWithPosition.x - nodeWidth / 2,
            y: nodeWithPosition.y - nodeHeight / 2,
        };

        return node;
    });

    return {nodes, edges};
};

const LayoutFlow = (props) => {
    const {_nodes, _edges} = DataConstructor(props)
    const {nodes: layoutedNodes, edges: layoutedEdges} = getLayoutedElements(
        _nodes,
        _edges
    );

    const [nodes, setNodes, onNodesChange] = useNodesState(layoutedNodes);
    const [edges, setEdges, onEdgesChange] = useEdgesState(layoutedEdges);
    const onConnect = useCallback(
        (params) => setEdges((eds) => addEdge({...params, type: ConnectionLineType.SmoothStep, animated: true}, eds)),
        []
    );
    const onLayout = useCallback(
        (direction) => {
            const {nodes: layoutedNodes, edges: layoutedEdges} = getLayoutedElements(
                nodes,
                edges,
                direction
            );

            setNodes([...layoutedNodes]);
            setEdges([...layoutedEdges]);
        },
        [nodes, edges]
    );


    return (
        <Box sx={{border: 1}} className="layoutflow" style={{width: 1200, height: 560}}>
            <ReactFlow
                nodes={nodes}
                edges={edges}
                onNodesChange={onNodesChange}
                onEdgesChange={onEdgesChange}
                onConnect={onConnect}
                connectionLineType={ConnectionLineType.SmoothStep}
                fitView
            />
            <div className="controls">
                <button onClick={() => onLayout('TB')}>vertical layout</button>
                <button onClick={() => onLayout('LR')}>horizontal layout</button>
            </div>
        </Box>
    );
};

export default LayoutFlow;