package org.dbsp.sqlCompiler.compiler.backend.rust.multi;

import org.dbsp.sqlCompiler.circuit.DBSPCircuit;
import org.dbsp.sqlCompiler.circuit.annotation.OperatorHash;
import org.dbsp.sqlCompiler.circuit.operator.DBSPControlledKeyFilterOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPNestedOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPSimpleOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPSinkOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPSourceBaseOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPSourceMultisetOperator;
import org.dbsp.sqlCompiler.circuit.operator.DBSPViewBaseOperator;
import org.dbsp.sqlCompiler.circuit.operator.IInputOperator;
import org.dbsp.sqlCompiler.compiler.DBSPCompiler;
import org.dbsp.sqlCompiler.compiler.backend.rust.BaseRustCodeGenerator;
import org.dbsp.sqlCompiler.compiler.backend.rust.SourcePositionResource;
import org.dbsp.sqlCompiler.compiler.backend.rust.ToRustInnerVisitor;
import org.dbsp.sqlCompiler.compiler.backend.rust.ToRustVisitor;
import org.dbsp.sqlCompiler.compiler.visitors.outer.CircuitVisitor;
import org.dbsp.sqlCompiler.compiler.visitors.outer.CollectSourcePositions;
import org.dbsp.sqlCompiler.compiler.visitors.outer.LateMaterializations;
import org.dbsp.sqlCompiler.ir.type.DBSPType;
import org.dbsp.util.HashString;
import org.dbsp.util.IndentStream;
import org.dbsp.util.IndentStreamBuilder;
import org.dbsp.util.Linq;

import java.util.HashSet;
import java.util.List;
import java.util.Set;

/** Code generator for a circuit which instantiates each operator using a function call
 * generated by a {@link SingleOperatorWriter}. */
public final class CircuitWriter extends BaseRustCodeGenerator {
    public static final String SOURCE_MAP_VARIABLE_NAME = "sourceMap";
    final LateMaterializations materializations;

    public CircuitWriter(LateMaterializations materializations) {
        this.materializations = materializations;
    }

    private void processChild(DBSPOperator node, boolean useHandles) {
        DBSPOperator op = node.to(DBSPOperator.class);
        String name = op.getNodeName(false);
        String hash = op.getNodeName(true);
        HashString merkle = OperatorHash.getHash(node, true);
        this.builder().append("let ");
        if (node.is(DBSPSimpleOperator.class)) {
            if (node.is(DBSPSourceBaseOperator.class))
                this.builder().append("(");
            this.builder().append(name);
            if (node.is(DBSPSourceBaseOperator.class))
                this.builder().append(", handle_").append(name).append(")");
        } else {
            this.builder().append("(");
            for (int i = 0; i < node.outputCount(); i++) {
                if (node.hasOutput(i)) {
                    String portName = node.getOutput(i).getName(false);
                    this.builder().append(portName)
                            .append(",");
                }
            }
            this.builder().append(")");
        }
        this.builder().append(" = ");
        this.builder().append("create_")
                .append(hash)
                .append("(&circuit, ");
        if (!node.is(DBSPNestedOperator.class)) {
            if (merkle != null) {
                this.builder().append("Some(\"")
                        .append(merkle.toString())
                        .append("\"), ");
            } else {
                this.builder().append("None, ");
            }
        }
        this.builder().append("&").append(SOURCE_MAP_VARIABLE_NAME).append(", ");
        if (!useHandles)
            this.builder().append("&mut catalog, ");
        for (var input: op.inputs) {
            name = input.getName(false);
            this.builder().append("&")
                    .append(name)
                    .append(",");
        }
        if (node.is(DBSPControlledKeyFilterOperator.class)) {
            DBSPControlledKeyFilterOperator filter = node.to(DBSPControlledKeyFilterOperator.class);
            if (this.materializations.hasRight(filter)) {
                DBSPSourceMultisetOperator source = this.materializations.getLeft(filter);
                String sourceName = source.getNodeName(false);
                this.builder().append("handle_").append(sourceName);
            }
        }
        this.builder().append(");").newline();
    }

    void writeCircuit(DBSPCompiler compiler, DBSPCircuit circuit) {
        boolean useHandles = compiler.options.ioOptions.emitHandles;
        IndentStream signature = new IndentStreamBuilder();
        ToRustInnerVisitor inner = new ToRustInnerVisitor(compiler, signature, null, false);

        if (!useHandles) {
            signature.append("Catalog");
        } else {
            signature.append("(").increase();
            for (IInputOperator input: circuit.sourceOperators.values()) {
                DBSPType type = input.getHandleType();
                type.accept(inner);
                signature.append(",").newline();
            }
            for (DBSPViewBaseOperator output: circuit.sinkOperators.values()) {
                DBSPType outputType = output.input().outputType();
                signature.append("OutputHandle<SpineSnapshot<");
                outputType.accept(inner);
                signature.append(">>,").newline();
            }
            signature.decrease().append(")");
        }

        this.builder().append("pub fn ")
                .append(circuit.getName());

        this.builder()
                .append("(cconf: CircuitConfig) -> Result<(DBSPHandle, ")
                .append(signature.toString())
                .append("), Error> {")
                .increase()
                .newline()
                .append("let (circuit, streams) = ");
        if (!useHandles)
            this.builder().append("dbsp_adapters::server::init_circuit(cconf, Box::new(");
        else
            this.builder().append("Runtime::init_circuit(cconf, ");
        this.builder().append("|circuit| {")
                .increase();
        if (!useHandles)
            this.builder().append("let mut catalog = Catalog::new();").newline();

        ToRustVisitor.registerPreprocessors(compiler, this.builder());

        SourcePositionResource sourcePositionResource = new SourcePositionResource();
        CircuitVisitor collector = new CollectSourcePositions(compiler, sourcePositionResource)
                .getCircuitVisitor(true);
        collector.apply(circuit);
        sourcePositionResource.generateInitializer(this.builder());
        SourcePositionResource.generateReference(this.builder(), CircuitWriter.SOURCE_MAP_VARIABLE_NAME);

        // Process sources first
        for (DBSPOperator node : circuit.getAllOperators())
            if (node.is(DBSPSourceBaseOperator.class))
                this.processChild(node, useHandles);

        for (DBSPOperator node : circuit.getAllOperators())
            if (!node.is(DBSPSourceBaseOperator.class))
                this.processChild(node, useHandles);

        if (!useHandles) {
            this.builder().append("Ok(catalog)");
        } else {
            this.builder().append("Ok((");
            for (IInputOperator operator: circuit.sourceOperators.values()) {
                String name = operator.asOperator().getNodeName(false);
                this.builder().append("handle_").append(name).append(", ");
            }
            for (DBSPSinkOperator operator: circuit.sinkOperators.values()) {
                String name = operator.getNodeName(false);
                this.builder().append(name).append(", ");
            }
            this.builder().append("))");
        }
        this.builder().newline()
                .decrease()
                .append("})")
                .append(!useHandles ? ")" : "")
                .append("?;")
                .newline();
        this.builder()
                .append("Ok((circuit, streams))")
                .newline()
                .decrease()
                .append("}")
                .newline();
    }

    @Override
    public void write(DBSPCompiler compiler) {
        this.builder().append(COMMON_PREAMBLE);
        this.builder().append(STANDARD_PREAMBLE);
        this.builder().append(ALLOC_PREAMBLE);
        Set<String> deps = new HashSet<>(this.dependencies);
        List<String> sorted = Linq.list(deps);
        sorted.sort(String::compareTo);
        for (String dep: sorted)
            this.builder().append("use ")
                    .append(dep)
                    .append("::*;")
                    .newline();
        for (var node: this.toWrite) {
            this.writeCircuit(compiler, node.to(DBSPCircuit.class));
        }

        boolean generateMainFunction = compiler.options.generateMultiCrateMain();
        if (generateMainFunction) {
            this.builder().append("""
                    fn main() {
                        dbsp_adapters::server::server_main(|cconfig| {
                            circuit(cconfig)
                                .map(|(dbsp, catalog)| {
                                    (
                                        dbsp,
                                        Box::new(catalog) as Box<dyn dbsp_adapters::CircuitCatalog>,
                                    )
                                })
                                .map_err(|e| dbsp_adapters::ControllerError::dbsp_error(e))
                        })
                        .unwrap_or_else(|e| {
                            eprintln!("{e}");
                            std::process::exit(1);
                        });
                    }
                    """);
            if (compiler.options.ioOptions.enterprise) {
                this.builder().append("extern crate dbsp_enterprise;")
                        .newline()
                        .append("extern crate sync_checkpoint;")
                        .newline();
            }
        }
    }
}
