/**
 * Copyright (C) <2019>  <chen junwen>
 * <p>
 * This program is free software: you can redistribute it and/or modify it under the terms of the
 * GNU General Public License as published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 * <p>
 * This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without
 * even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 * <p>
 * You should have received a copy of the GNU General Public License along with this program.  If
 * not, see <http://www.gnu.org/licenses/>.
 */
package io.mycat.calcite;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import io.mycat.calcite.table.MycatLogicTable;
import io.mycat.calcite.table.MycatPhysicalTable;
import io.mycat.calcite.table.MycatReflectiveSchema;
import io.mycat.metadata.SchemaHandler;
import io.mycat.TableHandler;
import io.mycat.upondb.*;
import org.apache.calcite.DataContext;
import org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.calcite.jdbc.CalciteSchema;
import org.apache.calcite.linq4j.QueryProvider;
import org.apache.calcite.materialize.SqlStatisticProvider;
import org.apache.calcite.plan.Context;
import org.apache.calcite.plan.RelOptCostFactory;
import org.apache.calcite.plan.RelOptTable;
import org.apache.calcite.plan.RelTraitDef;
import org.apache.calcite.rel.type.RelDataTypeSystem;
import org.apache.calcite.rex.RexExecutor;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.schema.Table;
import org.apache.calcite.schema.impl.AbstractSchema;
import org.apache.calcite.sql.SqlOperatorTable;
import org.apache.calcite.sql.parser.SqlParser;
import org.apache.calcite.sql.validate.SqlValidator;
import org.apache.calcite.sql2rel.SqlRexConvertletTable;
import org.apache.calcite.sql2rel.SqlToRelConverter;
import org.apache.calcite.tools.FrameworkConfig;
import org.apache.calcite.tools.Program;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;

/**
 * @author Junwen Chen
 **/

public class MycatCalciteDataContext implements DataContext, FrameworkConfig {
    private final MycatDBContext uponDBContext;
    private Map<String, Object> variables;

    public MycatCalciteDataContext(MycatDBContext uponDBContext) {
        this.uponDBContext = uponDBContext;
    }

    private ImmutableMap<String, Object> getCalciteLocalVariable() {
        final long time = System.currentTimeMillis();
        TimeZone timeZone = TimeZone.getDefault();
        final long localOffset = timeZone.getOffset(time);
        final long currentOffset = localOffset;
        final String systemUser = System.getProperty("user.name");
        final String user = "sa";
        final Locale locale = Locale.getDefault();
        ImmutableMap.Builder<String, Object> builder = ImmutableMap.builder();
        builder.put(Variable.UTC_TIMESTAMP.camelName, time)
                .put(Variable.CURRENT_TIMESTAMP.camelName, time + currentOffset)
                .put(Variable.LOCAL_TIMESTAMP.camelName, time + localOffset)
                .put(Variable.TIME_ZONE.camelName, timeZone)
                .put(Variable.USER.camelName, user)
                .put(Variable.SYSTEM_USER.camelName, systemUser)
                .put(Variable.LOCALE.camelName, locale)
                .put(Variable.STDIN.camelName, System.in)
                .put(Variable.STDOUT.camelName, System.out)
                .put(Variable.STDERR.camelName, System.err)
                .put(Variable.CANCEL_FLAG.camelName, uponDBContext.cancelFlag());
        return builder.build();
    }

    public SchemaPlus getRootSchema() {
        MycatDBSharedServer uponDBSharedServer = uponDBContext.getUponDBSharedServer();
        Function<Byte, SchemaPlus> function = aByte -> getSchema(uponDBContext);
        if (uponDBContext.config().isCache()) {
            return uponDBSharedServer.getComponent(Components.SCHEMA, function);
        } else {
            return function.apply(Components.SCHEMA);
        }
    }

    public JavaTypeFactory getTypeFactory() {
        return MycatCalciteSupport.INSTANCE.TypeFactory;
    }

    public QueryProvider getQueryProvider() {
        return null;
    }

    public Object get(String name) {
        Object o = uponDBContext.getVariable(name);
        if (o == null) {
            Map<String, Object> variables = uponDBContext.variables();
            if (variables != null) {
                Object o1 = variables.get(name);
                if (o1 != null) {
                    return o1;
                }
            }
        }
        if (variables == null) {
            variables = getCalciteLocalVariable();
        }
        return variables.get(name);
    }


//    public void preComputation(PreComputationSQLTable preComputationSQLTable) {
//        uponDBContext.cache(preComputationSQLTable, preComputationSQLTable.getTargetName(),preComputationSQLTable.getSql(),
//                Collections.emptyList(),()->preComputationSQLTable.scan(this).toList());
//    }
//
//    public Enumerable<Object[]> getPreComputation(PreComputationSQLTable preComputationSQLTable) {
//        Object o = uponDBContext.getCache(preComputationSQLTable,preComputationSQLTable.getTargetName(),preComputationSQLTable.getSql(),Collections.emptyList());
//        if (o != null) {
//            return Linq4j.asEnumerable((List<Object[]>) o);
//        } else {
//            return null;
//        }
//    }
//
//    public UpdateRowIteratorResponse getUpdateRowIterator(String targetName, List<String> sqls) {
//        return uponDBContext.update(targetName, sqls);
//    }

    public static SchemaPlus getSchema(MycatDBClientBased based) {
        SchemaPlus plus = CalciteSchema.createRootSchema(true).plus();
        MycatDBClientBasedConfig config = based.config();
        for (Map.Entry<String, SchemaHandler> stringConcurrentHashMapEntry : config.getSchemaMap().entrySet()) {
            SchemaPlus schemaPlus = plus.add(stringConcurrentHashMapEntry.getKey(), new AbstractSchema());
            for (Map.Entry<String, TableHandler> entry : stringConcurrentHashMapEntry.getValue().logicTables().entrySet()) {
                TableHandler logicTable = entry.getValue();
                MycatLogicTable mycatLogicTable = new MycatLogicTable(logicTable);
                schemaPlus.add(entry.getKey(), mycatLogicTable);
            }
        }
        config.getReflectiveSchemas().forEach((key, value) -> plus.add(key, new MycatReflectiveSchema(value)));
        return plus;
    }

    @Override
    public SqlParser.Config getParserConfig() {
        return MycatCalciteSupport.INSTANCE.config.getParserConfig();
    }

    @Override
    public SqlValidator.Config getSqlValidatorConfig() {
        return SqlValidator.Config.DEFAULT;
    }

    @Override
    public SqlToRelConverter.Config getSqlToRelConverterConfig() {
        return MycatCalciteSupport.INSTANCE.config.getSqlToRelConverterConfig();
    }

    @Override
    public SchemaPlus getDefaultSchema() {
        String schema = uponDBContext.getSchema();
        if (schema == null) {
            return getRootSchema();
        } else {
            return getRootSchema().getSubSchema(schema);
        }
    }

    @Override
    public RexExecutor getExecutor() {
        return (rexBuilder, constExps, reducedValues) -> {
            RexExecutor executor = MycatCalciteSupport.INSTANCE.config.getExecutor();
            if (executor != null) {
                executor.reduce(rexBuilder, constExps, reducedValues);
            }
        };
    }

    @Override
    public ImmutableList<Program> getPrograms() {
        return MycatCalciteSupport.INSTANCE.config.getPrograms();
    }

    @Override
    public SqlOperatorTable getOperatorTable() {
        return MycatCalciteSupport.INSTANCE.config.getOperatorTable();
    }

    @Override
    public RelOptCostFactory getCostFactory() {
        return MycatCalciteSupport.INSTANCE.config.getCostFactory();
    }

    @Override
    public ImmutableList<RelTraitDef> getTraitDefs() {
        return MycatCalciteSupport.INSTANCE.config.getTraitDefs();
    }

    @Override
    public SqlRexConvertletTable getConvertletTable() {
        return MycatCalciteSupport.INSTANCE.config.getConvertletTable();
    }

    @Override
    public Context getContext() {
        return MycatCalciteSupport.INSTANCE;
    }

    @Override
    public RelDataTypeSystem getTypeSystem() {
        return MycatCalciteSupport.INSTANCE.TypeSystem;
    }

    @Override
    public boolean isEvolveLattice() {
        return MycatCalciteSupport.INSTANCE.config.isEvolveLattice();
    }

    @Override
    public SqlStatisticProvider getStatisticProvider() {
        return MycatCalciteSupport.INSTANCE.config.getStatisticProvider();
    }

    @Override
    public RelOptTable.ViewExpander getViewExpander() {
        return MycatCalciteSupport.INSTANCE.config.getViewExpander();
    }


    public MycatDBContext getUponDBContext() {
        return uponDBContext;
    }

    public MycatLogicTable getLogicTable(String targetName, String schema, String table) {
        String uniqueName = targetName + "." + schema + "." + table;
        SchemaPlus rootSchema = getRootSchema();
        Set<String> subSchemaNames = rootSchema.getSubSchemaNames();
        for (String subSchemaName : subSchemaNames) {
            SchemaPlus subSchema = rootSchema.getSubSchema(subSchemaName);
            log.debug("schemaName:{}", subSchemaName);
            Set<String> tableNames = subSchema.getTableNames();
            log.debug("tableNames:{}", tableNames);
            for (String tableName : tableNames) {
                Table table1 = subSchema.getTable(tableName);
                if (table1 instanceof MycatLogicTable) {
                    Map<String, MycatPhysicalTable> dataNodeMap = ((MycatLogicTable) table1).getDataNodeMap();
                    log.debug("dataNodeMap:{}", dataNodeMap);
                    if (dataNodeMap.containsKey(uniqueName)) {
                        return Objects.requireNonNull((MycatLogicTable) table1);
                    }
                }
            }
        }
        return null;
    }

    final static Logger log = LoggerFactory.getLogger(MycatCalciteDataContext.class);

    public AtomicBoolean getCancelFlag() {
        return DataContext.Variable.CANCEL_FLAG.get(this);
    }


}