package ru.yandex.webmaster3.storage.importer.model.importing;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.UUID;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import com.datastax.driver.core.utils.UUIDs;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.core.type.TypeReference;
import com.google.common.base.Preconditions;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Getter;
import lombok.Value;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.joda.time.DateTime;
import org.joda.time.Duration;

import ru.yandex.webmaster3.core.util.RetryUtils;
import ru.yandex.webmaster3.core.util.json.JsonMapping;
import ru.yandex.webmaster3.storage.clickhouse.system.data.ClickhouseSystemTableInfo;
import ru.yandex.webmaster3.storage.importer.model.ImportColumnDefinition;
import ru.yandex.webmaster3.storage.importer.model.ImportContext;
import ru.yandex.webmaster3.storage.importer.model.ImportDefinition;
import ru.yandex.webmaster3.storage.importer.model.ImportStage;
import ru.yandex.webmaster3.storage.importer.model.ImportTask;
import ru.yandex.webmaster3.storage.importer.model.MdbClickhouseTableInfo;
import ru.yandex.webmaster3.storage.importer.model.MdbClickhouseTableState;
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.yt.YtException;
import ru.yandex.webmaster3.storage.util.yt.YtLockMode;
import ru.yandex.webmaster3.storage.util.yt.YtNode;
import ru.yandex.webmaster3.storage.util.yt.YtPath;
import ru.yandex.webmaster3.storage.util.yt.transfer.ClickhouseCopyToolSettings;
import ru.yandex.webmaster3.storage.util.yt.transfer.YtTransferManagerAddTaskCommand.ClickhouseCopyCommand;
import ru.yandex.webmaster3.storage.util.yt.transfer.YtTransferManagerAddTaskCommand.ClickhouseCopyOptions;
import ru.yandex.webmaster3.storage.util.yt.transfer.YtTransferManagerMutation;
import ru.yandex.webmaster3.storage.util.yt.transfer.YtTransferTaskInfo;
import ru.yandex.webmaster3.storage.util.yt.transfer.YtTransferTaskState;

/**
 * Created by Oleg Bazdyrev on 25/09/2020.
 */
@Slf4j
@Value
@AllArgsConstructor(onConstructor_ = @JsonCreator)
@JsonIgnoreProperties(ignoreUnknown = true)
@Builder
public class ImportWithTransferManager implements ImportImportingPolicy {

    private static final Pattern DISTRIBUTED_TABLE_ENGINE_PATTERN = Pattern.compile("Distributed\\s*\\(\\s*'([a-zA-Z0-9_-]+)'\\s*,\\s*'([a-zA-Z0-9_-]+)'\\s*,\\s*'([a-zA-Z0-9_-]+)'\\s*\\)");
    private static final Pattern MERGE_TABLE_ENGINE_PATTERN = Pattern.compile("Merge\\s*\\(\\s*'([a-zA-Z0-9_-]+)'\\s*,\\s*'([^']+)'\\s*\\)");
    private static final int MAX_ATTEMPTS = 3;
    private static final Duration LOCK_TIMEOUT = Duration.standardMinutes(5L);
    private static final String ATTR_LOCK_TRANSFER_MANAGER = "transfer-manager";

    public static final String LOCKS_NODE = "transfer-manager";

    boolean append;
    String primaryKey;
    String shardingKey;
    DateColumn dateColumn;
    ImportPriority priority;
    boolean disableTableInfoSave;

    @Override
    public ImportImportingType getType() {
        return ImportImportingType.TRANSFER_MANAGER;
    }

    public ImportPriority getPriority() {
        return Objects.requireNonNullElse(priority, ImportPriority.NORMAL);
    }

    @Override
    public ImportTask apply(ImportContext context) {
        ImportTask result = context.getTask();
        Data data = context.getTask().getData(ImportStage.IMPORTING, Data.class).orElse(Data.builder().attempt(0).mutationId(UUIDs.timeBased()).build());
        List<Data> currentImports = getImportRecords(context, data);
        if (currentImports == null || currentImports.size() >= getPriority().getMaxImports()) {
            log.info("No available space for transfer-manager import {}", context.getTask().getId());
            return result;
        }
        if (data.getState() == null || data.getState().isFinal()) { // start or restart import
            result = startImport(context, data);
        } else {
            result = waitForImport(context, data);
        }
        updateImportRecords(context, currentImports, result.getData(ImportStage.IMPORTING, Data.class).orElseThrow());
        return result;
    }

    private ImportTask startImport(ImportContext context, Data data) {
        ImportTask task = context.getTask();
        ImportDefinition definition = context.getDefinition();
        String tableName = context.createSubstitutor(null).replace(definition.getTableNamePattern());
        if (!append) {
            // try to clear tables
            clearImportedTables(context, tableName);
        }
        // keys
        String key = definition.getColumns() == null ? null : definition.getColumns().stream()
                .filter(column -> column.getKind() == ImportColumnDefinition.ImportColumnKind.KEY)
                .map(ImportColumnDefinition::getChName).collect(Collectors.joining(",", "(", ")"));
        var clickhouseCopyOptions = ClickhouseCopyOptions.builder();
        clickhouseCopyOptions.resetState(true);
        clickhouseCopyOptions.primaryKey(Objects.requireNonNullElse(primaryKey, key));
        clickhouseCopyOptions.shardingKey(Objects.requireNonNullElse(shardingKey, key));
        if (dateColumn != null) {
            clickhouseCopyOptions.dateColumn(dateColumn.getName()).date(dateColumn.getValue());
        }
        clickhouseCopyOptions.command(ClickhouseCopyCommand.CREATE);
        if (append) {
            // проверим, есть ли итоговая таблица
            if (context.getClickhouseSystemTablesCHDao().getTable(
                    context.getClickhouseServer().pickRandomAliveHost(), task.getDatabase(), tableName).isPresent()) {
                clickhouseCopyOptions.command(ClickhouseCopyCommand.APPEND);
            }
        }
        // schema
        if (!CollectionUtils.isEmpty(definition.getColumns())) {
            String chSchema = definition.getColumns().stream().map(col -> col.getChName() + " " + col.getChType().toString())
                    .collect(Collectors.joining(", ", "(", ")"));
            clickhouseCopyOptions.readSchema(chSchema);
        }
        try {
            ClickhouseCopyToolSettings copyToolSettings = ClickhouseCopyToolSettings.builder()
                    .ytTaskState(ClickhouseCopyToolSettings.YtTaskState.builder()
                            .tasksBasePath(YtPath.path(context.getWorkDir(), "transfer-manager").getPathWithoutCluster())
                            .removeOnCompletion(true)
                            .build())
                    .shardUploader(ClickhouseCopyToolSettings.ShardUploader.builder()
                            .insertBlockTimeout(300L).maxTooManyPartsPause(300L).tooManyPartsPause(30L)
                            .commitBatchSize(10L).threadsPerHost(1L)
                            .insertBlockSleepTime(1L)
                            .build())
                    .build();
            String taskId = context.getYtTransferManager().copyFromYtToClickhouse(task.getIntermediateTable(),
                    definition.getDatabase() + "." + tableName, context.getClickhouseServer().getClusterId(), clickhouseCopyOptions.build(),
                    new YtTransferManagerMutation(data.getMutationId(), data.getAttempt() > 0), copyToolSettings);
            log.info("Started import with taskId = {}", taskId);
            Data newData = data.toBuilder().taskId(taskId).attempt(data.getAttempt() + 1).state(YtTransferTaskState.PENDING)
                    .lastUpdate(DateTime.now()).build();
            return task.toBuilder().updated(DateTime.now()).data(task.putData(ImportStage.IMPORTING, newData)).distributedTableName(tableName).build();
        } catch (Exception e) {
            // reset data
            Data newData = data.toBuilder().state(YtTransferTaskState.FAILED).build();
            return task.updateError(e.getMessage()).data(task.putData(ImportStage.IMPORTING, newData)).build();
        }
    }

    private ImportTask waitForImport(ImportContext context, Data data) {
        ImportTask task = context.getTask();
        YtTransferTaskState taskState = YtTransferTaskState.PENDING;
        String errorMessage;
        try {
            YtTransferTaskInfo taskInfo = RetryUtils.query(RetryUtils.linearBackoff(5, Duration.standardMinutes(1L)),
                    () -> context.getYtTransferManager().getTask(data.getTaskId()));
            if (taskInfo == null) {
                taskState = YtTransferTaskState.FAILED;
                errorMessage = "Transfer task " + data.getTaskId() + " was lost";
            } else {
                taskState = taskInfo.getState();
                errorMessage = taskInfo.getError() == null ? null : taskInfo.getError().getMessage();
            }
        } catch (YtException e) {
            errorMessage = e.getMessage();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            return task;
        }

        Data newData = data.toBuilder().state(taskState).lastUpdate(DateTime.now()).build();
        task = task.toBuilder().updated(DateTime.now()).data(task.putData(ImportStage.IMPORTING, newData)).build();
        if (!taskState.isFinal()) {
            log.info("Task {} is still processing", data.getTaskId());
            return task;
        }
        log.info("Task {} finished with result {}", data.getTaskId(), taskState);
        if (taskState != YtTransferTaskState.COMPLETED) {
            if (data.getAttempt() < MAX_ATTEMPTS) {
                // попробуем еще
                return task;
            }
            if (!append) {
                // try to clear tables
                clearImportedTables(context, task.getDistributedTableName());
            }
            return task.updateError(errorMessage).build();
        }
        log.info("Task {} finished successfully. Collecting result tables", data.getTaskId());
        task = collectImportedTables(context.toBuilder().task(task).build());
        // save info about table
        if (!disableTableInfoSave) {
            context.getClickhouseTablesYDao().update(MdbClickhouseTableInfo.fromTask(task, context.getClickhouseServer().getShardsCount())
                    .state(MdbClickhouseTableState.IMPORTED).build());
        }
        return task.withNextStage().build();
    }

    private void clearImportedTables(ImportContext context, String tableName) {
        ImportTask task = context.getTask();
        log.info("Trying to clean tables after failed import {}", task.getId());
        context.getClickhouseServer().getHosts().forEach(host -> {
            // только по живым хостам
            if (host.isDown()) {
                return;
            }
            try {
                List<ClickhouseSystemTableInfo> tables = context.getClickhouseSystemTablesCHDao().getTablesForPrefix(host, task.getDatabase(), tableName);
                if (tables.isEmpty()) {
                    log.info("Nothing to clean");
                    return;
                }
                log.info("Removing tables {}", tables);
                for (ClickhouseSystemTableInfo t : tables) {
                    String query = "DROP TABLE IF EXISTS " + task.getDatabase() + "." + t.getName() + ";";
                    RetryUtils.execute(RetryUtils.linearBackoff(3, Duration.standardMinutes(10L)), () ->
                            context.getClickhouseServer().execute(ClickhouseQueryContext.useDefaults().setHost(host).setTimeout(Duration.standardMinutes(10L)), query)
                    );
                }
            } catch (ClickhouseException e) {
                log.error("Error when cleaning after failed import", e);
            } catch (InterruptedException ignore) {
                Thread.currentThread().interrupt();
            }
        });
    }

    private ImportTask collectImportedTables(ImportContext context) {
        ImportTask task = context.getTask();
        List<String> allTableNames = new ArrayList<>();
        // look table engine, parse it and collect all underlying tables
        String database = task.getDatabase();
        ClickhouseHost host = context.getClickhouseServer().pickRandomAliveHost();
        String distributedTableName = task.getDistributedTableName();
        allTableNames.add(distributedTableName);
        log.info("Dist table name: {}" , distributedTableName);
        ClickhouseSystemTableInfo tableInfo = context.getClickhouseSystemTablesCHDao().getTable(host,
                database, distributedTableName).orElseThrow();
        Matcher matcher = DISTRIBUTED_TABLE_ENGINE_PATTERN.matcher(tableInfo.getEngineFull());
        Preconditions.checkState(matcher.matches(), "Bad engine format for distributed table: " + tableInfo.getEngineFull());
        String localTable = matcher.group(3);
        allTableNames.add(localTable);
        // check localTable
        ClickhouseSystemTableInfo localTableInfo = context.getClickhouseSystemTablesCHDao().getTable(host, database, localTable).orElseThrow();
        matcher = MERGE_TABLE_ENGINE_PATTERN.matcher(localTableInfo.getEngineFull());
        if (matcher.matches()) {
            String tablePattern = matcher.group(2);
            // search for all tables
            List<ClickhouseSystemTableInfo> underlyingTables = context.getClickhouseSystemTablesCHDao().getTablesForPattern(host, database, tablePattern);
            for (ClickhouseSystemTableInfo underlyingTable : underlyingTables) {
                allTableNames.add(underlyingTable.getName());
            }
        }
        return task.toBuilder().distributedTableName(distributedTableName).localTableName(localTable).allTableNames(allTableNames).build();
    }

    /*
     *   Мы не должны запускать в параллель слишком много импортов, иначе мы убьем Clickhouse.
     *   Поэтому перед каждым запуском импорта мы пытаемся занять свое место (с учетом приоритетов)
     */
    private List<Data> getImportRecords(ImportContext context, Data currentImport) {
        // сперва залочим
        YtPath locksNode = YtPath.path(context.getLocksNode(), LOCKS_NODE);
        try {
            context.getLocksCypressService().lock(locksNode, YtLockMode.EXCLUSIVE);
        } catch (Exception e) {
            return null;
        }
        YtNode node = context.getLocksCypressService().getNode(locksNode);
        List<Data> importTasks = new ArrayList<>();
        if (node != null && node.getNodeMeta().has(ATTR_LOCK_TRANSFER_MANAGER)) {
            importTasks = JsonMapping.readValue(node.getNodeMeta().get(ATTR_LOCK_TRANSFER_MANAGER).traverse(), Data.LIST_REFERENCE);
        }
        // почистим старые локи, обновим даты и т.д.
        DateTime minLockDate = DateTime.now().minus(LOCK_TIMEOUT);
        return importTasks.stream()
                .filter(lock -> lock.getLastUpdate().isAfter(minLockDate))
                .filter(lock -> !lock.getTaskId().equals(currentImport.getTaskId()))
                .collect(Collectors.toList());
    }

    /*
        Мы не должны запускать в параллель слишком много импортов, иначе мы убьем Clickhouse.
        Поэтому перед каждым запуском импорта мы пытаемся занять свое место (с учетом приоритетов)
     */
    private void updateImportRecords(ImportContext context, List<Data> imports, Data currentImport) {
        List<Data> newImports = new ArrayList<>(imports);
        if (!currentImport.getState().isFinal()) {
            newImports.add(currentImport);
        }
        // сохраним
        YtPath locksNode = YtPath.path(context.getLocksNode(), LOCKS_NODE);
        context.getLocksCypressService().set(YtPath.attribute(locksNode, ATTR_LOCK_TRANSFER_MANAGER), JsonMapping.OM.valueToTree(newImports));
    }

    @Value
    @Builder
    public static final class DateColumn {
        String name;
        String value;
    }

    @Value
    @AllArgsConstructor(onConstructor_ = @JsonCreator)
    @Builder(toBuilder = true)
    private static final class Data {
        public static final TypeReference<List<Data>> LIST_REFERENCE = new TypeReference<>() {
        };
        String taskId;
        UUID mutationId;
        @Builder.Default
        int attempt = 0;
        YtTransferTaskState state;
        DateTime lastUpdate;
    }

    @Getter
    @AllArgsConstructor
    public enum ImportPriority {
        LOW(1),
        NORMAL(1),
        HIGH(2);

        private final int maxImports;
    }
}
