package ru.yandex.webmaster3.worker.searchurl;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Random;
import java.util.UUID;
import java.util.stream.Collectors;

import com.google.common.collect.Range;
import lombok.Setter;
import org.apache.commons.lang3.mutable.MutableObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import ru.yandex.webmaster3.core.data.WebmasterHostId;
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.MdbClickhouseReplicationManager;
import ru.yandex.webmaster3.storage.clickhouse.system.dao.ClickhouseSystemTablesCHDao;
import ru.yandex.webmaster3.storage.clickhouse.system.data.ClickhouseSystemTableInfo;
import ru.yandex.webmaster3.storage.host.AllHostsCacheService;
import ru.yandex.webmaster3.storage.searchurl.offline.dao.SearchBaseImportTablesRepository;
import ru.yandex.webmaster3.storage.searchurl.offline.data.SearchBaseImportInfo;
import ru.yandex.webmaster3.storage.searchurl.offline.data.SearchBaseImportStageEnum;
import ru.yandex.webmaster3.storage.searchurl.offline.data.SearchBaseImportTaskType;
import ru.yandex.webmaster3.storage.searchurl.samples.dao.SearchUrlEmailSamplesCHDao;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseException;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseHost;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseServer;
import ru.yandex.webmaster3.worker.PeriodicTask;
import ru.yandex.webmaster3.worker.TaskSchedule;

/**
 * @author aherman
 */
public class CreateSearchUrlEmailSamplesTask extends PeriodicTask<CreateSearchUrlEmailSamplesTask.State> {
    private static final Logger log = LoggerFactory.getLogger(CreateSearchUrlEmailSamplesTask.class);

    private static final int SAMPLE_SIZE = 150_000;
    private static final int BATCH_SIZE = 5_000;

    @Setter
    private AllHostsCacheService allHostsCacheService;
    @Setter
    private ClickhouseServer clickhouseServer;
    @Setter
    private MdbClickhouseReplicationManager clickhouseReplicationManager;
    @Setter
    private ClickhouseSystemTablesCHDao clickhouseSystemTablesCHDao;
    @Setter
    private SearchBaseImportTablesRepository searchBaseImportTablesCDao;
    @Setter
    private SearchUrlEmailSamplesCHDao searchUrlEmailSamplesCHDao;
    @Setter
    private PeriodicTaskType type;

    @Override
    public Result run(UUID runId) throws Exception {
        setState(new State());
        Optional<SearchBaseImportInfo> eventSamplesO = searchBaseImportTablesCDao.listBases(SearchBaseImportTaskType.SEARCH_URL_EVENT_SAMPLES)
                .stream()
                .filter(b -> b.getStage() == SearchBaseImportStageEnum.READY)
                .max(Comparator.comparing(SearchBaseImportInfo::getSearchBaseDate));

        if (eventSamplesO.isEmpty()) {
            return new Result(TaskResult.SUCCESS);
        }
        SearchBaseImportInfo eventSamples = eventSamplesO.get();
        log.info("Latest event samples: {}", eventSamples);
        getState().latestEventSamples = eventSamples;

        Optional<SearchBaseImportInfo> emailSamplesO = searchBaseImportTablesCDao.listBases(SearchBaseImportTaskType.SEARCH_URL_EMAIL_SAMPLES)
                .stream()
                .filter(b -> Objects.equals(b.getSearchBaseDate(), eventSamples.getSearchBaseDate()))
                .max(Comparator.comparing(SearchBaseImportInfo::getSearchBaseDate));

        SearchBaseImportInfo emailSamples;

        if (emailSamplesO.isPresent()) {
            emailSamples = emailSamplesO.get();
        } else {
            String dbName = eventSamples.getDbName();
            String tableName = searchUrlEmailSamplesCHDao.createTableName(eventSamples.getSearchBaseDate());
            Optional<String> dcO = clickhouseServer.pickAliveDC();
            if (dcO.isEmpty()) {
                log.error("Unable to find full DC");
                return new Result(TaskResult.FAIL);
            }
            String dc = dcO.get();

            emailSamples = new SearchBaseImportInfo(
                    SearchBaseImportTaskType.SEARCH_URL_EMAIL_SAMPLES,
                    eventSamples.getSearchBaseDate(),
                    dbName,
                    SearchBaseImportStageEnum.DISCOVERED,
                    tableName,
                    1, 1,
                    dc,
                    Collections.emptyList(),
                    false,
                    eventSamples.getSourceTable(),
                    null,
                    null
            );
            searchBaseImportTablesCDao.insertStage(emailSamples);
            getState().emailSamples = emailSamples;
        }

        while (emailSamples.getStage().isToBeProcessed()) {
            switch (emailSamples.getStage()) {
                case DISCOVERED:
                    emailSamples = createTables(emailSamples);
                    break;
                case IMPORTING:
                    emailSamples = importData(eventSamplesO.get(), emailSamples);
                    break;
                case IMPORTED:
                    emailSamples = replicate(emailSamples);
                    break;
                case REPLICATED:
                    emailSamples = emailSamples.withStage(SearchBaseImportStageEnum.READY);
                    break;
            }
            getState().emailSamples = emailSamples;
            searchBaseImportTablesCDao.insertStage(emailSamples);
        }

        return new Result(TaskResult.SUCCESS);
    }

    private SearchBaseImportInfo createTables(SearchBaseImportInfo emailSamples) {
        log.info("Create tables: {}", emailSamples);
        ClickhouseHost clickhouseHost = selectHost(emailSamples, true).orElseThrow();

        clickhouseServer.executeWithFixedHost(clickhouseHost, () -> {
            searchUrlEmailSamplesCHDao.dropAllTables(emailSamples);
            searchUrlEmailSamplesCHDao.createTable(emailSamples);
            return null;
        });

        return emailSamples.withStage(SearchBaseImportStageEnum.IMPORTING);
    }

    private SearchBaseImportInfo importData(SearchBaseImportInfo events, SearchBaseImportInfo emails) {
        Optional<ClickhouseHost> clickhouseHostO = selectHost(emails, false);
        if (clickhouseHostO.isEmpty()) {
            log.error("Import host is down, reimport");
            return emails.withStage(SearchBaseImportStageEnum.DISCOVERED);
        }
        log.info("Import data: {}", emails);
        List<WebmasterHostId> sample = getSample(SAMPLE_SIZE);
        Collections.sort(sample);

        List<Range<WebmasterHostId>> ranges = new ArrayList<>();
        WebmasterHostId lastHostId = sample.get(0);
        int i = BATCH_SIZE;
        while (i < sample.size()) {
            WebmasterHostId hostId = sample.get(i);
            ranges.add(Range.closedOpen(lastHostId, hostId));
            lastHostId = hostId;
            i += BATCH_SIZE;
        }
        ranges.add(Range.atLeast(lastHostId));
        for (Range<WebmasterHostId> range : ranges) {
            log.info("Range to import: {}", range);
        }

        clickhouseServer.executeWithFixedHost(clickhouseHostO.get(), () -> {
            clickhouseServer.executeByServiceUser(() -> {
                for (Range<WebmasterHostId> range : ranges) {
                    log.info("Import range: {}", range);
                    searchUrlEmailSamplesCHDao.insertEvents(events, emails, range);
                }
                return null;
            });
            return null;
        });

        return emails.withStage(SearchBaseImportStageEnum.IMPORTED);
    }

    private SearchBaseImportInfo replicate(SearchBaseImportInfo emailSamples)
            throws ClickhouseException, InterruptedException
    {
        log.info("Replicate: {}", emailSamples);

        var command = clickhouseReplicationManager.nonShardedReplication(
                emailSamples.getDbName(),
                emailSamples.getReadTableName(),
                SearchUrlEmailSamplesCHDao.getCreateStatement(emailSamples, false)
        );
        clickhouseReplicationManager.replicateSynchronously(command);
        return emailSamples.withStage(SearchBaseImportStageEnum.REPLICATED);
    }

    private Optional<ClickhouseHost> selectHost(SearchBaseImportInfo emailSamples, boolean first) throws ClickhouseException {
        String dc = emailSamples.getClickhouseDc();
        List<ClickhouseHost> hosts = clickhouseServer.getHosts().stream()
                .filter(h -> Objects.equals(h.getDcName(), dc))
                .filter(ClickhouseHost::isUp)
                .collect(Collectors.toList());
        if (first) {
            return Optional.of(hosts.get(0));
        } else {
            for (ClickhouseHost host : hosts) {
                Optional<ClickhouseSystemTableInfo> tableO = clickhouseSystemTablesCHDao
                        .getTable(host, emailSamples.getDbName(), emailSamples.getReadTableName());
                if (tableO.isPresent()) {
                    return Optional.of(host);
                }
            }

            return Optional.empty();
        }
    }

    public List<WebmasterHostId> getSample(int sampleCount) {
        int sampleSize = sampleCount - 2;
        MutableObject<WebmasterHostId> minHost = new MutableObject<>();
        MutableObject<WebmasterHostId> maxHost = new MutableObject<>();
        List<WebmasterHostId> sample = new ArrayList<>();
        int batchSize = 100_000;
        int countFromBatch = sampleSize / 130;
        allHostsCacheService.processKeyBatch(batch -> {
            for (WebmasterHostId host : batch) {
                if (minHost.getValue() == null || minHost.getValue().compareTo(host) == 1) {
                    minHost.setValue(host);
                }
                if (maxHost.getValue() == null || maxHost.getValue().compareTo(host) == -1) {
                    maxHost.setValue(host);
                }
            }
            int needed = countFromBatch;
            if (sample.size() + 2 * countFromBatch > sampleSize) {
                needed = sampleSize - sample.size();
            }
            Random rand = new Random();
            for (int i = 0;i < needed;i++) {
                int index = rand.nextInt(batch.size());
                sample.add(batch.get(index));
            }
        }, batchSize);

        sample.add(minHost.getValue());
        sample.add(maxHost.getValue());
        return sample;
    }

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

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

    public static class State implements PeriodicTaskState {
        private SearchBaseImportInfo latestEventSamples;
        private SearchBaseImportInfo emailSamples;

        public SearchBaseImportInfo getLatestEventSamples() {
            return latestEventSamples;
        }

        public SearchBaseImportInfo getEmailSamples() {
            return emailSamples;
        }
    }

}
