package ru.yandex.webmaster3.worker.indexing;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.TreeMap;
import java.util.UUID;
import java.util.stream.Collectors;

import com.datastax.driver.core.utils.UUIDs;
import com.fasterxml.jackson.databind.node.BooleanNode;
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.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.ClickhouseReplicationState;
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.IndexingSamplesTablesRepository;
import ru.yandex.webmaster3.storage.indexing2.internal.data.IndexingSamplesImportType;
import ru.yandex.webmaster3.storage.indexing2.internal.data.IndexingSamplesTableInfo;
import ru.yandex.webmaster3.storage.indexing2.internal.data.IndexingSamplesTableState;
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.YtException;
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.YtOperationId;
import ru.yandex.webmaster3.storage.util.yt.YtPath;
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 AbstractImportIndexingSamplesTask extends PeriodicTask<PeriodicTaskState> {
    private static final Logger log = LoggerFactory.getLogger(AbstractImportIndexingSamplesTask.class);

    private static final LocalDate TAIL_TABLE_DATE = LocalDate.parse("1970-01-01");

    private static final long REDUCER_MEMORY_LIMIT = 1L * 1024 * 1024 * 1024; // 1Gb
    private static final int COMPACTION_THRESHOLD = 3;

    private static final String YT_TAIL_TABLE_NAME = "tail";
    private static final DateTimeFormatter TABLE_NAME_DATE_FORMATTER = DateTimeFormat.forPattern("yyyyMMdd");
    private static final String TABLE_NAME_PREFIX = "week_";
    private static final String ATTR_LAST_PROCESSED = "last_processed";

    @Setter
    private ClickhouseServer clickhouseServer;
    @Setter
    private YtService ytService;
    @Setter
    private MdbYtClickhouseImportManager ytClickhouseImportManager;
    @Setter
    private MdbClickhouseReplicationManager clickhouseReplicationManager;
    @Setter
    private IndexingSamplesTablesRepository indexingSamplesTablesCDao;
    @Setter
    private ClickhouseSystemTablesCHDao clickhouseSystemTablesCHDao;
    @Setter
    private YtPath ytDataDir;
    @Setter
    private YtPath ytWorkDir;
    @Setter
    private YtPath ytPrepareExecutable;
    @Setter
    private PeriodicTaskType type;

    protected abstract IndexingSamplesImportType getImportType();

    protected abstract String getMrTaskName();

    protected abstract String getInsertSpec();

    protected abstract String getDBName();

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

    protected abstract String getTableNamePrefix();

    protected abstract int getPartsCount();

    @Override
    public Result run(UUID runId) throws Exception {
        TreeMap<LocalDate, YtTableInfo> ytTablesMap = getYtTables();
        if (ytTablesMap.size() >= COMPACTION_THRESHOLD) {
            requestCompaction(ytTablesMap.navigableKeySet().headSet(LocalDate.now().minusDays(1)));
        }
        log.info("Count tables on YT: {}", ytTablesMap.size());
        TreeMap<LocalDate, IndexingSamplesTableInfo> curTablesMap = getCurrentTableStates();

        boolean processedAbandonedTasks = false;
        for (IndexingSamplesTableInfo table : curTablesMap.values()) {
            if (!table.getState().isTerminal()) {
                processTable(table);
                processedAbandonedTasks = true;
            }
        }
        log.info("ProcessedAbandonedTasks: {}", processedAbandonedTasks);
        if (processedAbandonedTasks) {
            curTablesMap = getCurrentTableStates();
        }
        log.info("Count cur tables: {}", curTablesMap.size());
        log.debug("Current table states: {}", curTablesMap);

        for (YtTableInfo ytTable : ytTablesMap.values()) {
            boolean isTailTable = ytTable.getDate().equals(TAIL_TABLE_DATE);
            IndexingSamplesTableInfo localTable = curTablesMap.get(ytTable.getDate());
            if (localTable == null || localTable.getUpdateTime().isBefore(ytTable.getUpdateDate())) {
                Set<LocalDate> replacing;
                if (isTailTable) {
                    replacing = new HashSet<>();
                    // Tail-таблица заменяет предыдущую tail-таблицу,
                    // плюс все недельные таблицы, которые есть локально, но отсутствуют на YT
                    replacing.add(TAIL_TABLE_DATE);
                    curTablesMap.keySet()
                            .stream()
                            .filter(tableDate -> !ytTablesMap.containsKey(tableDate))
                            .forEach(replacing::add);
                } else {
                    replacing = Collections.singleton(ytTable.getDate());
                }
                log.info("Count replacing table: {}", replacing.size());
                log.debug("Replacing table: {}", replacing);
                // Если какая-то из таблиц, которые мы собираемся заменить, еще не прогружена -
                // не будем пока запускать этот импорт.
                // Пытаемся избежать дурацкой ситуации,
                // когда мы вообще ничего прогрузить не успеваем и начинаем грузить следующую таблицу
                boolean skipImport = false;
                for (LocalDate replacingTable : replacing) {
                    if (curTablesMap.containsKey(replacingTable) && !curTablesMap.get(replacingTable).getState().isTerminal()) {
                        skipImport = true;
                        break;
                    }
                }
                log.info("Skip import: {}", skipImport);
                if (!skipImport) {
                    IndexingSamplesTableInfo newTableInfo = new IndexingSamplesTableInfo(
                            getImportType(),
                            ytTable.getDate(),
                            isTailTable,
                            ytTable.getUpdateDate(),
                            replacing,
                            ytTable.getPath(),
                            IndexingSamplesTableState.PREPARING,
                            null,
                            null,
                            null,
                            null);
                    storeTableState(newTableInfo);
                    processTable(newTableInfo);
                }
            }
        }

        return new Result(TaskResult.SUCCESS);
    }

    private void requestCompaction(Set<LocalDate> dates) throws YtException, InterruptedException {
        ytService.withoutTransaction(cypressService -> {
            List<YtPath> items = cypressService.list(YtPath.path(ytDataDir, "weekly"));
            for (YtPath table : items) {
                LocalDate tableDate = null;
                if (table.getName().startsWith(TABLE_NAME_PREFIX)) {
                    tableDate = LocalDate.parse(table.getName().substring("week_20170101_".length()), TABLE_NAME_DATE_FORMATTER);
                    if (dates.contains(tableDate)) {
                        cypressService.set(YtPath.path(table, "@to_compact"), BooleanNode.TRUE);
                    }
                }
            }
            return true;
        });

    }

    private void processTable(IndexingSamplesTableInfo table) throws YtException, InterruptedException, ClickhouseException, WebmasterYdbException {
        assureWorkDirExists();
        IndexingSamplesTableState prevState;
        do {
            StageResult stageResult;
            prevState = table.getState();
            switch (table.getState()) {
                case PREPARING:
                    stageResult = runPrepare(table)
                            .nextStageIfFinished(IndexingSamplesTableState.IMPORTING);
                    break;
                case IMPORTING:
                    stageResult = runImport(table)
                            .nextStageIfFinished(IndexingSamplesTableState.REPLICATING);
                    break;
                case REPLICATING:
                    stageResult = runReplication(table)
                            .nextStageIfFinished(IndexingSamplesTableState.REPLACING);
                    break;
                case REPLACING:
                    stageResult = runReplace(table)
                            .nextStageIfFinished(IndexingSamplesTableState.CLEANUP);
                    break;
                case CLEANUP:
                    stageResult = runCleanup(table)
                            .nextStageIfFinished(IndexingSamplesTableState.DONE);
                    break;
                default:
                    throw new RuntimeException("Unknown non-terminal state " + table.getState());
            }
            table = stageResult.tableInfo;
            storeTableState(table);
        } while (prevState != table.getState() && !table.getState().isTerminal());
    }

    private StageResult runCleanup(IndexingSamplesTableInfo table) throws YtException, InterruptedException, ClickhouseException {
        ytService.withoutTransactionQuery(cypressService -> {
            for (YtPath path : table.getPreparedYtTables()) {
                if (cypressService.exists(path)) {
                    cypressService.remove(path);
                }
            }
            return true;
        });
        for (LocalDate oldTableDate : table.getReplacingCHTables()) {
            for (String tableName : getTableNamesWithParts(oldTableDate)) {
                clickhouseServer.executeOnAllHosts(
                        ClickhouseQueryContext.useDefaults().setTimeout(Duration.standardMinutes(1)),
                        "DROP TABLE IF EXISTS " + getDBName() + "." + getOldTableName(tableName)
                );
            }
        }
        return StageResult.finished(table);
    }

    private StageResult runReplace(IndexingSamplesTableInfo table) throws ClickhouseException {
        for (ClickhouseHost host : clickhouseServer.getHosts()) {
            ClickhouseQueryContext.Builder chQCtx = ClickhouseQueryContext.useDefaults()
                    .setHost(host);

            Set<String> existTables = clickhouseSystemTablesCHDao.getTables(host, getDBName())
                    .stream()
                    .map(ClickhouseSystemTableInfo::getName)
                    .collect(Collectors.toSet());
            for (int part = 0; part < getPartsCount(); part++) {
                String tmpTableName = table.getTmpTableName() + "_" + part;
                if (existTables.contains(tmpTableName)) {
                    int finalPart = part;
                    Set<String> existingOldTables = table.getReplacingCHTables()
                            .stream()
                            .map(tbl -> getTableName(tbl, finalPart))
                            .filter(existTables::contains)
                            .collect(Collectors.toSet());
                    StringBuilder replaceQ = new StringBuilder("RENAME TABLE ");
                    String targetTableName = getTableName(table.getTableDate());
                    for (String oldTable : existingOldTables) {
                        replaceQ.append(getDBName()).append(".").append(oldTable)
                                .append(" TO ")
                                .append(getDBName()).append(".").append(getOldTableName(oldTable))
                                .append(", ");
                    }
                    replaceQ.append(getDBName() + "." + table.getTmpTableName() + "_" + part)
                            .append(" TO ")
                            .append(getDBName() + "." + targetTableName + "_" + part);
                    clickhouseServer.execute(chQCtx, ClickhouseServer.QueryType.INSERT, replaceQ.toString(), Optional.empty(), Optional.empty());
                }
            }
        }
        return StageResult.finished(table);
    }


    private StageResult runReplication(IndexingSamplesTableInfo table)  {
        if (table.getReplicationManagerTaskId() == null) {
            UUID replicationTaskId = UUIDs.timeBased();
            table = table.withReplicationManagerTaskId(replicationTaskId);
            storeTableState(table);
            List<ClickhouseReplicationCommand.TableInfo> tables = new ArrayList<>();
            for (int chShard = 0; chShard < getClickhouseShardsCount(); chShard++) {
                for (int part = 0; part < getPartsCount(); part++){
                    tables.add(new ClickhouseReplicationCommand.TableInfo(
                            table.getTmpTableName() + "_" + part,
                            getTableCreateSpec(table.getTableDate(), table.getUpdateTime(), part),
                            chShard
                    ));
                }
            }
            clickhouseReplicationManager.enqueueReplication(new ClickhouseReplicationCommand(
                    replicationTaskId,
                    getDBName(),
                    table.isTail() ? ClickhouseReplicationPriority.OFFLINE : ClickhouseReplicationPriority.ONLINE,
                    tables
            ));
        }
        ClickhouseReplicationTaskGroup taskState = clickhouseReplicationManager.getReplicationState(table.getReplicationManagerTaskId());
        if (taskState == null) {
            String message = "Replication task " + table.getReplicationManagerTaskId() + " was lost";
            log.error(message);
            storeTableState(table.failed());
            throw new RuntimeException(message);
        }
        if (taskState.getState() == ClickhouseReplicationState.DONE) {
            return StageResult.finished(table);
        } else {
            return StageResult.inProgress(table);
        }
    }

    private StageResult runImport(IndexingSamplesTableInfo table) throws YtException, InterruptedException, WebmasterYdbException {
        if (table.getImportManagerTaskId() == null) {
            UUID importTaskId = UUIDs.timeBased();
            List<YtClickhouseTableRelation> ytClickhouseTableRelations = new ArrayList<>();

            String tmpTableName = "tmp_" + getTableName(table.getTableDate()) + "_" + table.getUpdateTime().getMillis();
            table = table.withImportManagerTaskId(importTaskId).withTmpTable(tmpTableName);
            storeTableState(table);

            int i = 0;
            for (int shard = 0; shard < getClickhouseShardsCount(); shard++) {
                for (int part = 0; part < getPartsCount(); part++) {
                    ytClickhouseTableRelations.add(new YtClickhouseTableRelation(
                            table.getPreparedYtTables().get(i++),
                            shard,
                            tmpTableName + "_" + part,
                            getTableCreateSpec(table.getTableDate(), table.getUpdateTime(), part)
                    ));
                }
            }
            YtClickhouseImportCommand importCommand = new YtClickhouseImportCommand(
                    importTaskId,
                    ytClickhouseTableRelations,
                    getDBName(),
                    getInsertSpec(),
                    table.isTail() ? ImportPriority.OFFLINE : ImportPriority.ONLINE
            );
            ytClickhouseImportManager.startImport(importCommand);
        }

        YtClickhouseImportState importState = ytClickhouseImportManager.getImportState(table.getImportManagerTaskId());
        if (importState == null || importState.getState() == YtClickhouseImportStateEnum.FAILED ||
                importState.getState() == YtClickhouseImportStateEnum.CANCELLED) {
            String message = importState == null
                    ? "Import task " + table.getImportManagerTaskId() + " was lost"
                    : "Import task " + table.getImportManagerTaskId() + " was failed";
            log.error(message);
            storeTableState(table.failed());
            throw new RuntimeException(message);
        } else if (importState.getState() == YtClickhouseImportStateEnum.DONE) {
            return StageResult.finished(table);
        } else if (!importState.getState().isTerminal()) {
            return StageResult.inProgress(table);
        } else {
            throw new RuntimeException("Unexpected terminal state " + importState.getState());
        }
    }

    private StageResult runPrepare(IndexingSamplesTableInfo table) throws YtException {
        int shardsCount = getClickhouseShardsCount();
        List<YtPath> outTables = new ArrayList<>();
        String tableNamePrefix = table.getTableDate().toString(TABLE_NAME_DATE_FORMATTER) + "_" + table.getUpdateTime().getMillis() + "_";
        for (int shard = 0; shard < shardsCount; shard++) {
            for (int part = 0; part < getPartsCount(); part++) {
                outTables.add(YtPath.path(ytWorkDir, tableNamePrefix + shard + "_" + part));
            }
        }
        ytService.inTransaction(ytWorkDir).execute(cypressService -> {
            YtUtils.recreateTables(cypressService, outTables, new YtNodeAttributes().setCompressionCodec("none"));
            YtMapReduceCommand mrCommand = YtUtils.newPrepareTablesForImportBuilder()
                    .addInputTable(table.getSourceYtPath())
                    .setOutputTables(outTables)
                    .setTask(getMrTaskName())
                    .setBinary(ytPrepareExecutable)
                    .setReducerMemoryLimit(REDUCER_MEMORY_LIMIT)
                    .buildV2();
            YtOperationId operationId = cypressService.mapReduce(mrCommand);
            if (!cypressService.waitFor(operationId)) {
                throw new RuntimeException("YT task failed " + operationId);
            }
            //YtUtils.sortAndWait(cypressService, outTables);
            return true;
        });
        return StageResult.finished(table.withPreparedTables(outTables));
    }

    protected String getTableName(LocalDate tableDate) {
        LocalDate startDate = tableDate.minusDays(6);
        return getTableNamePrefix() +
                startDate.toString(TABLE_NAME_DATE_FORMATTER) + "_" + tableDate.toString(TABLE_NAME_DATE_FORMATTER);
    }

    protected String getTableName(LocalDate tableDate, int part) {
        LocalDate startDate = tableDate.minusDays(6);
        return getTableNamePrefix() +
                startDate.toString(TABLE_NAME_DATE_FORMATTER) + "_" + tableDate.toString(TABLE_NAME_DATE_FORMATTER) + "_" + part;
    }

    protected List<String> getTableNamesWithParts(LocalDate tableDate) {
        List<String> result = new ArrayList<>();
        result.add(getTableName(tableDate));
        for (int part = 0; part < getPartsCount(); part++) {
            result.add(getTableName(tableDate, part));
        }
        return result;
    }

    private String getOldTableName(String tableName) {
        return "old_" + tableName;
    }

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

    private TreeMap<LocalDate, IndexingSamplesTableInfo> getCurrentTableStates()  {
        Map<LocalDate, List<IndexingSamplesTableInfo>> tables = indexingSamplesTablesCDao.listAllTables()
                .stream()
                .filter(tableInfo -> tableInfo.getImportType() == getImportType())
                .collect(Collectors.groupingBy(IndexingSamplesTableInfo::getTableDate,
                        Collectors.toList()
                ));
        TreeMap<LocalDate, IndexingSamplesTableInfo> result = new TreeMap<>();
        for (Map.Entry<LocalDate, List<IndexingSamplesTableInfo>> entry : tables.entrySet()) {
            result.put(
                    entry.getKey(),
                    entry.getValue().stream().sorted(Comparator.comparing(IndexingSamplesTableInfo::getUpdateTime).reversed()).findFirst().orElseThrow()
            );
        }
        return result;
    }

    private TreeMap<LocalDate, YtTableInfo> getYtTables() throws YtException, InterruptedException {
        return ytService.withoutTransactionQuery(cypressService -> {
            List<YtPath> items = new ArrayList<>();
            YtPath tailTablePath = YtPath.path(ytDataDir, YT_TAIL_TABLE_NAME);
            if (cypressService.exists(tailTablePath)) {
                items.add(tailTablePath);
            }
            items.addAll(cypressService.list(YtPath.path(ytDataDir, "weekly")));
            TreeMap<LocalDate, YtTableInfo> result = new TreeMap<>();
            for (YtPath table : items) {
                LocalDate tableDate = null;
                if (table.getName().startsWith(TABLE_NAME_PREFIX)) {
                    tableDate = LocalDate.parse(table.getName().substring("week_20170101_".length()), TABLE_NAME_DATE_FORMATTER);
                } else if (table.getName().equals(YT_TAIL_TABLE_NAME)) {
                    tableDate = TAIL_TABLE_DATE;
                }
                YtNode ytNode = cypressService.getNode(table);
                Instant updateTime = new Instant(ytNode.getNodeMeta().get(ATTR_LAST_PROCESSED).asLong() * 1000L);
                result.put(tableDate, new YtTableInfo(table, tableDate, updateTime));
            }
            return result;
        });
    }

    private void storeTableState(IndexingSamplesTableInfo newTableInfo)  {
        indexingSamplesTablesCDao.insertState(newTableInfo);
    }

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

    private static class YtTableInfo {
        private final YtPath path;
        private final LocalDate date;
        private final Instant updateDate;

        public YtTableInfo(YtPath path, LocalDate date, Instant updateDate) {
            this.path = path;
            this.date = date;
            this.updateDate = updateDate;
        }

        public YtPath getPath() {
            return path;
        }

        public LocalDate getDate() {
            return date;
        }

        public Instant getUpdateDate() {
            return updateDate;
        }
    }

    private static class StageResult {
        private final boolean finished;
        private final IndexingSamplesTableInfo tableInfo;

        public StageResult(boolean finished, IndexingSamplesTableInfo tableInfo) {
            this.finished = finished;
            this.tableInfo = tableInfo;
        }

        public StageResult nextStageIfFinished(IndexingSamplesTableState state) {
            return finished
                    ? new StageResult(finished, tableInfo.withState(state))
                    : this;
        }

        static StageResult finished(IndexingSamplesTableInfo tableInfo) {
            return new StageResult(true, tableInfo);
        }

        static StageResult inProgress(IndexingSamplesTableInfo tableInfo) {
            return new StageResult(false, tableInfo);
        }
    }

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

    @Override
    public TaskSchedule getSchedule() {
        return TaskSchedule.startByCron("0 4/5 * * * *");
    }

}
