/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.shardingsphere.data.pipeline.core.ingest.dumper.inventory;

import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.data.pipeline.core.channel.PipelineChannel;
import org.apache.shardingsphere.data.pipeline.core.constant.PipelineSQLOperationType;
import org.apache.shardingsphere.data.pipeline.core.datasource.PipelineDataSource;
import org.apache.shardingsphere.data.pipeline.core.exception.IngestException;
import org.apache.shardingsphere.data.pipeline.core.exception.param.PipelineInvalidParameterException;
import org.apache.shardingsphere.data.pipeline.core.execute.AbstractPipelineLifecycleRunnable;
import org.apache.shardingsphere.data.pipeline.core.ingest.dumper.Dumper;
import org.apache.shardingsphere.data.pipeline.core.ingest.dumper.inventory.column.InventoryColumnValueReaderEngine;
import org.apache.shardingsphere.data.pipeline.core.ingest.dumper.inventory.position.InventoryDataRecordPositionCreator;
import org.apache.shardingsphere.data.pipeline.core.ingest.dumper.inventory.query.QueryType;
import org.apache.shardingsphere.data.pipeline.core.ingest.dumper.inventory.query.Range;
import org.apache.shardingsphere.data.pipeline.core.ingest.dumper.inventory.query.StreamingRangeType;
import org.apache.shardingsphere.data.pipeline.core.ingest.dumper.inventory.query.calculator.AbstractRecordTableInventoryCalculator;
import org.apache.shardingsphere.data.pipeline.core.ingest.dumper.inventory.query.calculator.TableInventoryCalculateParameter;
import org.apache.shardingsphere.data.pipeline.core.ingest.position.IngestPosition;
import org.apache.shardingsphere.data.pipeline.core.ingest.position.type.finished.IngestFinishedPosition;
import org.apache.shardingsphere.data.pipeline.core.ingest.position.type.pk.UniqueKeyIngestPosition;
import org.apache.shardingsphere.data.pipeline.core.ingest.record.Column;
import org.apache.shardingsphere.data.pipeline.core.ingest.record.DataRecord;
import org.apache.shardingsphere.data.pipeline.core.ingest.record.FinishedRecord;
import org.apache.shardingsphere.data.pipeline.core.ingest.record.NormalColumn;
import org.apache.shardingsphere.data.pipeline.core.ingest.record.Record;
import org.apache.shardingsphere.data.pipeline.core.query.JDBCStreamQueryBuilder;
import org.apache.shardingsphere.data.pipeline.core.ratelimit.JobRateLimitAlgorithm;
import org.apache.shardingsphere.data.pipeline.core.sqlbuilder.sql.PipelineInventoryDumpSQLBuilder;
import org.apache.shardingsphere.data.pipeline.core.util.PipelineJdbcUtils;
import org.apache.shardingsphere.database.connector.core.type.DatabaseType;
import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation;
import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions;
import org.apache.shardingsphere.infra.metadata.database.schema.QualifiedTable;
import org.apache.shardingsphere.infra.metadata.identifier.ShardingSphereIdentifier;
import org.apache.shardingsphere.infra.util.close.QuietlyCloser;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Inventory dumper.
 */
@HighFrequencyInvocation
@Slf4j
public final class InventoryDumper extends AbstractPipelineLifecycleRunnable implements Dumper {
    
    private final InventoryDumperContext dumperContext;
    
    private final PipelineChannel channel;
    
    private final PipelineDataSource dataSource;
    
    private final InventoryDataRecordPositionCreator positionCreator;
    
    private final PipelineInventoryDumpSQLBuilder sqlBuilder;
    
    private final InventoryColumnValueReaderEngine columnValueReaderEngine;
    
    private final AtomicReference<Statement> runningStatement = new AtomicReference<>();
    
    public InventoryDumper(final InventoryDumperContext dumperContext, final PipelineChannel channel, final PipelineDataSource dataSource, final InventoryDataRecordPositionCreator positionCreator) {
        this.dumperContext = dumperContext;
        this.channel = channel;
        this.dataSource = dataSource;
        this.positionCreator = positionCreator;
        DatabaseType databaseType = dumperContext.getCommonContext().getDataSourceConfig().getDatabaseType();
        sqlBuilder = new PipelineInventoryDumpSQLBuilder(databaseType);
        columnValueReaderEngine = new InventoryColumnValueReaderEngine(databaseType);
    }
    
    @Override
    protected void runBlocking() {
        IngestPosition position = dumperContext.getCommonContext().getPosition();
        if (position instanceof IngestFinishedPosition) {
            log.info("Ignored because of already finished.");
            return;
        }
        try {
            if (dumperContext.hasUniqueKey()) {
                dumpByCalculator();
            } else {
                dumpWithStreamingQuery();
            }
            // CHECKSTYLE:OFF
        } catch (final SQLException | RuntimeException ex) {
            // CHECKSTYLE:ON
            log.error("Inventory dump failed on {}", dumperContext.getActualTableName(), ex);
            throw new IngestException("Inventory dump failed on " + dumperContext.getActualTableName(), ex);
        }
    }
    
    private void dumpByCalculator() {
        String schemaName = dumperContext.getCommonContext().getTableAndSchemaNameMapper().getSchemaName(dumperContext.getLogicTableName());
        QualifiedTable table = new QualifiedTable(schemaName, dumperContext.getActualTableName());
        IngestPosition initialPosition = dumperContext.getCommonContext().getPosition();
        log.info("Dump by calculator start, dataSource={}, table={}, initialPosition={}", dumperContext.getCommonContext().getDataSourceName(), table, initialPosition);
        List<String> columnNames = dumperContext.getQueryColumnNames();
        TableInventoryCalculateParameter calculateParam = new TableInventoryCalculateParameter(dataSource, table,
                columnNames, dumperContext.getUniqueKeyColumns(), QueryType.RANGE_QUERY, null);
        Range<?> range = Range.closed(((UniqueKeyIngestPosition<?>) initialPosition).getLowerBound(), ((UniqueKeyIngestPosition<?>) initialPosition).getUpperBound());
        calculateParam.setRange(range);
        RecordTableInventoryDumpCalculator dumpCalculator = new RecordTableInventoryDumpCalculator(dumperContext.getBatchSize());
        long rowCount = 0L;
        try {
            JobRateLimitAlgorithm rateLimitAlgorithm = dumperContext.getRateLimitAlgorithm();
            String firstUniqueKey = calculateParam.getFirstUniqueKey().getName();
            for (List<DataRecord> each : dumpCalculator.calculate(calculateParam)) {
                if (null != rateLimitAlgorithm) {
                    rateLimitAlgorithm.intercept(PipelineSQLOperationType.SELECT, 1);
                }
                channel.push(Collections.unmodifiableList(each));
                IngestPosition position = UniqueKeyIngestPosition.newInstance(Range.closed(dumpCalculator.getFirstUniqueKeyValue(each.get(each.size() - 1), firstUniqueKey), range.getUpperBound()));
                dumperContext.getCommonContext().setPosition(position);
                rowCount += each.size();
            }
        } finally {
            QuietlyCloser.close(calculateParam.getCalculationContext());
        }
        IngestPosition position = new IngestFinishedPosition();
        channel.push(Collections.singletonList(new FinishedRecord(position)));
        dumperContext.getCommonContext().setPosition(position);
        log.info("Dump by calculator done, rowCount={}, dataSource={}, table={}, initialPosition={}", rowCount, dumperContext.getCommonContext().getDataSourceName(), table, initialPosition);
    }
    
    private void dumpWithStreamingQuery() throws SQLException {
        DatabaseType databaseType = dumperContext.getCommonContext().getDataSourceConfig().getDatabaseType();
        try (Connection connection = dataSource.getConnection()) {
            fetchAllNoUniqueKeyQuery(connection, databaseType, dumperContext.getBatchSize());
        }
    }
    
    private void fetchAllNoUniqueKeyQuery(final Connection connection, final DatabaseType databaseType, final int batchSize) throws SQLException {
        log.info("Start to fetch all no unique key query, dataSource={}, actualTable={}", dumperContext.getCommonContext().getDataSourceName(), dumperContext.getActualTableName());
        try (PreparedStatement statement = JDBCStreamQueryBuilder.build(databaseType, connection, buildFetchAllNoUniqueKeySQL(), batchSize)) {
            runningStatement.set(statement);
            try (ResultSet resultSet = statement.executeQuery()) {
                consumeResultSetToChannel(resultSet, batchSize);
            } finally {
                runningStatement.set(null);
            }
        }
        log.info("End to fetch all no unique key query, dataSource={}, actualTable={}", dumperContext.getCommonContext().getDataSourceName(), dumperContext.getActualTableName());
    }
    
    private String buildFetchAllNoUniqueKeySQL() {
        String schemaName = dumperContext.getCommonContext().getTableAndSchemaNameMapper().getSchemaName(dumperContext.getLogicTableName());
        List<String> columnNames = dumperContext.getQueryColumnNames();
        return sqlBuilder.buildFetchAllSQL(schemaName, dumperContext.getActualTableName(), columnNames);
    }
    
    private void consumeResultSetToChannel(final ResultSet resultSet, final int batchSize) throws SQLException {
        long rowCount = 0;
        JobRateLimitAlgorithm rateLimitAlgorithm = dumperContext.getRateLimitAlgorithm();
        ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
        List<Record> dataRecords = new LinkedList<>();
        while (resultSet.next()) {
            if (dataRecords.size() >= batchSize) {
                channel.push(dataRecords);
                dataRecords = new LinkedList<>();
            }
            dataRecords.add(loadDataRecord(resultSet, resultSetMetaData));
            ++rowCount;
            if (!isRunning()) {
                log.info("Broke because of inventory dump is not running.");
                break;
            }
            if (null != rateLimitAlgorithm && 0 == rowCount % batchSize) {
                rateLimitAlgorithm.intercept(PipelineSQLOperationType.SELECT, 1);
            }
        }
        dataRecords.add(new FinishedRecord(new IngestFinishedPosition()));
        channel.push(dataRecords);
        log.info("Inventory dump with streaming query done, rowCount={}, dataSource={}, actualTable={}", rowCount, dumperContext.getCommonContext().getDataSourceName(),
                dumperContext.getActualTableName());
    }
    
    private DataRecord loadDataRecord(final ResultSet resultSet, final ResultSetMetaData resultSetMetaData) throws SQLException {
        int columnCount = resultSetMetaData.getColumnCount();
        String tableName = dumperContext.getLogicTableName();
        DataRecord result = new DataRecord(PipelineSQLOperationType.INSERT, tableName, positionCreator.create(dumperContext, resultSet), columnCount);
        List<String> insertColumnNames = Optional.ofNullable(dumperContext.getInsertColumnNames()).orElse(Collections.emptyList());
        ShardingSpherePreconditions.checkState(insertColumnNames.isEmpty() || insertColumnNames.size() == resultSetMetaData.getColumnCount(),
                () -> new PipelineInvalidParameterException("Insert column names count not equals ResultSet column count"));
        for (int i = 1; i <= columnCount; i++) {
            String columnName = insertColumnNames.isEmpty() ? resultSetMetaData.getColumnName(i) : insertColumnNames.get(i - 1);
            Column column = getColumn(resultSet, resultSetMetaData, columnName, i, dumperContext.getTargetUniqueKeysNames().contains(new ShardingSphereIdentifier(columnName)));
            result.addColumn(column);
        }
        result.setActualTableName(dumperContext.getActualTableName());
        return result;
    }
    
    private Column getColumn(final ResultSet resultSet, final ResultSetMetaData resultSetMetaData, final String columnName, final int columnIndex, final boolean isUniqueKey) throws SQLException {
        return new NormalColumn(columnName, columnValueReaderEngine.read(resultSet, resultSetMetaData, columnIndex), true, isUniqueKey);
    }
    
    @Override
    protected void doStop() {
        Optional.ofNullable(runningStatement.get()).ifPresent(PipelineJdbcUtils::cancelStatement);
    }
    
    private class RecordTableInventoryDumpCalculator extends AbstractRecordTableInventoryCalculator<List<DataRecord>, DataRecord> {
        
        RecordTableInventoryDumpCalculator(final int chunkSize) {
            super(chunkSize, StreamingRangeType.SMALL);
        }
        
        @Override
        protected DataRecord readRecord(final ResultSet resultSet, final ResultSetMetaData resultSetMetaData, final InventoryColumnValueReaderEngine columnValueReaderEngine) throws SQLException {
            return loadDataRecord(resultSet, resultSetMetaData);
        }
        
        @Override
        protected Object getFirstUniqueKeyValue(final DataRecord record, final String firstUniqueKey) {
            return record.getColumn(firstUniqueKey).getValue();
        }
        
        @Override
        protected List<DataRecord> convertRecordsToResult(final List<DataRecord> records, final Object maxUniqueKeyValue) {
            return records;
        }
    }
}
