package ru.yandex.webmaster3.worker.clickhouse;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Pair;
import org.joda.time.Duration;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import ru.yandex.webmaster3.core.util.W3Collectors;
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.clickhouse.replication.dao.ClickhouseReplicationQueueYDao;
import ru.yandex.webmaster3.storage.clickhouse.replication.data.ClickhouseReplicationCommand;
import ru.yandex.webmaster3.storage.clickhouse.replication.data.ClickhouseReplicationTaskGroup;
import ru.yandex.webmaster3.storage.clickhouse.system.dao.ClickhouseSystemReplicasCHDao;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseException;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseHost;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseHostLocation;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseQueryContext;
import ru.yandex.webmaster3.storage.util.clickhouse2.MdbClickhouseServer;
import ru.yandex.webmaster3.worker.PeriodicTask;
import ru.yandex.webmaster3.worker.TaskSchedule;

/**
 * @author avhaliullin
 */
@Slf4j
@Service
@RequiredArgsConstructor(onConstructor_ = @Autowired)
public class MdbClickhouseReplicateTablesTask extends PeriodicTask<PeriodicTaskState> {

    private static final int MAX_REPLICATION_PER_HOST = 2;
    private static final Pattern REPLICA_ALREADY_EXISTS_PATTERN = Pattern.compile(".*Replica ([^ ]+) already exists.*\\n?");
    public static final Pattern ENGINE_PATTERN =
            Pattern.compile("ReplicatedMergeTree\\( *\\'([^\\']+)\\' *, *\\'([^\\']*)\\' *(, *date *)?(?:, *\\(([a-zA-Z_]+ *(?:, *[a-zA-Z_]+ *)*)\\) *)?(, *[0-9]+ *)?\\)");


    private final MdbClickhouseServer legacyMdbClickhouseServer;
    private final ClickhouseReplicationQueueYDao clickhouseReplicationQueueYDao;
    private final ClickhouseSystemReplicasCHDao legacyMdbClickhouseSystemReplicasCHDao;

    @Override
    public Result run(UUID runId) throws Exception {
        TaskState taskState = new TaskState();
        setState(taskState);
        // сортируем таски по приоритету и дате создания и оставим top-10
        List<ClickhouseReplicationTaskGroup> taskGroups = clickhouseReplicationQueueYDao.listUnfinishedTasks();
        taskState.tasksCount = taskGroups.size();
        taskGroups = taskGroups.stream()
                .sorted(ClickhouseReplicationTaskGroup.BY_PRIORITY_AND_DATE)
                .limit(10)
                .collect(Collectors.toList());
        // все таблицы по шардам
        Map<Integer, List<String>> shard2Tables = taskGroups.stream()
                .flatMap(taskGroup ->
                        taskGroup.getCommand().getTables().stream()
                                .map(tableInfo -> Pair.of(tableInfo.getShard(), taskGroup.getCommand().getDatabase() + "." + tableInfo.getName()))
                ).collect(W3Collectors.groupByPair());
        log.info("Tables by shard {}", shard2Tables);
        Map<String, Map<Integer, Map<String, Integer>>> insertsByTableShardDc = new HashMap<>();
        Map<Integer, Map<String, Integer>> replicatingTablesByShardDc = new HashMap<>();
        // соберем инфу
        for (ClickhouseHost clickhouseHost : legacyMdbClickhouseServer.getHosts()) {
            int shard = clickhouseHost.getShard();
            List<String> tables = shard2Tables.get(shard);
            Map<String, Integer> insertsByDbTable = legacyMdbClickhouseSystemReplicasCHDao.getInsertsInQueue(clickhouseHost, tables);
            for (var entry : insertsByDbTable.entrySet()) {
                insertsByTableShardDc.computeIfAbsent(entry.getKey(), k -> new HashMap<>())
                        .computeIfAbsent(shard, k -> new HashMap<>()).put(clickhouseHost.getDcName(), entry.getValue());
            }
            replicatingTablesByShardDc.computeIfAbsent(shard, k -> new HashMap<>())
                    .put(clickhouseHost.getDcName(), (int) insertsByDbTable.values().stream().filter(i -> i > 0).count());
        }
        Set<String> dcs = legacyMdbClickhouseServer.getHosts().stream().map(ClickhouseHostLocation::getDcName).collect(Collectors.toSet());
        log.info("Current replicating tables: {}", replicatingTablesByShardDc);

        // проходимся по таскам и обновляем статус / создаем таблички
        taskCycle:
        for (ClickhouseReplicationTaskGroup taskGroup : taskGroups) {
            ClickhouseReplicationCommand command = taskGroup.getCommand();
            boolean allTablesReplicated = true;
            for (ClickhouseReplicationCommand.TableInfo tableInfo : command.getTables()) {
                String table = command.getDatabase() + "." + tableInfo.getName();
                int shard = tableInfo.getShard();
                Map<String, Integer> insertsByDc = insertsByTableShardDc.getOrDefault(table, Collections.emptyMap())
                        .getOrDefault(shard, Collections.emptyMap());
                boolean hasBaseTable;
                if (tableInfo.getCreateSpec().contains("{shard}")) {
                    hasBaseTable = !insertsByDc.isEmpty();
                } else {
                    hasBaseTable = insertsByTableShardDc.containsKey(table);
                }
                if (!hasBaseTable) {
                    // если таблицы нет, то просто сносим данную плохую репликацию, соответствующая таска импорта сама перезапустится
                    log.error("Fatal error: table " + table + " at shard " + shard + ", DCs holding data not found");
                    clickhouseReplicationQueueYDao.remove(taskGroup);
                    continue taskCycle;
                }
                for (String dc : dcs) {
                    Integer inserts = insertsByDc.get(dc);
                    if (inserts == null) {
                        // табличка отсутствует, создаем ее
                        if (replicatingTablesByShardDc.get(shard).get(dc) < MAX_REPLICATION_PER_HOST) {
                            ClickhouseHost clickhouseHost = legacyMdbClickhouseServer.getHosts().stream()
                                    .filter(host -> host.getDcName().equals(dc) && host.getShard() == shard).findAny().orElseThrow();
                            log.info("Creating table {} on shard {} dc {}", table, shard, dc);
                            startReplication(table, tableInfo.getCreateSpec(), clickhouseHost);
                            replicatingTablesByShardDc.get(shard).compute(dc, (k, v) -> v + 1);
                            taskState.newReplications.add(dc + "-" + shard + ": " + table);
                            taskState.newReplicationsCount++;
                        } else {
                            log.info("Clickhouse host {}-{} is full", dc, shard);
                        }
                        allTablesReplicated = false;
                    } else if (inserts > 0) {
                        // репликация данной таблицы еще не завершена
                        allTablesReplicated = false;
                    }
                }
            }
            if (allTablesReplicated) {
                log.info("Replicating task {} is finished", taskGroup);
                clickhouseReplicationQueueYDao.update(taskGroup.finished());
            }
        }

        return new Result(TaskResult.SUCCESS);
    }

    private void startReplication(String table, String createSpec, ClickhouseHost clickhouseHost) throws ClickhouseException {
        ClickhouseQueryContext.Builder chQCtx = ClickhouseQueryContext.useDefaults().setHost(clickhouseHost).setTimeout(Duration.standardMinutes(3));
        String createQuery = upgradeCreateQuery(String.format("CREATE TABLE %1$s %2$s", table, createSpec));
        try {
            legacyMdbClickhouseServer.execute(chQCtx, createQuery);
        } catch (ClickhouseException e) {
            if (e.getError() != null) {
                Matcher matcher = REPLICA_ALREADY_EXISTS_PATTERN.matcher(e.getError());
                if (matcher.matches()) {
                    String zkPath = matcher.group(1);
                    String suffix = "/replicas/" + clickhouseHost.getHostURI().getHost();
                    if (zkPath.endsWith(suffix)) {
                        zkPath = zkPath.substring(0, zkPath.length() - suffix.length());
                        legacyMdbClickhouseServer.execute(chQCtx,
                                String.format("SYSTEM DROP REPLICA '%1$s' FROM ZKPATH '%2$s'", clickhouseHost.getHostURI().getHost(), zkPath));
                        // повторим попытку создания таблицы
                        legacyMdbClickhouseServer.execute(chQCtx, createQuery);
                        return;
                    } else {
                        log.error("Bad replica name {}", zkPath);
                    }
                }
            }
            throw e;
        }
    }

    public static String upgradeCreateQuery(String query) {
        Matcher matcher = ENGINE_PATTERN.matcher(query);
        if (matcher.find() && matcher.group(3) != null) {
            return matcher.replaceAll(String.format("ReplicatedMergeTree('%s', '%s') PARTITION BY toYYYYMM(%s) ORDER BY (%s)",
                    matcher.group(1), matcher.group(2), "date", matcher.group(4)));
        }
        return query;
    }

    @Override
    public PeriodicTaskType getType() {
        return PeriodicTaskType.MDB_CLICKHOUSE_REPLICATE_TABLES;
    }

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

    @Getter
    public class TaskState implements PeriodicTaskState {
        int tasksCount;
        int newReplicationsCount;
        List<String> newReplications = new ArrayList<>();
    }

}
