package org.dbsp.sqlCompiler.compiler.backend.dot;

import org.dbsp.sqlCompiler.circuit.operator.DBSPAggregateOperatorBase;
import org.dbsp.sqlCompiler.circuit.operator.DBSPChainOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPConstantOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPFlatMapOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPIndexedTopKOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPInputMapWithWaterlineOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPInternOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPJoinFilterMapOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPLeftJoinFilterMapOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPNestedOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPOperatorWithError;
import org.dbsp.sqlCompiler.circuit.operator.DBSPPartitionedRollingAggregateWithWaterlineOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPSimpleOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPSourceBaseOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPViewBaseOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPWaterlineOperator;
import org.dbsp.sqlCompiler.circuit.operator.IGCOperator;
import org.dbsp.sqlCompiler.circuit.operator.IStateful;
import org.dbsp.sqlCompiler.compiler.DBSPCompiler;
import org.dbsp.sqlCompiler.compiler.backend.rust.ToRustInnerVisitor;
import org.dbsp.sqlCompiler.compiler.errors.SourcePositionRange;
import org.dbsp.sqlCompiler.compiler.visitors.VisitDecision;
import org.dbsp.sqlCompiler.compiler.visitors.outer.CircuitVisitor;
import org.dbsp.sqlCompiler.compiler.visitors.outer.LowerCircuitVisitor;
import org.dbsp.sqlCompiler.compiler.visitors.outer.ToJsonVisitor;
import org.dbsp.sqlCompiler.ir.IDBSPInnerNode;
import org.dbsp.sqlCompiler.ir.IDBSPOuterNode;
import org.dbsp.sqlCompiler.ir.expression.DBSPExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPFlatmap;
import org.dbsp.util.IndentStream;

/** Visitor which emits the circuit nodes in a graphviz file.
 * The compiler options control the detail.  On details=0 table and view names are omitted. */
public class ToDotNodesVisitor extends CircuitVisitor {
    protected final IndentStream stream;
    // A higher value -> more details
    protected final int details;
    static long lastCircuit = 0;

    public ToDotNodesVisitor(DBSPCompiler compiler, IndentStream stream, int details) {
        super(compiler);
        this.stream = stream;
        this.details = details;
    }

    static String isMultiset(DBSPSimpleOperator operator) {
        return operator.isMultiset ? "" : "*";
    }

    static String annotations(DBSPOperator operator) {
        return operator.annotations.toDotString();
    }

    long currentCircuit = 0;

    public Token startVisit(IDBSPOuterNode node) {
        this.currentCircuit = node.getId();
        return super.startVisit(node);
    }

    @Override
    public void endVisit() {
        super.endVisit();
        lastCircuit = this.currentCircuit;
    }

    @Override
    public VisitDecision preorder(DBSPSourceBaseOperator node) {
        String name = (this.details > 0 ? (node.tableName + " ") : "") + node.operation;
        String color = " style=filled fillcolor=lightgrey";
        this.stream.append(node.getNodeName(false))
                .append(" [ shape=box " + color + " label=\"")
                .append(node.getIdString())
                .append(isMultiset(node))
                .append(annotations(node))
                .append(" ")
                .append(name)
                .append("\" ]")
                .newline();
        return VisitDecision.STOP;
    }

    @Override
    public VisitDecision preorder(DBSPConstantOperator node) {
        this.stream.append(node.getNodeName(false))
                .append(" [ shape=box,label=\"")
                .append(node.getIdString())
                .append(isMultiset(node))
                .append(annotations(node))
                .append(" ");
        if (this.details >= 2)
            this.stream.append(getFunction(node));
         else
            this.stream.append("constant ...");
        this.stream.append("\" ]")
                .newline();
        return VisitDecision.STOP;
    }

    @Override
    public VisitDecision preorder(DBSPViewBaseOperator node) {
        this.stream.append(node.getNodeName(false))
                .append(" [ shape=box,label=\"")
                .append(node.getIdString())
                .append(isMultiset(node))
                .append(annotations(node))
                .append(this.details > 0 ? " " + node.viewName.name() : "")
                .append("\"")
                .append(" style=filled fillcolor=lightgrey")
                .append("]")
                .newline();
        return VisitDecision.STOP;
    }

    @Override
    public VisitDecision preorder(DBSPInternOperator node) {
        this.stream.append(node.getNodeName(false))
                .append(" [ shape=box,label=\"")
                .append(node.getIdString())
                .append(" ")
                .append(node.operation)
                .append(" ")
                .append("\"")
                .append(" style=filled fillcolor=lightgrey")
                .append("]")
                .newline();
        return VisitDecision.STOP;
    }

    static String escapeString(String input) {
        StringBuilder escapedString = new StringBuilder();
        for (char c : input.toCharArray()) {
            switch (c) {
                case '"':
                    escapedString.append("\\\"");
                    break;
                case '\r':
                    escapedString.append("\\r");
                    break;
                case '\\':
                    escapedString.append("\\\\");
                    break;
                default:
                    escapedString.append(c);
            }
        }

        return escapedString.toString();
    }

    String getPositions(IDBSPInnerNode node) {
        StringBuilder result = new StringBuilder();
        ToJsonVisitor.FindSourcePositions finder = new ToJsonVisitor.FindSourcePositions(this.compiler, true);
        finder.apply(node);
        for (SourcePositionRange r : finder.getPositions()) {
            result.append(r.toShortString());
        }
        return result.toString();
    }

    static String rustToDot(String rust) {
        String f = escapeString(rust);
        return f.replace("\n", "\\l");
    }

    String convertFunction(DBSPExpression expression) {
        String result = "";
        if (this.details > 3) {
            String f = ToRustInnerVisitor.toRustString(this.compiler(), expression, null, true);
            result = rustToDot(f);
        }
        if (this.details >= 3) {
            result += getPositions(expression);
        }
        return result;
    }

    String getFunction(DBSPSimpleOperator node) {
        DBSPExpression expression = node.function;
        if (node.is(DBSPAggregateOperatorBase.class)) {
            DBSPAggregateOperatorBase aggregate = node.to(DBSPAggregateOperatorBase.class);
            if (aggregate.aggregateList != null) {
                if (this.details > 3) {
                    return rustToDot(aggregate.aggregateList.toString());
                } else if (details >= 3) {
                    return this.getPositions(aggregate.aggregateList);
                }
            }
        } else if (node.is(DBSPPartitionedRollingAggregateWithWaterlineOperator.class)) {
            DBSPPartitionedRollingAggregateWithWaterlineOperator aggregate =
                    node.to(DBSPPartitionedRollingAggregateWithWaterlineOperator.class);
            if (aggregate.aggregateList != null) {
                if (this.details > 3) {
                    return rustToDot(aggregate.aggregateList.toString());
                } else if (this.details >= 3) {
                    return this.getPositions(aggregate.aggregateList);
                }
            }
        } else if (node.is(DBSPChainOperator.class) && this.details > 3) {
            String chain = node.to(DBSPChainOperator.class).chain.toString();
            return rustToDot(chain);
        }
        if (expression == null) {
            if (this.details >= 3) {
                StringBuilder builder = new StringBuilder();
                for (SourcePositionRange pos: node.getSourcePositions())
                    builder.append(pos.toShortString());
                return builder.toString();
            } else {
                return "";
            }
        }
        if (node.is(DBSPFlatMapOperator.class)) {
            if (expression.is(DBSPFlatmap.class)) {
                expression = LowerCircuitVisitor.rewriteFlatmap(
                        expression.to(DBSPFlatmap.class), this.compiler);
            }
        }
        if (node.is(DBSPJoinFilterMapOperator.class)) {
            DBSPJoinFilterMapOperator jfm = node.to(DBSPJoinFilterMapOperator.class);
            expression = LowerCircuitVisitor.lowerJoinFilterMapFunctions(
                    this.compiler(),
                    jfm.getClosureFunction(), jfm.filter, jfm.map);
        }
        if (node.is(DBSPLeftJoinFilterMapOperator.class)) {
            DBSPLeftJoinFilterMapOperator jfm = node.to(DBSPLeftJoinFilterMapOperator.class);
            expression = LowerCircuitVisitor.lowerJoinFilterMapFunctions(
                    this.compiler(),
                    jfm.getClosureFunction(), jfm.filter, jfm.map);
        }
        return this.convertFunction(expression);
    }

    String getColor(DBSPSimpleOperator operator) {
        /*
        // Uncomment various choices for debugging
        // Show in green projection operators
        if (annotations(operator).contains("IsProjection"))
            return " style=filled fillcolor=green";
        // Show in green operators added since last call
        if (operator.id > lastCircuit)
            return " style=filled fillcolor=green";
         */
        if (operator.is(IGCOperator.class)) {
            return " style=filled fillcolor=pink";
        }
        switch (operator.operation) {
            case "waterline": return " style=filled fillcolor=lightgreen";
            case "controlled_filter": return " style=filled fillcolor=cyan";
            case "apply", "apply2", "apply_n": return " style=filled fillcolor=yellow";
            default: break;
        }
        if (operator.is(IStateful.class)) {
            // There are more of these every day
            return " style=filled fillcolor=orangered";
        }
        return "";
    }

    String shorten(String operation) {
        if (operation.startsWith("accumulate_integrate_trace"))
            return operation.substring("accumulate_integrate_trace_".length());
        if (operation.equals("aggregate_linear_postprocess"))
            return "aggregate_linear";
        return operation;
    }

    @Override
    public VisitDecision preorder(DBSPSimpleOperator node) {
        this.stream.append(node.getNodeName(false))
                .append(" [ shape=box")
                .append(this.getColor(node))
                .append(" label=\"")
                .append(node.getIdString())
                .append(isMultiset(node))
                .append(annotations(node))
                .append(" ")
                .append(shorten(node.operation))
                .append(node.comment != null ? node.comment : "");
        if (this.details >= 3) {
            this.stream
                    .append("(")
                    .append(this.getFunction(node))
                    .append(")\\l");
        }
        this.stream.append("\" ]")
                .newline();
        return VisitDecision.STOP;
    }

    @Override
    public VisitDecision preorder(DBSPIndexedTopKOperator node) {
        this.stream.append(node.getNodeName(false))
                .append(" [ shape=box")
                .append(this.getColor(node))
                .append(" label=\"")
                .append(node.getIdString())
                .append(isMultiset(node))
                .append(annotations(node))
                .append(" ")
                .append(shorten(node.operation))
                .append(node.comment != null ? node.comment : "");
        if (this.details >= 3) {
            this.stream
                    .append("(")
                    .append(this.getFunction(node))
                    .append(", ")
                    .append(this.convertFunction(node.outputProducer))
                    .append(")\\l");
        }
        this.stream.append("\" ]")
                .newline();
        return VisitDecision.STOP;
    }

    @Override
    public VisitDecision preorder(DBSPNestedOperator node) {
        this.stream.append("subgraph cluster_")
                .append(node.id)
                .append(" {").increase()
                .append("color=black;")
                .newline();
        return VisitDecision.CONTINUE;
    }

    @Override
    public void postorder(DBSPNestedOperator node) {
        this.stream.decrease().append("}").newline();
    }

    String quoteSymbols(String str) {
        return str.replace("<", "\\<")
                .replace(">", "\\>")
                .replace("{", "\\{")
                .replace("}", "\\}");
    }

    @Override
    public VisitDecision preorder(DBSPInputMapWithWaterlineOperator node) {
        String name = node.getNodeName(false) + " " + node.operation + " " + node.tableName;
        this.stream.append(node.getNodeName(false))
                .append(" [ shape=record")
                .append(" label=\"<p0>")
                .append(name)
                .append("|<p1> E")
                .append("|<p2> W");
        this.stream.append("\" ]")
                .newline();
        return VisitDecision.STOP;
    }

    @Override
    public VisitDecision preorder(DBSPOperatorWithError node) {
        this.stream.append(node.getNodeName(false))
                .append(" [ shape=record")
                .append(" label=\"<pN>")
                .append(node.getIdString())
                .append("|<p0>")
                .append(shorten(node.operation))
                .append(annotations(node))
                .append("|<p1> E");
        if (this.details >= 3) {
            this.stream
                    .append("(")
                    // additional quoting needed in operators with ports
                    .append(quoteSymbols(this.convertFunction(node.function)))
                    .append(", ")
                    .append(quoteSymbols(this.convertFunction(node.error)))
                    .append(")\\l");
        }
        this.stream.append("\" ]")
                .newline();
        return VisitDecision.STOP;
    }

    @Override
    public VisitDecision preorder(DBSPWaterlineOperator node) {
        this.stream.append(node.getNodeName(false))
                .append(" [ shape=box")
                .append(this.getColor(node))
                .append(" label=\"")
                .append(node.getIdString())
                .append(isMultiset(node))
                .append(annotations(node))
                .append(" ")
                .append(shorten(node.operation))
                .append(node.comment != null ? node.comment : "");
        if (this.details >= 3) {
            this.stream
                    .append("(")
                    .append(this.convertFunction(node.init))
                    .append(", ")
                    .append(this.convertFunction(node.extractTs))
                    .append(", ")
                    .append(this.getFunction(node))
                    .append(")\\l");
        }
        this.stream.append("\" ]")
                .newline();
        return VisitDecision.STOP;
    }
}
