package ru.yandex.webmaster3.worker.digest;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Stream;

import com.datastax.driver.core.utils.UUIDs;
import lombok.Setter;
import org.apache.commons.lang3.tuple.Pair;
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.PeriodicTaskType;
import ru.yandex.webmaster3.storage.digest.DigestTableUtil;
import ru.yandex.webmaster3.storage.util.clickhouse2.AbstractClickhouseDao;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseHost;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseServer;
import ru.yandex.webmaster3.storage.util.yt.YtException;
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.YtUtils;
import ru.yandex.webmaster3.storage.ytimport.ImportPriority;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseDataLoad;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseDataLoadType;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseImportCommand;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseTableRelation;
import ru.yandex.webmaster3.worker.TaskSchedule;
import ru.yandex.webmaster3.worker.ytimport.AbstractYtClickhouseDataLoadTask;

import static ru.yandex.webmaster3.storage.ytimport.YtClickhouseDataLoadState.DONE;

/**
 * @author avhaliullin
 */
public class PrepareDigestNotificationTask extends AbstractYtClickhouseDataLoadTask {
    private static final Logger log = LoggerFactory.getLogger(PrepareDigestNotificationTask.class);

    private static final Pattern IN_TABLE_NAME_PATTERN = Pattern.compile("^(\\d{8})_(\\d{8})$");
    private static final DateTimeFormatter IN_TABLE_NAME_DATE_FORMATTER = DateTimeFormat.forPattern("yyyyMMdd");
    private static final DateTimeFormatter CH_INSERTION_DATE_FORMATTER = DateTimeFormat.forPattern("yyyy-MM-dd");
    private static final int ROWS_COUNT = 256;
    private static final String TASK_NAME = "DIGEST_MESSAGES";

    @Setter
    private ClickhouseServer clickhouseServer;
    @Setter
    private YtPath mrBinary;
    @Setter
    private YtPath workDir;

    @Override
    protected YtClickhouseDataLoad init(YtClickhouseDataLoad latestImport) throws Exception {
        return ytService.withoutTransactionQuery(cypressService -> {
            LocalDate lastSent = latestImport.getDateTo();
            log.info("Last sent table is {}", lastSent);
            Optional<Pair<LocalDate, YtPath>> tableOpt = cypressService.list(tablePath)
                    .stream().flatMap(table -> {
                        Matcher m = IN_TABLE_NAME_PATTERN.matcher(table.getName());
                        if (m.find()) {
                            LocalDate dateTo = IN_TABLE_NAME_DATE_FORMATTER.parseLocalDate(m.group(2));
                            if (lastSent == null || dateTo.isAfter(lastSent)) {
                                return Stream.of(Pair.of(dateTo, table));
                            }
                        }
                        return Stream.empty();
                    }).max(Comparator.comparing(Pair::getLeft));
            if (tableOpt.isPresent()) {
                LocalDate date = tableOpt.get().getLeft();
                YtPath table = tableOpt.get().getRight();
                log.info("Downloading table {} {}", date, table);
                return latestImport.withSourceTable(table, date, date);
            } else {
                return latestImport.withState(DONE);
            }
        });
    }

    @Override
    protected YtClickhouseDataLoad prepare(YtClickhouseDataLoad imprt) throws Exception {

        LocalDate date = imprt.getDateTo();
        YtPath preparedTable = preparedTable(imprt);

        try {
            ytService.inTransaction(mrBinary).execute(cypressService -> {

                YtUtils.recreateTables(cypressService, Collections.singletonList(preparedTable), new YtNodeAttributes().setCompressionCodec("none"));
                YtOperationId operationId = cypressService.mapReduce(
                        YtUtils.newPrepareTablesForImportBuilder()
                                .addInputTable(imprt.getSourceTable())
                                .setOutputTables(Collections.singletonList(preparedTable))
                                .setBinary(mrBinary)
                                .setTask(TASK_NAME)
                                .setLines(ROWS_COUNT)
                                .addMapperArg("--date", date.toString(CH_INSERTION_DATE_FORMATTER))
                                .build());

                if (!cypressService.waitFor(operationId)) {
                    throw new YtException("Prepare host threats failed. See " + operationId);
                }
                return true;
            });
        } catch (YtException e) {
            throw new RuntimeException(e);
        }
        return imprt.withPreparedTables(Collections.singletonList(preparedTable))
                .withNextState();
    }

    @Override
    protected YtClickhouseDataLoad doImport(YtClickhouseDataLoad imprt) throws Exception {
        String dateString = tableDate(imprt);
        YtPath preparedTable = preparedTable(imprt);
        UUID taskId = UUIDs.timeBased();
        log.info("Host threats import " + taskId);
        List<YtClickhouseTableRelation> tablesRels = new ArrayList<>();
        int idx = 0;
        tablesRels.add(new YtClickhouseTableRelation(
                preparedTable,
                0,
                DigestTableUtil.TABLE_SPEC.replicatedMergeTreeTableName(0, dateString),
                DigestTableUtil.TABLE_SPEC.createReplicatedMergeTreeSpec(0, dateString)
        ));
        YtClickhouseImportCommand command = new YtClickhouseImportCommand(
                taskId,
                tablesRels,
                AbstractClickhouseDao.DB_WEBMASTER3_NOTIFICATIONS,
                DigestTableUtil.TABLE_SPEC.importSpec(),
                ImportPriority.ONLINE
        );
        ytClickhouseImportManager.startImport(command);
        return imprt.withImportTaskIds(taskId)
                .withNextState();
    }

    @Override
    protected YtClickhouseDataLoad replicate(YtClickhouseDataLoad imprt) throws Exception {
        String tableDate = tableDate(imprt);
        int atLeastDCs = 2;
        Map<String, ClickhouseHost> dc2Host = new HashMap<>();
        List<ClickhouseHost> hosts = new ArrayList<>(clickhouseServer.getHosts());
        Collections.shuffle(hosts);
        for (ClickhouseHost host : hosts) {
            if (host.isUp()) {
                dc2Host.putIfAbsent(host.getDcName(), host);
            }
        }
        if (dc2Host.size() < atLeastDCs) {
            throw new RuntimeException("Cannot find " + 2 + " DCs with at least one alive CH host");
        }

        String tableName = DigestTableUtil.TABLE_SPEC.replicatedMergeTreeTableName(0, tableDate);
        var command = clickhouseReplicationManager.nonShardedReplication(
                AbstractClickhouseDao.DB_WEBMASTER3_NOTIFICATIONS,
                tableName,
                DigestTableUtil.TABLE_SPEC.createReplicatedMergeTreeSpec(0, tableDate)
        );
        clickhouseReplicationManager.enqueueReplication(command);
        return imprt.withReplicationTaskIds(command.getReplicationTaskId()).withNextState();
    }

    private YtPath preparedTable(YtClickhouseDataLoad imprt) {
        return YtPath.path(workDir, tableDate(imprt));
    }

    private static String tableDate(YtClickhouseDataLoad imprt) {
        return DigestTableUtil.tableDateString(imprt.getDateTo());
    }

    @Override
    protected YtClickhouseDataLoadType getImportType() {
        return YtClickhouseDataLoadType.DIGEST_MESSAGES;
    }

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

    @Override
    public TaskSchedule getSchedule() {
        return TaskSchedule.startByCron("50 2/10 * * * *");
    }
}
