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

import java.io.File;
import java.io.IOException;
import java.io.PrintStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

import org.dbsp.sqlCompiler.circuit.DBSPCircuit;
import org.dbsp.sqlCompiler.compiler.CompilerOptions;
import org.dbsp.sqlCompiler.compiler.DBSPCompiler;
import org.dbsp.sqlCompiler.compiler.backend.rust.BaseRustCodeGenerator;
import org.dbsp.sqlCompiler.compiler.backend.rust.RustFileWriter;
import org.dbsp.sqlCompiler.compiler.backend.rust.RustWriter;
import org.dbsp.sqlCompiler.compiler.errors.CompilationError;
import org.dbsp.sqlCompiler.compiler.errors.UnimplementedException;
import org.dbsp.sqlCompiler.compiler.visitors.outer.LateMaterializations;
import org.dbsp.sqlCompiler.ir.IDBSPInnerNode;
import org.dbsp.util.IIndentStream;
import org.dbsp.util.Utilities;

/**
 * This class helps generate Rust code. It is given a set of circuit and
 * functions and generates code in multiple crates. */
public final class MultiCratesWriter extends RustWriter {
    public final String outputDirectory;
    public final String pipelineName;
    File rootDirectory;

    /**
     * Create a writer which will generate code in the specified directory
     */
    public MultiCratesWriter(String outputDirectory, String pipelineName, boolean preserve) {
        this.outputDirectory = outputDirectory;
        this.pipelineName = pipelineName;
        this.rootDirectory = this.rootDirectory(preserve);
    }

    File rootDirectory(boolean preserve) {
        if (this.rootDirectory == null) {
            File file = new File(this.outputDirectory);
            if (!file.exists() || !file.isDirectory()) {
                throw new CompilationError("Output directory does not exist "
                        + Utilities.singleQuote(this.outputDirectory));
            }
            if (!preserve) {
                Utilities.deleteContents(file);
            }
            this.rootDirectory = file;
        }
        return this.rootDirectory;
    }

    @Override
    public void setOutputBuilder(IIndentStream stream) {
        throw new UnimplementedException();
    }

    public void write(List<CrateGenerator> generators, CompilerOptions options) throws IOException {
        File cargo = new File(this.rootDirectory, CrateGenerator.CARGO);
        PrintStream cargoStream = new PrintStream(Files.newOutputStream(cargo.toPath()));
        cargoStream.println("# Automatically generated by sql-to-dbsp");
        cargoStream.println("[workspace]");
        cargoStream.println("members = [");
        boolean first = true;
        for (CrateGenerator gen: generators) {
            Path path = Paths.get(gen.directory, gen.crateName);
            cargoStream.println((first ? "  " : ", ") + Utilities.doubleQuote(path.toString(), false));
            first = false;
        }
        cargoStream.println("]");

        String deps = """
                resolver = "2"

                [profile.unoptimized]
                inherits = "release"
                opt-level = 0
                lto = "off"
                codegen-units = 256

                [profile.optimized]
                inherits = "release"

                [profile.optimized_symbols]
                inherits = "release"
                debug = "line-tables-only"

                [workspace.dependencies]
                arcstr = { version = "1.2.0" }
                paste = { version = "1.0.12" }
                derive_more = { version = "1.0.0", features = ["add", "not", "from"] }
                dbsp = { path = "$ROOT/crates/dbsp", features = ["backend-mode"] }
                dbsp_adapters = { path = "$ROOT/crates/adapters"$FEATURES }
                feldera-macros = { path = "$ROOT/crates/feldera-macros" }
                feldera-types = { path = "$ROOT/crates/feldera-types" }
                feldera-adapterlib = { path = "$ROOT/crates/adapterlib" }
                feldera-sqllib = { path = "$ROOT/crates/sqllib" }
                serde = { version = "1.0.213", features = ["derive"] }
                compare = { version = "0.1.0" }
                size-of = { version = "0.1.7", package = "feldera-size-of" }
                serde_json = { version = "1.0.132", features = ["arbitrary_precision"] }
                rkyv = { version = "0.7.45", default-features = false, features = ["std", "size_64"] }
                seq-macro = { version = "0.3.6" }
                tikv-jemallocator = { version = "0.6.0", features = ["profiling", "unprefixed_malloc_on_supported_platforms"] }""";

        String relativePath = "../..";
        if (!options.ioOptions.runtimePath.isEmpty())
            relativePath = options.ioOptions.runtimePath;
        deps = deps.replace("$ROOT", relativePath);
        String enterpriseFeatures = options.ioOptions.enterprise ?
                ", features = [\"feldera-enterprise\"] " : "";
        deps = deps.replace("$FEATURES", enterpriseFeatures);
        cargoStream.println(deps);

        if (options.ioOptions.enterprise) {
            cargoStream.println("""
                dbsp-enterprise = { path = "$ROOT/crates/dbsp-enterprise" }
                sync-checkpoint = { path = "$ROOT/crates/sync-checkpoint" }
            """.replace("$ROOT", relativePath));
        }

        cargoStream.close();
    }

    @Override
    public void write(DBSPCompiler compiler) throws IOException {
        StructuresUsed used = this.analyze(compiler);
        LateMaterializations materializations = new LateMaterializations(compiler);
        for (var node: this.toWrite) {
            if (node.is(DBSPCircuit.class))
                materializations.apply(node.to(DBSPCircuit.class));
        }
        MultiCrates crates = new MultiCrates(this.rootDirectory, this.pipelineName, compiler, used, materializations);
        crates.addNodes(this.toWrite);
        List<CrateGenerator> topLevel = new ArrayList<>();
        topLevel.add(crates.main);

        if (!this.testNodes.isEmpty()) {
            BaseRustCodeGenerator writer = new RustFileWriter(materializations)
                    .withUdf(false).withMalloc(false).withGenerateTuples(false);
            CrateGenerator test = new CrateGenerator(
                    this.rootDirectory, MultiCrates.CRATES_DIRECTORY, MultiCratesWriter.getTestName(), writer,
                    crates.enterprise(), true);
            RustWriter.StructuresUsed locallyUsed = new RustWriter.StructuresUsed();
            FindInnerResources finder = new FindInnerResources(compiler, locallyUsed);

            for (var node : this.testNodes) {
                if (node.is(IDBSPInnerNode.class))
                    finder.apply(node.to(IDBSPInnerNode.class));
                test.add(node);
            }
            test.addDependency(crates.main);
            for (int i : locallyUsed.tupleSizesUsed) {
                if (!used.isPredefined(i)) {
                    CrateGenerator gen = Utilities.getExists(crates.tupleCrates, i);
                    test.addDependency(gen);
                }
            }

            test.write(compiler);
            topLevel.add(test);
        }

        crates.write();
        this.write(topLevel, compiler.options);
    }

    public String getGlobalsName() {
        // Keep in sync with MultiCrates.getGlobalsName
        return MultiCrates.FILE_PREFIX + this.pipelineName + "_globals";
    }

    public static String getTestName() {
        return "feldera_test";
    }

    public Path stubsDirectory() {
        return new File(new File(
                new File(this.rootDirectory, MultiCrates.CRATES_DIRECTORY), this.getGlobalsName()), "src")
                .toPath();
    }
}
