package application.editor.stages;

import algorithms.DataTransferModelAnalyzer;
import application.editor.Editor;
import application.editor.FlowCellEditor;
import application.editor.Stage;
import com.mxgraph.model.mxCell;
import com.mxgraph.model.mxGeometry;
import com.mxgraph.model.mxGraphModel;
import com.mxgraph.swing.mxGraphComponent;
import com.mxgraph.util.mxEventSource.mxIEventListener;
import com.mxgraph.util.mxPoint;
import generators.JavaCodeGenerator;
import models.Edge;
import models.algebra.*;
import models.dataConstraintModel.Channel;
import models.dataConstraintModel.ChannelMember;
import models.dataConstraintModel.ResourcePath;
import models.dataConstraintModel.Selector;
import models.dataFlowModel.*;

import java.awt.event.MouseListener;
import java.util.*;

public class PushPullSelectionStage extends Stage {

    private DataFlowGraph dataFlowGraph = null;

    public PushPullSelectionStage(mxGraphComponent graphComponent) {
        super(graphComponent);
    }

    @Override
    public void init(Stage prevStage) {
        if (prevStage instanceof DataFlowModelingStage) {
            model = prevStage.getModel();
            dataFlowGraph = analyzeDataTransferModel(model);
        }
    }

    @Override
    public FlowCellEditor createCellEditor(mxGraphComponent graphComponent) {
        return new PushPullSelectionCellEditor(this, graphComponent);
    }

    @Override
    public mxIEventListener createChangeEventListener(Editor editor) {
        return (sender, eventObject) -> {
            List<mxCell> terminals = new ArrayList<>();
            mxCell cell = null;
            for (Object change : ((List) eventObject.getProperties().get("changes"))) {
                if (change instanceof mxGraphModel.mxTerminalChange) {
                    mxGraphModel.mxTerminalChange terminalChange = (mxGraphModel.mxTerminalChange) change;
                    cell = (mxCell) terminalChange.getCell();
                    mxCell terminal = (mxCell) terminalChange.getTerminal();
                    terminals.add(terminal);
                }
            }
            if (terminals.size() == 2) {
                // cancel connect
                graph.removeCells(new mxCell[]{cell});
                graph.clearSelection();
            }
        };
    }

    @Override
    public MouseListener createMouseEventListener(Editor editor) {
        return null;
    }

    @Override
    public boolean canChangeFrom(Stage prevStage) {
        if (prevStage instanceof DataFlowModelingStage) {
            return ((DataFlowModelingStage) prevStage).isValid();
        }
        return false;
    }

    private DataFlowGraph analyzeDataTransferModel(DataTransferModel model) {
        DataFlowGraph flowGraph = DataTransferModelAnalyzer.createDataFlowGraphWithStateStoringAttribute(model);
        DataFlowGraph dataFlowGraph = DataTransferModelAnalyzer.annotateWithSelectableDataTransferAttiribute(flowGraph);
        updateEdgeAttributes(dataFlowGraph);
        return dataFlowGraph;
    }

    private void updateEdgeAttributes(DataFlowGraph dataFlowGraph) {
        mxCell root = (mxCell) graph.getDefaultParent();

        graph.getModel().beginUpdate();
        try {
            // add input, output and reference edges
            for (Edge e : dataFlowGraph.getEdges()) {
                if (e instanceof DataFlowEdge) {
                    DataFlowEdge dataFlow = (DataFlowEdge) e;
                    if (!dataFlow.isChannelToResource()) {
                        ResourceNode srcRes = (ResourceNode) dataFlow.getSource();
                        DataTransferChannel channel = ((ChannelNode) dataFlow.getDestination()).getChannel();
                        // input edge
                        for (Object edge : graph.getChildEdges(root)) {
                            mxCell edgeCell = (mxCell) edge;
                            if (edgeCell.getValue() instanceof Editor.SrcDstAttribute) {
                                Editor.SrcDstAttribute edgeAttr = (Editor.SrcDstAttribute) edgeCell.getValue();
                                if (srcRes.getPrimaryResourcePath().equals(edgeAttr.getSource()) && channel.equals(edgeAttr.getDestination())) {
                                    edgeCell.setValue(dataFlow.getAttribute());
                                    break;
                                }
                            }
                        }
                    }
                }
            }
        } finally {
            graph.getModel().endUpdate();
        }
        graph.refresh();
    }

    /**
     * Construct a mxGraph from DataFlowGraph
     */
    public void constructGraph() {
        ((mxGraphModel) graph.getModel()).clear();
        mxCell parent = (mxCell) graph.getDefaultParent();

        graph.getModel().beginUpdate();
        try {
            Map<DataTransferChannel, Object> channelsIn = new HashMap<>();
            Map<DataTransferChannel, Object> channelsOut = new HashMap<>();
            Map<ResourceNode, Object> resources = new HashMap<>();

            mxGeometry geoPortIn = new mxGeometry(0, 0.5, PORT_DIAMETER, PORT_DIAMETER);
            geoPortIn.setOffset(new mxPoint(-PORT_RADIUS, -PORT_RADIUS));
            geoPortIn.setRelative(true);

            mxGeometry geoPortOut = new mxGeometry(1.0, 0.5, PORT_DIAMETER, PORT_DIAMETER);
            geoPortOut.setOffset(new mxPoint(-PORT_RADIUS, -PORT_RADIUS));
            geoPortOut.setRelative(true);

            // create resource vertices
            for (ResourceNode resourceNode : dataFlowGraph.getRootResourceNodes()) {
                int w = 80;
                int h = 30;
                ResourcePath resourcePath = resourceNode.getPrimaryResourcePath();
                Object resource = graph.insertVertex(parent, null, resourcePath.getLeafResourceName(), 20, 20, w, h, "shape=ellipse;perimeter=ellipsePerimeter"); // insert a resource as a vertex
                resources.put(resourceNode, resource);
                createChildResourceVertices(resource, resourceNode, resources, w, h);
            }

            // create channel vertices
            for (ChannelNode channelNode : dataFlowGraph.getRootChannelNodes()) {
                DataTransferChannel channel = channelNode.getChannel();
                if (!channel.getInputResources().isEmpty()) {
                    // Normal channel
                    if (channelsIn.get(channel) == null || channelsOut.get(channel) == null) {
                        if (Objects.equals(channel.getSelectors().toString(), "[]")) {
                            Object channelCell = graph.insertVertex(parent, null, channel.getChannelName(), 150, 20, 30, 30); // insert a channel as a vertex

                            mxCell portIn = new mxCell(null, geoPortIn, "shape=ellipse;perimter=ellipsePerimeter");
                            portIn.setVertex(true);
                            graph.addCell(portIn, channelCell);        // insert the input port of a channel
                            channelsIn.put(channel, portIn);

                            mxCell port_out = new mxCell(null, geoPortOut, "shape=ellipse;perimter=ellipsePerimeter");
                            port_out.setVertex(true);
                            graph.addCell(port_out, channelCell); // insert the output port of a channel
                            channelsOut.put(channel, port_out);
                        } else {
                            for (Selector s : channel.getSelectors()) {
                                Expression exp = s.getExpression();
                                String selectorName = exp.toString();
                                String channelName = channel.getChannelName();
                                String cellName = channelName + " (" + selectorName + ")";

                                Object channelCell = graph.insertVertex(parent, null, cellName, 150, 20, 60, 30); // insert a channel as a vertex

                                mxCell portIn = new mxCell(null, geoPortIn, "shape=ellipse;perimter=ellipsePerimeter");
                                portIn.setVertex(true);
                                graph.addCell(portIn, channelCell); // insert the input port of a channel
                                channelsIn.put(channel, portIn);

                                mxCell portOut = new mxCell(null, geoPortOut, "shape=ellipse;perimter=ellipsePerimeter");
                                portOut.setVertex(true);
                                graph.addCell(portOut, channelCell); // insert the output port of a channel
                                channelsOut.put(channel, portOut);
                            }
                        }
                    }
                } else {
                    // I/O channel
                    if (channelsOut.get(channel) == null) {
                        if (Objects.equals(channel.getSelectors().toString(), "[]")) {
                            Object chCell = graph.insertVertex(parent, null, channel.getChannelName(), 150, 20, 30, 30); // insert a channel as a vertex
                            mxCell port_in = new mxCell(null, geoPortIn, "shape=ellipse;perimter=ellipsePerimeter");
                            port_in.setVertex(true);
                            graph.addCell(port_in, chCell);        // insert the input port of a channel
                            mxCell port_out = new mxCell(null, geoPortOut, "shape=ellipse;perimter=ellipsePerimeter");
                            port_out.setVertex(true);
                            graph.addCell(port_out, chCell);        // insert the output port of a channel
                            channelsIn.put(channel, port_in);
                            channelsOut.put(channel, port_out);
                        } else {
                            for (Selector s : channel.getSelectors()) {
                                Expression exp = s.getExpression();
                                String selectorName = exp.toString();
                                String channelName = channel.getChannelName();
                                String cellName = channelName + " (" + selectorName + ")";
                                Object chCell = graph.insertVertex(parent, null, cellName, 150, 20, 60, 30); // insert a channel as a vertex

                                mxCell port_in = new mxCell(null, geoPortIn, "shape=ellipse;perimter=ellipsePerimeter");
                                port_in.setVertex(true);
                                graph.addCell(port_in, chCell); // insert the input port of a channel
                                channelsIn.put(channel, port_in);

                                mxCell port_out = new mxCell(null, geoPortOut, "shape=ellipse;perimter=ellipsePerimeter");
                                port_out.setVertex(true);
                                graph.addCell(port_out, chCell); // insert the output port of a channel
                                channelsOut.put(channel, port_out);
                            }
                        }
                    }
                }
            }

            // add input, output and reference edges
            for (Edge edge : dataFlowGraph.getEdges()) {
                DataFlowEdge dataFlowEdge = (DataFlowEdge) edge;
                if (dataFlowEdge.isChannelToResource()) {
                    // output edge
                    DataTransferChannel channel = ((ChannelNode) dataFlowEdge.getSource()).getChannel();
                    ResourcePath dstRes = ((ResourceNode) dataFlowEdge.getDestination()).getInSideResource(channel);
                    graph.insertEdge(parent, null, new Editor.SrcDstAttribute(channel, dstRes), channelsOut.get(channel), resources.get((ResourceNode) dataFlowEdge.getDestination()), "movable=false");
                } else {
                    // input edge
                    DataTransferChannel channel = ((ChannelNode) dataFlowEdge.getDestination()).getChannel();
                    ResourcePath srcRes = ((ResourceNode) dataFlowEdge.getSource()).getOutSideResource(channel);
                    Set<Map.Entry<ResourceNode, ResourceNode>> toRes = getResourceDependencyForChannel(channel, dataFlowGraph);
                    for (Map.Entry<ResourceNode, ResourceNode> RtoR : toRes) {
                        graph.insertEdge(parent, null, null, resources.get(RtoR.getValue()), resources.get(RtoR.getKey()), "dashed=true;movable=false");
                    }

                    graph.insertEdge(parent, null, new Editor.SrcDstAttribute(srcRes, channel), resources.get((ResourceNode) dataFlowEdge.getSource()), channelsIn.get(channel), "movable=false");
                }
            }

            for (Channel ch : model.getChannels()) {
                // reference edges
                DataTransferChannel channel = (DataTransferChannel) ch;
                for (ResourcePath refRes : channel.getReferenceResources()) {
                    graph.insertEdge(parent, null, null, resources.get(dataFlowGraph.getResourceNode(refRes)), channelsIn.get(channel), "dashed=true;movable=false");
                }
            }
        } finally {
            graph.getModel().endUpdate();
        }
    }

    private void createChildResourceVertices(Object resource, ResourceNode resNode, Map<ResourceNode, Object> resources, int w, int h) {
        for (ResourceNode childNode : resNode.getChildren()) {
            ResourcePath childRes = childNode.getPrimaryResourcePath();
            Object childResource = graph.insertVertex(resource, null, childRes.getName(), 0, 0, w, h, "shape=ellipse;perimeter=ellipsePerimeter"); // insert a resource as a vertex
            resources.put(childNode, childResource);
            createChildResourceVertices(childResource, childNode, resources, w, h);
        }
    }

    private Set<Map.Entry<ResourceNode, ResourceNode>> getResourceDependencyForChannel(DataTransferChannel ch, DataFlowGraph dataFlowGraph) {
        Set<Map.Entry<ResourceNode, ResourceNode>> resourceDependency = new HashSet<>();
        if (!ch.getOutputChannelMembers().isEmpty()) {
            try {
                Map<ChannelMember, Map.Entry<ResourcePath, Set<ChannelMember>>> dependency = ch.fillOutsideResourcePaths(ch.getOutputChannelMembers().iterator().next(), JavaCodeGenerator.pullAccessor);
                for (ChannelMember srcMem : dependency.keySet()) {
                    ResourceNode srcNode = dataFlowGraph.getResourceNode(srcMem.getResource());
                    if (srcNode != null) {
                        for (ChannelMember dstMem : dependency.get(srcMem).getValue()) {
                            ResourceNode dstNode = dataFlowGraph.getResourceNode(dstMem.getResource());
                            while (srcNode.getResourceHierarchy().getNumParameters() == 0 && srcNode.getParent() != null) {
                                srcNode = srcNode.getParent();
                            }
                            resourceDependency.add(new AbstractMap.SimpleEntry<>(srcNode, dstNode));
                        }
                    }
                }
            } catch (ParameterizedIdentifierIsFutureWork | ResolvingMultipleDefinitionIsFutureWork | InvalidMessage |
                     UnificationFailed | ValueUndefined e) {
                e.printStackTrace();
            }
        }
        return resourceDependency;
    }

    public DataFlowGraph getDataFlowGraph() {
        return dataFlowGraph;
    }
}
