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

import org.dbsp.sqlCompiler.compiler.DBSPCompiler;
import org.dbsp.sqlCompiler.compiler.backend.rust.ICodeGenerator;
import org.dbsp.sqlCompiler.ir.IDBSPNode;
import org.dbsp.util.IndentStream;
import org.dbsp.util.IndentStreamBuilder;
import org.dbsp.util.Linq;
import org.dbsp.util.Utilities;

import java.io.File;
import java.io.IOException;
import java.io.PrintStream;
import java.nio.file.Files;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/** Generates a Rust crate in a specified directory.
 * Calls a ICodeGenerator to generate the Rust code. */
public final class CrateGenerator {
    public final File baseDirectory;
    public final String directory;
    public final String crateName;
    public final boolean enterprise;
    public final boolean lib;

    /** Cargo file name */
    public static final String CARGO = "Cargo.toml";
    /** Rust file name */
    public static final String LIB = "lib.rs";
    public static final String MAIN = "main.rs";

    /** Crates that we depend on */
    private final Set<CrateGenerator> dependencies;
    /** Generates the actual Rust code */
    final ICodeGenerator codeGenerator;

    public CrateGenerator(File baseDirectory, String directory, String crateName, ICodeGenerator codeGenerator,
                          boolean enterprise, boolean lib) {
        this.crateName = crateName;
        this.directory = directory;
        this.baseDirectory = baseDirectory;
        this.dependencies = new HashSet<>();
        this.codeGenerator = codeGenerator;
        this.enterprise = enterprise;
        this.lib = lib;
    }

    @Override
    public boolean equals(Object o) {
        if (o == null || getClass() != o.getClass()) return false;

        CrateGenerator that = (CrateGenerator) o;
        return this.baseDirectory.equals(that.baseDirectory) &&
                this.crateName.equals(that.crateName) &&
                this.directory.equals(that.directory);
    }

    @Override
    public int hashCode() {
        int result = this.baseDirectory.hashCode();
        result = 31 * result + this.crateName.hashCode();
        result = 31 * result + this.directory.hashCode();
        return result;
    }

    public void addDependency(CrateGenerator generator) {
        this.dependencies.add(generator);
        this.codeGenerator.addDependency(generator.crateName);
    }

    boolean isMain() {
        return this.crateName.contains("main");
    }

    void generateCargo(PrintStream stream) {
        stream.println("# Automatically generated by sql-to-dbsp");
        stream.println("[package]");
        stream.print("name = \"");
        stream.print(this.crateName);
        stream.println("\"");

        stream.println("""
                version = "0.1.0"
                edition = "2021"
                publish = false

                [dev-dependencies]
                uuid = { version = "1.6.1" }""");
        stream.println();
        if (this.lib) {
            stream.println("[lib]");
            stream.println("path = \"src/lib.rs\"");
        } else {
            stream.println("[[bin]]");
            stream.println("name = \"" + this.crateName + "\"");
            stream.println("path = \"src/main.rs\"");
        }
        stream.println("""
                doctest = false

                [dependencies]
                arcstr = { workspace = true }
                paste = { workspace = true }
                derive_more = { workspace = true }
                dbsp = { workspace = true }
                dbsp_adapters = { workspace = true }
                feldera-macros = { workspace = true }
                feldera-adapterlib = { workspace = true }
                feldera-types = { workspace = true }
                feldera-sqllib = { workspace = true }
                serde = { workspace = true }
                compare = { workspace = true }
                size-of = { workspace = true }
                serde_json = { workspace = true }
                seq-macro = { workspace = true }
                rkyv = { workspace = true }""");
        if (this.enterprise) {
            stream.println("dbsp-enterprise = { workspace = true }");
            stream.println("sync-checkpoint = { workspace = true }");
        }
        if (isMain()) {
            stream.println("""
                    [target.'cfg(not(target_env = "msvc"))'.dependencies]
                    tikv-jemallocator = { workspace = true }
                    """);
        }
        List<CrateGenerator> deps = Linq.list(this.dependencies);
        deps.sort(Comparator.comparing(a -> a.crateName));
        for (CrateGenerator dep: deps) {
            stream.println(dep.crateName + " = { path = " + Utilities.doubleQuote("../" + dep.crateName, false) + " }");
        }
    }

    public void add(IDBSPNode node) {
        this.codeGenerator.add(node);
    }

    public void write(DBSPCompiler compiler) throws IOException {
        if (!this.baseDirectory.exists())
            throw new RuntimeException(
                    "Directory " + Utilities.singleQuote(this.baseDirectory.getPath()) + " does not exist");
        if (!this.baseDirectory.isDirectory())
            throw new RuntimeException(
                    Utilities.singleQuote(this.baseDirectory.getPath()) + " is not a directory");
        File crateRoot = new File(new File(this.baseDirectory, this.directory), this.crateName);
        if (!crateRoot.exists()) {
            boolean success = crateRoot.mkdirs();
            if (!success)
                throw new RuntimeException("Could not create directory " + Utilities.singleQuote(crateRoot.getPath()));
        }
        File cargo = new File(crateRoot, CARGO);
        PrintStream cargoStream = new PrintStream(Files.newOutputStream(cargo.toPath()));
        this.generateCargo(cargoStream);
        cargoStream.close();

        File src = new File(crateRoot, "src");
        if (!src.exists()) {
            boolean success = src.mkdirs();
            if (!success)
                throw new RuntimeException("Could not create directory " + Utilities.singleQuote(src.getPath()));
        }
        File lib;
        if (!this.isMain() || !compiler.options.generateMultiCrateMain())
            lib = new File(src, LIB);
        else
            lib = new File(src, MAIN);
        PrintStream rustStream = new PrintStream(Files.newOutputStream(lib.toPath()));
        this.codeGenerator.setOutputBuilder(new IndentStream(rustStream));
        this.codeGenerator.write(compiler);
        rustStream.close();
    }

    public String dump(DBSPCompiler compiler) throws IOException {
        IndentStreamBuilder builder = new IndentStreamBuilder();
        this.codeGenerator.setOutputBuilder(builder);
        this.codeGenerator.write(compiler);
        return builder.toString();
    }
}
