package io.activej.dataflow.calcite;

import io.activej.codegen.DefiningClassLoader;
import io.activej.dataflow.DataflowClient;
import io.activej.dataflow.ISqlDataflow;
import io.activej.dataflow.calcite.RelToDatasetConverter.ConversionResult;
import io.activej.dataflow.calcite.optimizer.FilterScanTableRule;
import io.activej.dataflow.calcite.optimizer.SortScanTableRule;
import io.activej.dataflow.collector.AbstractCollector;
import io.activej.dataflow.collector.ICollector;
import io.activej.dataflow.collector.MergeCollector;
import io.activej.dataflow.collector.UnionCollector;
import io.activej.dataflow.dataset.Dataset;
import io.activej.dataflow.dataset.LocallySortedDataset;
import io.activej.dataflow.exception.DataflowException;
import io.activej.dataflow.graph.DataflowGraph;
import io.activej.dataflow.graph.Partition;
import io.activej.datastream.processor.transformer.impl.Limiter;
import io.activej.datastream.supplier.StreamSupplier;
import io.activej.datastream.supplier.StreamSuppliers;
import io.activej.promise.Promise;
import io.activej.reactor.AbstractReactive;
import io.activej.reactor.Reactor;
import io.activej.record.Record;
import org.apache.calcite.jdbc.CalciteSchema;
import org.apache.calcite.plan.RelOptPlanner;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.plan.hep.HepProgramBuilder;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.RelRoot;
import org.apache.calcite.rel.metadata.DefaultRelMetadataProvider;
import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider;
import org.apache.calcite.rel.metadata.RelMetadataQueryBase;
import org.apache.calcite.rel.rules.CoreRules;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.parser.SqlParseException;
import org.apache.calcite.sql.parser.SqlParser;
import org.apache.calcite.sql.validate.SqlValidator;
import org.apache.calcite.sql2rel.SqlToRelConverter;
import org.apache.calcite.tools.Program;
import org.apache.calcite.tools.Programs;

import java.util.Collections;
import java.util.List;

import static io.activej.common.Checks.checkNotNull;
import static io.activej.reactor.Reactive.checkInReactorThread;

public final class SqlDataflow extends AbstractReactive implements ISqlDataflow {
	private final DataflowClient client;
	private final List<Partition> partitions;

	private final SqlParser.Config parserConfig;
	private final SqlToRelConverter converter;
	private final SqlValidator validator;
	private final RelOptPlanner planner;
	private final DefiningClassLoader classLoader;

	private final RelTraitSet traits = RelTraitSet.createEmpty();

	private SqlDataflow(
		Reactor reactor, DataflowClient client, List<Partition> partitions, SqlParser.Config parserConfig,
		SqlToRelConverter converter, RelOptPlanner planner, DefiningClassLoader classLoader
	) {
		super(reactor);
		this.client = client;
		this.partitions = partitions;
		this.parserConfig = parserConfig;
		this.converter = converter;
		this.validator = checkNotNull(converter.validator);
		this.planner = planner;
		this.classLoader = classLoader;
	}

	public static SqlDataflow create(
		Reactor reactor, DataflowClient client, List<Partition> partitions, SqlParser.Config parserConfig,
		SqlToRelConverter converter, RelOptPlanner planner, DefiningClassLoader classLoader
	) {
		return new SqlDataflow(reactor, client, partitions, parserConfig, converter, planner, classLoader);
	}

	@Override
	public Promise<StreamSupplier<Record>> query(String sql) {
		checkInReactorThread(this);
		try {
			return Promise.of(queryDataflow(convertToDataset(sql)));
		} catch (DataflowException | SqlParseException e) {
			return Promise.ofException(e);
		}
	}

	public RelNode convertToNode(String sql) throws SqlParseException, DataflowException {
		SqlParser sqlParser = SqlParser.create(sql, parserConfig);
		SqlNode sqlNode = sqlParser.parseQuery();

		sqlNode = validator.validate(sqlNode);

		if (sqlNode.getKind() != SqlKind.SELECT && sqlNode.getKind() != SqlKind.UNION) {
			throw new DataflowException("Only 'SELECT' queries are allowed");
		}

		RelRoot root = convert(sqlNode);

		return optimize(root);
	}

	public Dataset<Record> convertToDataset(String sql, long maxRows) throws SqlParseException, DataflowException {
		RelNode node = convertToNode(sql);
		ConversionResult transformed = convert(node, maxRows);

		return transformed.unmaterializedDataset().materialize(Collections.emptyList());
	}

	public Dataset<Record> convertToDataset(String sql) throws SqlParseException, DataflowException {
		return convertToDataset(sql, -1);
	}

	public ConversionResult convert(RelNode node, long maxRows) {
		return RelToDatasetConverter.convert(classLoader, node, maxRows, partitions.size());
	}

	public ConversionResult convert(RelNode node) {
		return convert(node, -1);
	}

	public StreamSupplier<Record> queryDataflow(Dataset<Record> dataset) {
		checkInReactorThread(this);
		return queryDataflow(dataset, Limiter.NO_LIMIT);
	}

	public StreamSupplier<Record> queryDataflow(Dataset<Record> dataset, long limit) {
		checkInReactorThread(this);
		if (limit == 0) {
			return StreamSuppliers.empty();
		}

		ICollector<Record> calciteCollector = createCollector(dataset, limit);

		DataflowGraph graph = new DataflowGraph(reactor, client, partitions);
		StreamSupplier<Record> result = calciteCollector.compile(graph);

		graph.execute();
		return result;
	}

	private RelRoot convert(SqlNode sqlNode) {
		if (RelMetadataQueryBase.THREAD_PROVIDERS.get() == null) {
			RelMetadataQueryBase.THREAD_PROVIDERS.set(JaninoRelMetadataProvider.DEFAULT);
		}
		return converter.convertQuery(sqlNode, false, true);
	}

	private RelNode optimize(RelRoot root) {
		HepProgramBuilder builder = new HepProgramBuilder();
		builder.addRuleCollection(List.of(CoreRules.FILTER_INTO_JOIN,
			FilterScanTableRule.create(),
			SortScanTableRule.create())
		);
		Program prog = Programs.of(builder.build(), true, DefaultRelMetadataProvider.INSTANCE);
		Program program = Programs.sequence(prog);
		return program.run(planner, root.project(), traits, Collections.emptyList(), Collections.emptyList());
	}

	public CalciteSchema getSchema() {
		return validator.getCatalogReader().getRootSchema();
	}

	public RelDataTypeFactory getTypeFactory() {
		return validator.getTypeFactory();
	}

	public String explainPlan(String query) throws SqlParseException, DataflowException {
		return convertToNode(query).explain();
	}

	public String explainGraph(String query, long maxRows) throws SqlParseException, DataflowException {
		return convertToDataset(query, maxRows).toGraphViz();
	}

	public String explainGraph(String query) throws SqlParseException, DataflowException {
		return explainGraph(query, -1);
	}

	public String explainNodes(String query, long maxRows) throws SqlParseException, DataflowException {
		Dataset<Record> dataset = convertToDataset(query, maxRows);
		ICollector<Record> calciteCollector = createCollector(dataset, maxRows);

		DataflowGraph graph = new DataflowGraph(reactor, client, partitions);
		calciteCollector.compile(graph);

		return graph.toGraphViz();
	}

	public String explainNodes(String query) throws SqlParseException, DataflowException {
		return explainNodes(query, -1);
	}

	private <
		B extends AbstractCollector<Record, ?>.Builder<B, C>,
		C extends AbstractCollector<Record, ?>
		>
	ICollector<Record> createCollector(Dataset<Record> dataset, long limit) {
		//noinspection unchecked
		B collectorBuilder = (B) (dataset instanceof LocallySortedDataset<?, Record> sortedDataset ?
			MergeCollector.builder(reactor, sortedDataset, client) :
			UnionCollector.builder(reactor, dataset, client));

		if (limit != Limiter.NO_LIMIT) {
			collectorBuilder.withLimit(limit);
		}

		return collectorBuilder.build();
	}
}
