package org.dbsp.sqlCompiler.circuit.operator;

import com.fasterxml.jackson.databind.JsonNode;
import org.dbsp.sqlCompiler.circuit.OutputPort;
import org.dbsp.sqlCompiler.compiler.DBSPCompiler;
import org.dbsp.sqlCompiler.compiler.backend.JsonDecoder;
import org.dbsp.sqlCompiler.compiler.errors.InternalCompilerError;
import org.dbsp.sqlCompiler.compiler.frontend.calciteObject.CalciteEmptyRel;
import org.dbsp.sqlCompiler.compiler.frontend.calciteObject.CalciteObject;
import org.dbsp.sqlCompiler.compiler.frontend.calciteObject.CalciteRelNode;
import org.dbsp.sqlCompiler.compiler.visitors.VisitDecision;
import org.dbsp.sqlCompiler.compiler.visitors.inner.EquivalenceContext;
import org.dbsp.sqlCompiler.compiler.visitors.inner.Expensive;
import org.dbsp.sqlCompiler.compiler.visitors.inner.InnerVisitor;
import org.dbsp.sqlCompiler.compiler.visitors.inner.Simplify;
import org.dbsp.sqlCompiler.compiler.visitors.outer.CircuitVisitor;
import org.dbsp.sqlCompiler.ir.expression.DBSPBlockExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPClosureExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPIfExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPRawTupleExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPReturnExpression;
import org.dbsp.sqlCompiler.ir.expression.DBSPVariablePath;
import org.dbsp.sqlCompiler.ir.statement.DBSPExpressionStatement;
import org.dbsp.sqlCompiler.ir.statement.DBSPLetStatement;
import org.dbsp.sqlCompiler.ir.statement.DBSPStatement;
import org.dbsp.sqlCompiler.ir.type.DBSPType;
import org.dbsp.sqlCompiler.ir.type.derived.DBSPTypeRawTuple;
import org.dbsp.sqlCompiler.ir.type.user.DBSPTypeIndexedZSet;
import org.dbsp.sqlCompiler.ir.type.user.DBSPTypeZSet;
import org.dbsp.util.IndentStreamBuilder;
import org.dbsp.util.Linq;
import org.dbsp.util.Maybe;
import org.dbsp.util.Utilities;

import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.List;

/** A Chain operator performs a linear chain of Map/Filter/MapIndex operations.
 * In the end it is lowered to a {@link DBSPFlatMapOperator} or
 * {@link DBSPFlatMapIndexOperator} operator, depending on the last operation. */
public class DBSPChainOperator extends DBSPUnaryOperator implements ILinear {
    public final ComputationChain chain;

    public DBSPChainOperator(CalciteRelNode node, ComputationChain chain, boolean isMultiset, OutputPort source) {
        super(node, "chain" + chain.summary(), null, chain.getOutputType(), isMultiset, source);
        this.chain = chain;
        Utilities.enforce(this.chain.size() > 1);
    }

    public enum ComputationKind {
        Map,
        Filter,
        MapIndex
    }

    public record Computation(ComputationKind kind, DBSPClosureExpression closure) {
        public Computation {
            Utilities.enforce(closure.parameters.length == 1);
        }

        void checkInputType(DBSPType inputType) {
            DBSPTypeZSet zSet = inputType.as(DBSPTypeZSet.class);
            DBSPTypeIndexedZSet iZSet = inputType.as(DBSPTypeIndexedZSet.class);
            DBSPType argType;
            if (zSet != null) {
                argType = zSet.elementType.ref();
            } else {
                Utilities.enforce(iZSet != null);
                argType = iZSet.getKVRefType();
            }
            DBSPType paramType = closure.parameters[0].getType();
            if (!paramType.sameType(argType))
                throw new InternalCompilerError("Expected function to accept\n" + argType +
                        " as argument " + 0 + " but it expects\n" + paramType);
        }

        boolean equivalent(Computation o) {
            return this.kind == o.kind && EquivalenceContext.equiv(this.closure, o.closure);
        }

        public static Computation fromJson(JsonNode node, JsonDecoder decoder) {
            ComputationKind kind = ComputationKind.valueOf(Utilities.getStringProperty(node, "kind"));
            DBSPClosureExpression closure = fromJsonInner(node, "closure", decoder, DBSPClosureExpression.class);
            return new Computation(kind, closure);
        }

        public char summary() {
            return switch (this.kind) {
                case Map -> 'M';
                case Filter -> 'F';
                case MapIndex -> 'I';
            };
        }
    }

    /** A chain of functions; if the chain is empty, it is equivalent to the identity function */
    public record ComputationChain(DBSPType inputType, List<Computation> computations) {
        public ComputationChain(DBSPType inputType) {
            this(inputType, new ArrayList<>());
        }

        DBSPType getOutputType(int index) {
            Computation computation = this.computations.get(index);
            return switch (computation.kind) {
                case Map -> new DBSPTypeZSet(computation.closure.getResultType());
                case Filter -> (index == 0) ? this.inputType : this.getOutputType(index - 1);
                case MapIndex -> new DBSPTypeIndexedZSet(
                        computation.closure.getResultType().to(DBSPTypeRawTuple.class));
            };
        }

        DBSPType getOutputType() {
            if (this.computations.isEmpty())
                return inputType;
            return this.getOutputType(this.computations.size() - 1);
        }

        public ComputationChain add(Computation computation) {
            computation.checkInputType(this.getOutputType());
            return new ComputationChain(this.inputType, Linq.append(this.computations, computation));
        }

        public static ComputationChain fromJson(JsonNode node, JsonDecoder decoder) {
            DBSPType inputType = fromJsonInner(node, "inputType", decoder, DBSPType.class);
            List<Computation> computations = Linq.list(Linq.map(
                    node.elements(), e -> Computation.fromJson(e, decoder)));
            return new ComputationChain(inputType, computations);
        }

        DBSPExpression call(DBSPCompiler compiler, DBSPClosureExpression closure, DBSPExpression arg) {
            if (arg.getType().is(DBSPTypeRawTuple.class)) {
                return closure.call(new DBSPRawTupleExpression(
                        arg.field(0).borrow(),
                        arg.field(1).borrow()))
                        .reduce(compiler);
            } else {
                return closure.call(arg.borrow()).reduce(compiler);
            }
        }

        public boolean containsFilter() {
            return Linq.any(this.computations, c -> c.kind() == ComputationKind.Filter);
        }

        /** Return a function that is equivalent to the composition of the functions in the chain.
         * If the chain contains any filters the result will have an Option type result. */
        public DBSPClosureExpression collapse(DBSPCompiler compiler) {
            if (this.computations.size() == 1) {
                Computation comp = this.computations.get(0);
                if (comp.kind != ComputationKind.Filter)
                    return comp.closure;
            }

            DBSPVariablePath inputVar;
            DBSPExpression currentArg;
            if (this.inputType.is(DBSPTypeZSet.class)) {
                inputVar = this.inputType.to(DBSPTypeZSet.class).getElementType().ref().var();
                currentArg = inputVar.deref();
            } else {
                inputVar = this.inputType.to(DBSPTypeIndexedZSet.class).getKVRefType().var();
                currentArg = new DBSPRawTupleExpression(
                        inputVar.field(0).deref(), inputVar.field(1).deref());
            }
            // Value returned when a filter fails
            DBSPExpression none;
            DBSPType resultType = this.getOutputType();
            if (resultType.is(DBSPTypeZSet.class))
                resultType = resultType.to(DBSPTypeZSet.class).elementType;
            else
                resultType = resultType.to(DBSPTypeIndexedZSet.class).getKVType();
            none = resultType.withMayBeNull(true).none();
            boolean foundFilter = false;
            boolean foundNonFilter = false;
            List<DBSPStatement> statements = new ArrayList<>();
            for (Computation comp: this.computations) {
                CalciteObject node = comp.closure().getNode();
                DBSPStatement stat;
                switch (comp.kind) {
                    case Map, MapIndex: {
                        foundNonFilter = true;
                        DBSPVariablePath nextVar = comp.closure.getResultType().var();
                        stat = new DBSPLetStatement(nextVar.variable,
                                this.call(compiler, comp.closure, currentArg));
                        currentArg = nextVar;
                        break;
                    }
                    case Filter: {
                        foundFilter = true;
                        DBSPExpression condition = this.call(compiler, comp.closure, currentArg).not();
                        Simplify simplify = new Simplify(compiler);
                        condition = simplify.apply(condition).to(DBSPExpression.class);
                        DBSPExpression cond = new DBSPIfExpression(node, condition,
                                new DBSPReturnExpression(node, none),
                                null).simplify();
                        stat = new DBSPExpressionStatement(cond);
                        break;
                    }
                    default: {
                        throw new InternalCompilerError("unreachable");
                    }
                }
                statements.add(stat);
            }
            if (!foundNonFilter) {
                currentArg = currentArg.applyCloneIfNeeded();
            }
            DBSPExpression last = (foundFilter) ? currentArg.some() : currentArg;
            DBSPBlockExpression block = new DBSPBlockExpression(statements, last);
            return block.closure(inputVar);
        }

        public int size() {
            return this.computations.size();
        }

        boolean equivalent(ComputationChain chain) {
            if (this.computations.size() != chain.size())
                return false;
            for (int i = 0; i < this.computations.size(); i++) {
                Computation c = this.computations.get(i);
                Computation o = chain.computations.get(i);
                if (!c.equivalent(o))
                    return false;
            }
            return true;
        }

        public String summary() {
            StringBuilder builder = new StringBuilder();
            for (var op: this.computations)
                builder.append(op.summary());
            return builder.toString();
        }

        @Override
        public String toString() {
            IndentStreamBuilder builder = new IndentStreamBuilder();
            builder.append("Chain[")
                    .append(this.size())
                    .append("]: [").increase();
            for (var op: this.computations)
                builder.append(op.toString()).newline();
            builder.decrease().newline();
            return builder.toString();
        }

        /** Compose pairs of maps that can be efficiently composed, taking into advantage
         * the fact that function composition is associative. */
        public DBSPChainOperator.ComputationChain shrinkMaps(DBSPCompiler compiler) {
            List<DBSPChainOperator.Computation> result = new ArrayList<>();
            for (DBSPChainOperator.Computation comp: this.computations()) {
                if (result.isEmpty() || comp.kind() == DBSPChainOperator.ComputationKind.Filter) {
                    result.add(comp);
                } else {
                    DBSPChainOperator.Computation last = Utilities.removeLast(result);
                    if (last.kind() == DBSPChainOperator.ComputationKind.Filter) {
                        result.add(last);
                        result.add(comp);
                        continue;
                    }

                    boolean expensive = Expensive.isExpensive(compiler, last.closure());
                    if (expensive) {
                        result.add(last);
                    } else {
                        DBSPClosureExpression composed;
                        if (last.kind() == DBSPChainOperator.ComputationKind.Map) {
                            composed = comp.closure().applyAfter(compiler, last.closure(), Maybe.MAYBE);
                        } else {
                            DBSPClosureExpression lastFunction = last.closure();
                            DBSPExpression argument = new DBSPRawTupleExpression(
                                    lastFunction.body.field(0).borrow(),
                                    lastFunction.body.field(1).borrow());
                            DBSPExpression apply = comp.closure().call(argument);
                            composed = apply.reduce(compiler)
                                    .closure(lastFunction.parameters);
                        }
                        comp = new DBSPChainOperator.Computation(comp.kind(), composed);
                    }
                    result.add(comp);
                }
            }

            if (result.size() == this.size())
                return this;
            return new DBSPChainOperator.ComputationChain(this.inputType(), result);
        }

        /** Convert Map(m1) -> Filter(f) -> Map(m2) into Filter(f \circ m1) -> Map(m2 \circ m1) if m1 is simple */
        public ComputationChain shrinkMapFilterMap(DBSPCompiler compiler) {
            List<DBSPChainOperator.Computation> result = new ArrayList<>();
            if (this.size() < 3) {
                return this;
            }

            // Find a sequence Map -> Filter -> Map/MapIndex
            int startIndex = -1;
            for (int i = 0; i < this.size() - 2; i++) {
                if (this.computations().get(i).kind() == DBSPChainOperator.ComputationKind.Map &&
                        this.computations().get(i+1).kind() == DBSPChainOperator.ComputationKind.Filter &&
                        this.computations().get(i+2).kind() != DBSPChainOperator.ComputationKind.Filter) {
                    DBSPClosureExpression map = this.computations().get(i).closure();
                    if (this.computations().get(i+1).closure().shouldInlineComposition(compiler, map) &&
                            this.computations().get(i+2).closure().shouldInlineComposition(compiler, map)) {
                        startIndex = i;
                        break;
                    }
                }
                result.add(this.computations().get(i));
            }

            if (startIndex < 0)
                return this;

            DBSPChainOperator.Computation first = this.computations().get(startIndex);
            DBSPChainOperator.Computation filter = this.computations().get(startIndex + 1);
            DBSPChainOperator.Computation third = this.computations().get(startIndex + 2);

            DBSPClosureExpression filterMap = filter.closure().applyAfter(compiler, first.closure(), Maybe.MAYBE);
            DBSPClosureExpression mapMap = third.closure().applyAfter(compiler, first.closure(), Maybe.MAYBE);
            result.add(new DBSPChainOperator.Computation(DBSPChainOperator.ComputationKind.Filter, filterMap));
            result.add(new DBSPChainOperator.Computation(third.kind(), mapMap));
            // Keep the subsequent unchanged
            for (int i = startIndex + 3; i < this.size(); i++) {
                result.add(this.computations().get(i));
            }

            return new DBSPChainOperator.ComputationChain(this.inputType(), result);
        }
    }

    @Override
    public DBSPSimpleOperator with(
            @Nullable DBSPExpression function, DBSPType outputType,
            List<OutputPort> newInputs, boolean force) {
        if (this.mustReplace(force, function, newInputs, outputType)) {
            Utilities.enforce(function == null);
            return new DBSPChainOperator(this.getRelNode(), this.chain, this.isMultiset, newInputs.get(0))
                    .copyAnnotations(this);
        }
        return this;
    }

    @Override
    public void accept(CircuitVisitor visitor) {
        visitor.push(this);
        VisitDecision decision = visitor.preorder(this);
        if (!decision.stop())
            visitor.postorder(this);
        visitor.pop(this);
    }

    @Override
    public void accept(InnerVisitor visitor) {
        visitor.startArrayProperty("chain");
        for (var operation: this.chain.computations())
            operation.closure().accept(visitor);
        visitor.endArrayProperty("chain");
    }

    @Override
    public boolean equivalent(DBSPOperator other) {
        if (!super.equivalent(other))
            return false;
        DBSPChainOperator otherOperator = other.as(DBSPChainOperator.class);
        if (otherOperator == null)
            return false;
        return this.chain.equivalent(otherOperator.chain);
    }

    @SuppressWarnings("unused")
    public static DBSPChainOperator fromJson(JsonNode node, JsonDecoder decoder) {
        CommonInfo info = DBSPSimpleOperator.commonInfoFromJson(node, decoder);
        ComputationChain chain = ComputationChain.fromJson(node.get("chain"), decoder);
        return new DBSPChainOperator(CalciteEmptyRel.INSTANCE,
                chain, info.isMultiset(), info.getInput(0))
                .addAnnotations(info.annotations(), DBSPChainOperator.class);
    }
}
