package ru.yandex.webmaster3.worker.indexing;

import java.util.ArrayList;
import java.util.List;
import java.util.NavigableSet;
import java.util.Optional;
import java.util.Set;
import java.util.TreeSet;
import java.util.UUID;
import java.util.stream.Collectors;

import com.datastax.driver.core.utils.UUIDs;
import lombok.Setter;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.joda.time.LocalDate;
import org.joda.time.format.DateTimeFormat;
import org.joda.time.format.DateTimeFormatter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import ru.yandex.webmaster3.core.util.RetryUtils;
import ru.yandex.webmaster3.core.worker.task.PeriodicTaskState;
import ru.yandex.webmaster3.core.worker.task.PeriodicTaskType;
import ru.yandex.webmaster3.core.worker.task.TaskResult;
import ru.yandex.webmaster3.storage.util.ydb.exception.WebmasterYdbException;
import ru.yandex.webmaster3.storage.clickhouse.replication.MdbClickhouseReplicationManager;
import ru.yandex.webmaster3.storage.clickhouse.replication.data.ClickhouseReplicationCommand;
import ru.yandex.webmaster3.storage.clickhouse.replication.data.ClickhouseReplicationPriority;
import ru.yandex.webmaster3.storage.clickhouse.replication.data.ClickhouseReplicationTaskGroup;
import ru.yandex.webmaster3.storage.clickhouse.system.dao.ClickhouseSystemTablesCHDao;
import ru.yandex.webmaster3.storage.clickhouse.system.data.ClickhouseSystemTableInfo;
import ru.yandex.webmaster3.storage.indexing2.internal.dao.IndexingHistoryTablesRepository;
import ru.yandex.webmaster3.storage.indexing2.internal.data.IndexingHistoryImportType;
import ru.yandex.webmaster3.storage.indexing2.internal.data.IndexingHistoryTableImport;
import ru.yandex.webmaster3.storage.indexing2.internal.data.IndexingHistoryTableImportState;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseException;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseHost;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseQueryContext;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseServer;
import ru.yandex.webmaster3.storage.util.yt.Format;
import ru.yandex.webmaster3.storage.util.yt.YtException;
import ru.yandex.webmaster3.storage.util.yt.YtJobSpec;
import ru.yandex.webmaster3.storage.util.yt.YtMapIOSpec;
import ru.yandex.webmaster3.storage.util.yt.YtMapReduceCommand;
import ru.yandex.webmaster3.storage.util.yt.YtNode;
import ru.yandex.webmaster3.storage.util.yt.YtNodeAttributes;
import ru.yandex.webmaster3.storage.util.yt.YtOperationFilePathSpec;
import ru.yandex.webmaster3.storage.util.yt.YtOperationId;
import ru.yandex.webmaster3.storage.util.yt.YtPath;
import ru.yandex.webmaster3.storage.util.yt.YtReduceIOSpec;
import ru.yandex.webmaster3.storage.util.yt.YtService;
import ru.yandex.webmaster3.storage.util.yt.YtUtils;
import ru.yandex.webmaster3.storage.ytimport.ImportPriority;
import ru.yandex.webmaster3.storage.ytimport.MdbYtClickhouseImportManager;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseImportCommand;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseImportState;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseImportStateEnum;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseTableRelation;
import ru.yandex.webmaster3.worker.PeriodicTask;
import ru.yandex.webmaster3.worker.TaskSchedule;

/**
 * @author avhaliullin
 */
public abstract class AbstractImportIndexingHistoryTask extends PeriodicTask<PeriodicTaskState> {
    private static final Logger log = LoggerFactory.getLogger(AbstractImportIndexingHistoryTask.class);

    private static final YtReduceIOSpec YT_REDUCE_IO_SPEC = new YtReduceIOSpec(1024 * 1024 * 128);//128 Mb
    private static final YtMapIOSpec YT_MAP_IO_SPEC = new YtMapIOSpec(1024 * 1024 * 128);//128 Mb

    private static final DateTimeFormatter TABLE_NAME_DATE_FORMATTER = DateTimeFormat.forPattern("yyyyMMdd");
    private static final DateTimeFormatter CH_INSERTION_DATE_FORMATTER = DateTimeFormat.forPattern("yyyy-MM-dd");
    private static final String TABLE_NAME_PREFIX = "day_";
    private static final String ATTR_LAST_PROCESSED = "last_processed";
    private static final int PREPARED_TABLE_LINES_COUNT = 128;
    private static final int PARTITIONS_COUNT = 16;

    @Setter
    private YtService ytService;
    @Setter
    private ClickhouseServer clickhouseServer;
    @Setter
    private MdbYtClickhouseImportManager ytClickhouseImportManager;
    @Setter
    private MdbClickhouseReplicationManager clickhouseReplicationManager;
    @Setter
    private ClickhouseSystemTablesCHDao clickhouseSystemTablesCHDao;
    @Setter
    private IndexingHistoryTablesRepository indexingHistoryTablesCDao;
    @Setter
    private YtPath ytDataDir;
    @Setter
    private YtPath ytWorkDir;
    @Setter
    private YtPath ytPrepareExecutable;
    @Setter
    private PeriodicTaskType type;
    @Setter
    private String schedule;
    @Setter
    private boolean newMergeSystem;
    @Setter
    private String fullTableName;

    protected abstract String getMrTaskName();

    protected abstract String getTableName(LocalDate date);

    protected abstract String getTableCreateSpec(LocalDate tableDate, Instant updateTime);

    protected abstract String getDBName();

    protected abstract String getInsertSpec();

    protected abstract IndexingHistoryImportType getImportType();

    @Override
    public PeriodicTaskType getType() {
        return type;
    }

    @Override
    public TaskSchedule getSchedule() {
        return TaskSchedule.startByCron(schedule);
    }

    @Override
    public Result run(UUID runId) throws Exception {
        IndexingHistoryTableImport latestImport = getLatestImport();
        if (latestImport != null && !latestImport.getState().isTerminal()) {
            // Последняя таска недоделана
            log.info("Continue processing abandoned task {}", latestImport);
            processImport(latestImport);
        } else if (latestImport != null && latestImport.getState() == IndexingHistoryTableImportState.FAILED) {
            // Последняя таска развалилась
            log.info("Restarting failed task {}", latestImport);
            latestImport = storeState(latestImport.restarted());
            processImport(latestImport);
        } else {
            if (latestImport != null) {
                Instant updateTimeForCurrentTable = getDataUpdateTimeForTableOnYT(latestImport.getTableDate());
                if (updateTimeForCurrentTable.isAfter(latestImport.getDataUpdateTime())) {
                    // Найдено обновление для последней таблицы
                    log.info("Found new data from {} for table {}", latestImport.getTableDate(), updateTimeForCurrentTable);
                    latestImport = IndexingHistoryTableImport.create(getImportType(), latestImport.getTableDate(), updateTimeForCurrentTable);
                    storeState(latestImport);
                    processImport(latestImport);
                } else {
                    // Пытаемся начать импортировать данные за следующий день
                    NavigableSet<LocalDate> tablesOnYT = listTablesFromYT();
                    LocalDate nextTableDate = tablesOnYT.higher(latestImport.getTableDate());
                    if (nextTableDate == null) {
                        log.info("No new data found");
                        return new Result(TaskResult.SUCCESS);
                    }
                    Instant updateTime = getDataUpdateTimeForTableOnYT(nextTableDate);
                    log.info("Starting import for next day's table {}", nextTableDate);
                    latestImport = IndexingHistoryTableImport.create(getImportType(), nextTableDate, updateTime);
                    storeState(latestImport);
                    processImport(latestImport);
                }
            } else {
                // Первый импорт, грузим доступные таблицы с самого начала
                NavigableSet<LocalDate> availableDates = listTablesFromYT();
                if (availableDates.isEmpty()) {
                    log.info("No data tables found on YT");
                    return new Result(TaskResult.SUCCESS);
                }
                LocalDate firstTable = availableDates.first();
                Instant updateDate = getDataUpdateTimeForTableOnYT(firstTable);
                latestImport = IndexingHistoryTableImport.create(getImportType(), firstTable, updateDate);
                storeState(latestImport);
                processImport(latestImport);
            }
        }
        return new Result(TaskResult.SUCCESS);
    }

    private void processImport(IndexingHistoryTableImport imprt) throws YtException, InterruptedException, ClickhouseException, WebmasterYdbException {
        assureWorkDirExists();
        while (!imprt.getState().isTerminal()) {
            switch (imprt.getState()) {
                case PREPARING:
                    imprt = prepareTable(imprt)
                            .withState(IndexingHistoryTableImportState.IMPORTING);
                    break;
                case IMPORTING:
                    imprt = importTable(imprt)
                            .withState(IndexingHistoryTableImportState.REPLICATING);
                    break;
                case REPLICATING:
                    imprt = replicateTable(imprt)
                            .withState(IndexingHistoryTableImportState.RENAMING);
                    break;
                case RENAMING:
                    imprt = renameTable(imprt)
                            .withState(IndexingHistoryTableImportState.CLEANING);
                    break;
                case CLEANING:
                    imprt = cleanup(imprt)
                            .withState(IndexingHistoryTableImportState.DONE);
                    break;
                default:
                    throw new RuntimeException("Unknown non-terminal state " + imprt.getState());
            }
            storeState(imprt);
        }
    }

    private int getShardsCount() {
        return clickhouseServer.getShardsCount();
    }

    private IndexingHistoryTableImport prepareTable(IndexingHistoryTableImport imprt) throws YtException {
        int shardsCount = getShardsCount();
        List<YtPath> outTables = new ArrayList<>();
        String tableNamePrefix = imprt.getTableDate().toString(TABLE_NAME_DATE_FORMATTER) + "_" + imprt.getDataUpdateTime().getMillis() + "_";
        for (int shard = 0; shard < shardsCount; shard++) {
            outTables.add(YtPath.path(ytWorkDir, tableNamePrefix + shard));
        }
        ytService.inTransaction(ytWorkDir).execute(cypressService -> {
            YtUtils.recreateTables(cypressService, outTables, new YtNodeAttributes().setCompressionCodec("none"));
            YtMapReduceCommand mrCommand = YtMapReduceCommand.newBuilder()
                    .setOutputTables(outTables)
                    .addInputTable(getSourceTablePath(imprt.getTableDate()))
                    .reduceBy("part_id", "line_id")
                    .sortBy("part_id", "line_id", "row_id")
                    .setReduceJobIo(YT_REDUCE_IO_SPEC)
                    .setMapperJobIo(YT_MAP_IO_SPEC)
                    .setMapper(
                            YtJobSpec.newBuilder()
                                    .setCommand(
                                            YtOperationFilePathSpec.newBuilder()
                                                    .setName("mapper")
                                                    .setPath(ytPrepareExecutable)
                                                    .setExecutable(true)
                                                    .build(),
                                            "--task", getMrTaskName(),
                                            "--stage", "map",
                                            "--tables", String.valueOf(outTables.size()),
                                            "--lines", String.valueOf(PREPARED_TABLE_LINES_COUNT),
                                            "--date", imprt.getTableDate().toString(CH_INSERTION_DATE_FORMATTER)
                                    )
                                    .setInputFormat(Format.createBinaryYson())
                                    .setOutputFormat(Format.createBinaryYson())
                                    .build()
                    )
                    .setReducer(
                            YtJobSpec.newBuilder()
                                    .setCommand(
                                            YtOperationFilePathSpec.newBuilder()
                                                    .setName("reducer")
                                                    .setPath(ytPrepareExecutable)
                                                    .setExecutable(true)
                                                    .build(),
                                            "--task", getMrTaskName(),
                                            "--tables", String.valueOf(outTables.size()),
                                            "--stage", "reduce",
                                            "--date", imprt.getTableDate().toString(CH_INSERTION_DATE_FORMATTER)
                                    )
                                    .setInputFormat(Format.createBinaryYson())
                                    .setOutputFormat(Format.createBinaryYson())
                                    .build()
                    )
                    .build();
            YtOperationId operationId = cypressService.mapReduce(mrCommand);
            if (!cypressService.waitFor(operationId)) {
                throw new RuntimeException("YT task failed " + operationId);
            }
            return true;
        });
        return imprt.withPreparedTables(outTables);
    }

    private IndexingHistoryTableImport importTable(IndexingHistoryTableImport imprt) throws InterruptedException, WebmasterYdbException {
        if (imprt.getImportManagerTaskId() == null) {
            UUID importTaskId = UUIDs.timeBased();
            List<YtClickhouseTableRelation> ytClickhouseTableRelations = new ArrayList<>();

            if (imprt.getYtPreparedTables().size() != getShardsCount()) {
                String message = "Got " + imprt.getYtPreparedTables().size() + " prepared tables and " + getShardsCount() + " clickhouse shards";
                log.error(message);
                storeState(imprt.failed());
                throw new RuntimeException(message);
            }

            String tmpTableName = "tmp_" + getTableName(imprt.getTableDate()) + "_" + imprt.getDataUpdateTime().getMillis();
            imprt = storeState(imprt.withImportManagerTaskId(importTaskId).withTmpTable(tmpTableName));

            String createSpec = getTableCreateSpec(
                    imprt.getTableDate(),
                    imprt.getDataUpdateTime()
            );
            for (int i = 0; i < getShardsCount(); i++) {
                ytClickhouseTableRelations.add(new YtClickhouseTableRelation(
                        imprt.getYtPreparedTables().get(i),
                        i,
                        tmpTableName,
                        createSpec
                ));
            }
            YtClickhouseImportCommand importCommand = new YtClickhouseImportCommand(
                    importTaskId,
                    ytClickhouseTableRelations,
                    getDBName(),
                    getInsertSpec(),
                    ImportPriority.ONLINE
            );
            ytClickhouseImportManager.startImport(importCommand);
        }

        YtClickhouseImportState importState = ytClickhouseImportManager.waitForCompletionAndGetState(imprt.getImportManagerTaskId());
        if (importState == null || importState.getState() == YtClickhouseImportStateEnum.FAILED) {
            String message = importState == null
                    ? "Import task " + imprt.getImportManagerTaskId() + " was lost"
                    : "Import task " + imprt.getImportManagerTaskId() + " was failed";
            log.error(message);
            storeState(imprt.failed());
            throw new RuntimeException(message);
        }
        return imprt;
    }

    private IndexingHistoryTableImport replicateTable(IndexingHistoryTableImport imprt) throws InterruptedException, WebmasterYdbException {
        if (newMergeSystem) {
            return imprt;
        }
        if (imprt.getReplicationManagerTaskId() == null) {
            UUID replicationTaskId = UUIDs.timeBased();
            imprt = storeState(imprt.withReplicationManagerTaskId(replicationTaskId));
            List<ClickhouseReplicationCommand.TableInfo> tables = new ArrayList<>();
            for (int chShard = 0; chShard < getShardsCount(); chShard++) {
                tables.add(new ClickhouseReplicationCommand.TableInfo(
                        imprt.getTmpTableName(),
                        getTableCreateSpec(imprt.getTableDate(), imprt.getDataUpdateTime()),
                        chShard
                ));
            }
            clickhouseReplicationManager.enqueueReplication(new ClickhouseReplicationCommand(
                    replicationTaskId,
                    getDBName(),
                    ClickhouseReplicationPriority.ONLINE,
                    tables
            ));
        }
        ClickhouseReplicationTaskGroup taskState = clickhouseReplicationManager.waitForCompletionAndGetState(imprt.getReplicationManagerTaskId());
        if (taskState == null) {
            String message = "Replication task " + imprt.getReplicationManagerTaskId() + " was lost";
            log.error(message);
            storeState(imprt.failed());
            throw new RuntimeException(message);
        }
        return imprt;
    }

    private IndexingHistoryTableImport renameTable(IndexingHistoryTableImport imprt) throws InterruptedException {
        if (newMergeSystem) {
            return mergeTable(imprt);
        }
        String normalTableName = getTableName(imprt.getTableDate());
        String replacedTableName = imprt.getReplacedTableNewName();
        if (replacedTableName == null) {
            imprt = imprt.withReplacedTableNewName("toremove_" + normalTableName + "_" + imprt.getDataUpdateTime().getMillis());
            storeState(imprt);
            replacedTableName = imprt.getReplacedTableNewName();
        }
        String replaceOldQ = String.format("RENAME TABLE %1$s.%2$s TO %1$s.%3$s, %1$s.%4$s TO %1$s.%2$s",
                getDBName(),
                normalTableName,
                replacedTableName,
                imprt.getTmpTableName()
        );
        String switchOnNewQ = String.format("RENAME TABLE %1$s.%2$s TO %1$s.%3$s",
                getDBName(),
                imprt.getTmpTableName(),
                normalTableName
        );
        List<ClickhouseHost> hosts = clickhouseServer.getHosts();
        if (hosts.stream().anyMatch(ClickhouseHost::isDown)) {
            throw new RuntimeException("Will not start renaming: some clickhouse hosts are down");
        }
        for (ClickhouseHost host : hosts) {
            ClickhouseQueryContext.Builder chQCtx = ClickhouseQueryContext.useDefaults()
                    .setHost(host);
            Set<String> availableTables = clickhouseSystemTablesCHDao.getTables(host, getDBName())
                    .stream().map(ClickhouseSystemTableInfo::getName).collect(Collectors.toSet());
            boolean tmpTableExists = availableTables.contains(imprt.getTmpTableName());
            boolean normalTableExists = availableTables.contains(normalTableName);
            if (tmpTableExists) {
                if (normalTableExists) {
                    clickhouseServer.execute(chQCtx, ClickhouseServer.QueryType.INSERT, replaceOldQ, Optional.empty(), Optional.empty());
                } else {
                    clickhouseServer.execute(chQCtx, ClickhouseServer.QueryType.INSERT, switchOnNewQ, Optional.empty(), Optional.empty());
                }
            } else {
                if (!normalTableExists) {
                    String message = "Not found newly imported tmp table " + imprt.getTmpTableName() + ", neither normally-named table " + normalTableName + " on host " + host.getHostURI().toString();
                    log.error(message);
                    storeState(imprt.failed());
                    throw new RuntimeException(message);
                } // Иначе на этом хосте уже переименовали
            }
        }
        return imprt;
    }

    private IndexingHistoryTableImport mergeTable(IndexingHistoryTableImport imprt) throws InterruptedException {
        // удаляем из общей таблицы подливаемую дату и вливаем свежие данные
        for (ClickhouseHost host : clickhouseServer.getHosts()) {
            Set<String> availableTables = clickhouseSystemTablesCHDao.getTables(host, getDBName())
                    .stream().map(ClickhouseSystemTableInfo::getName).collect(Collectors.toSet());
            if (!availableTables.contains(imprt.getTmpTableName())) {
                continue;
            }
            ClickhouseQueryContext.Builder ctx = ClickhouseQueryContext.useDefaults().setHost(host).setTimeout(Duration.standardMinutes(10L));
            RetryUtils.execute(RetryUtils.linearBackoff(3, Duration.standardMinutes(1L)), () -> {
                // удаляем старое
                clickhouseServer.execute(ctx, String.format("ALTER TABLE %s.%s DELETE WHERE date = '%s'", getDBName(), fullTableName, imprt.getTableDate()));
                // вставляем новое
                StringBuilder sb = new StringBuilder();
                sb.append("ALTER TABLE ").append(getDBName()).append(".").append(fullTableName);
                for (int partition = 0; partition < PARTITIONS_COUNT; partition++) {
                    if (partition > 0) {
                        sb.append(",");
                    }
                    sb.append(" ATTACH PARTITION '").append(partition).append("' FROM ").append(getDBName()).append(".").append(imprt.getTmpTableName());
                }
                clickhouseServer.execute(ctx, sb.toString());
                // дропаем временную табличку
                clickhouseServer.execute(ctx, String.format("DROP TABLE %s.%s", getDBName(), imprt.getTmpTableName()));
            });
        }
        return imprt;
    }

    private IndexingHistoryTableImport cleanup(IndexingHistoryTableImport latestImport) throws InterruptedException, YtException, ClickhouseException {
        ytService.withoutTransaction(cypressService -> {
            for (YtPath tmpYtTable : latestImport.getYtPreparedTables()) {
                if (cypressService.exists(tmpYtTable)) {
                    cypressService.remove(tmpYtTable);
                }
            }
            return true;
        });
        String chDropQ = String.format("DROP TABLE IF EXISTS %s.%s", getDBName(), latestImport.getReplacedTableNewName());
        clickhouseServer.executeOnAllHosts(ClickhouseQueryContext.useDefaults(), chDropQ);
        return latestImport;
    }

    private Instant getDataUpdateTimeForTableOnYT(LocalDate table) throws YtException, InterruptedException {
        return ytService.withoutTransactionQuery(cypressService -> {
            YtPath tablePath = getSourceTablePath(table);
            if (!cypressService.exists(tablePath)) {
                return null;
            }
            YtNode result = cypressService.getNode(tablePath);
            return new Instant(result.getNodeMeta().get(ATTR_LAST_PROCESSED).asLong() * 1000L);
        });
    }

    private YtPath getSourceTablePath(LocalDate tableDate) {
        return YtPath.path(ytDataDir, TABLE_NAME_PREFIX + tableDate.toString(TABLE_NAME_DATE_FORMATTER));
    }

    private NavigableSet<LocalDate> listTablesFromYT() throws YtException, InterruptedException {
        return ytService.withoutTransactionQuery(cypressService ->
                cypressService.list(ytDataDir)
                        .stream()
                        .filter(node -> node.getName().startsWith(TABLE_NAME_PREFIX))
                        .map(node -> TABLE_NAME_DATE_FORMATTER.parseLocalDate(node.getName().substring(TABLE_NAME_PREFIX.length())))
                        .collect(Collectors.toCollection(TreeSet::new)));
    }

    private IndexingHistoryTableImport storeState(IndexingHistoryTableImport imprt)  {
        log.info("Storing new state {}", imprt);
        indexingHistoryTablesCDao.insertImportInfo(imprt);
        return imprt;
    }

    private IndexingHistoryTableImport getLatestImport()  {
        return indexingHistoryTablesCDao.getLatestInfoOfType(getImportType());
    }

    private void assureWorkDirExists() throws InterruptedException, YtException {
        ytService.withoutTransaction(cypressService -> {
            if (!cypressService.exists(ytWorkDir)) {
                cypressService.create(ytWorkDir, YtNode.NodeType.MAP_NODE, true);
            }
            return true;
        });
    }

}
