package ru.yandex.webmaster3.worker.abt;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.collect.Range;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Pair;
import org.joda.time.DateTime;
import org.joda.time.Duration;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import ru.yandex.webmaster3.core.data.WebmasterHostId;
import ru.yandex.webmaster3.core.util.IdUtils;
import ru.yandex.webmaster3.core.util.RetryUtils;
import ru.yandex.webmaster3.core.worker.task.PeriodicTaskState;
import ru.yandex.webmaster3.core.worker.task.PeriodicTaskType;
import ru.yandex.webmaster3.storage.abt.ExperimentMapperService;
import ru.yandex.webmaster3.storage.abt.dao.AbtHostExperimentYDao;
import ru.yandex.webmaster3.storage.abt.dao.AbtUserExperimentYDao;
import ru.yandex.webmaster3.storage.abt.model.ExperimentInfo;
import ru.yandex.webmaster3.storage.abt.model.IExperiment;
import ru.yandex.webmaster3.storage.host.CommonDataType;
import ru.yandex.webmaster3.storage.settings.SettingsService;
import ru.yandex.webmaster3.storage.util.yt.AsyncTableReader;
import ru.yandex.webmaster3.storage.util.yt.YtCypressService;
import ru.yandex.webmaster3.storage.util.yt.YtException;
import ru.yandex.webmaster3.storage.util.yt.YtPath;
import ru.yandex.webmaster3.storage.util.yt.YtService;
import ru.yandex.webmaster3.storage.util.yt.YtTableReadDriver;
import ru.yandex.webmaster3.worker.PeriodicTask;
import ru.yandex.webmaster3.worker.TaskSchedule;

/**
 * @author leonidrom
 * <p>
 * Импортирует данные про эксперименты из таблиц в //home/webmaster/prod|test/experiments/autoimport/
 * на Хане и Арнольде.
 * Все таблицы вычитываются за один раз, поэтому для больших таблиц эта таска не подходит.
 */
@Component("importExperimentsPeriodicTask")
@Slf4j
public class ImportExperimentsPeriodicTask extends PeriodicTask<ImportExperimentsPeriodicTask.TaskState> {
    private static final Duration TTL = Duration.standardDays(14);
    private static final RetryUtils.RetryPolicy RETRY_POLICY = RetryUtils.instantRetry(7);


    private final SettingsService settingsService;
    private final YtService ytService;
    private final AbtUserExperimentYDao abtUserExperimentYDao;
    private final AbtHostExperimentYDao abtHostExperimentYDao;
    private final ExperimentMapperService experimentMapperService;

    private final YtPath arnoldImportPath;
    private final YtPath hahnImportPath;

    @Autowired
    public ImportExperimentsPeriodicTask(
            SettingsService settingsService,
            YtService ytService,
            AbtUserExperimentYDao abtUserExperimentYDao,
            AbtHostExperimentYDao abtHostExperimentYDao,
            ExperimentMapperService experimentMapperService,
            @Value("${external.yt.service.arnold.root.default}/experiments/autoimport") String arnoldImportPath,
            @Value("${external.yt.service.hahn.root.default}/experiments/autoimport") String hahnImportPath) {
        this.settingsService = settingsService;
        this.ytService = ytService;
        this.abtUserExperimentYDao = abtUserExperimentYDao;
        this.abtHostExperimentYDao = abtHostExperimentYDao;
        this.arnoldImportPath = YtPath.fromString(arnoldImportPath);
        this.hahnImportPath = YtPath.fromString(hahnImportPath);
        this.experimentMapperService = experimentMapperService;
    }

    @Override
    public Result run(UUID runId) throws Exception {
        RetryUtils.execute(RETRY_POLICY, this::process);
        return Result.SUCCESS;
    }

    private void process() {
        setState(new TaskState());
        log.info("Started importing experiments data");

        List<Pair<YtPath, IExperiment>> experimentsToImport = getExperimentsToImport(hahnImportPath);
        experimentsToImport.addAll(getExperimentsToImport(arnoldImportPath));
        log.info("Experiments to import: {}", experimentsToImport);

        DateTime importDate = DateTime.now();
        for (Pair<YtPath, IExperiment> p : experimentsToImport) {
            YtPath tablePath = p.getLeft();
            IExperiment experiment = p.getRight();
            ytService.inTransaction(tablePath).execute(ytCypressService -> {
                switch (experiment.getScope()) {
                    case DOMAIN:
                    case HOST:
                        return processTable(ytCypressService, tablePath, experiment, importDate,
                                YtHostExperimentRow.class, this::processYtHostExperimentBatchRows);
                    case USER:
                        return processTable(ytCypressService, tablePath, experiment, importDate,
                                YtUserExperimentRow.class, this::processYtUserExperimentBatchRows);
                    default:
                        throw new IllegalStateException("Unexpected enum value " + experiment.getScope());
                }
            });

            getState().importedExperiments.add(p);
        }

        settingsService.update(CommonDataType.EXPERIMENTS_LAST_IMPORT_DATE, importDate.toString());
        log.info("Finished importing experiments data");
    }

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

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

    private interface RowProcessor<T> {
        void process(IExperiment experiment, T row, DateTime importDate);
    }

    private List<Pair<YtPath, IExperiment>> getExperimentsToImport(YtPath basePath) {
        log.info("Looking for tables in {}", basePath);
        List<Pair<YtPath, IExperiment>> res = new ArrayList<>();
        ytService.inTransaction(basePath).execute(ytCypressService -> {
            List<YtPath> tables = ytCypressService.list(basePath);
            log.info("Found tables: {}", Arrays.toString(tables.toArray()));

            tables.forEach(tablePath -> {
                String tableName = tablePath.getName().toUpperCase();

                IExperiment experiment = experimentMapperService.getExperiment(tableName);
                if (!experiment.getName().equals(tableName)) {
                    log.info("Skipping table {}", tableName);
                } else {
                    res.add(Pair.of(tablePath, experiment));
                }
            });

            return true;
        });

        return res;
    }

    private <T> boolean processTable(
            YtCypressService ytCypressService,
            YtPath tablePath,
            IExperiment experiment,
            DateTime importDate,
            Class<T> cls,
            RowProcessor<List<T>> rowProcessor) {

        log.info("Processing {}", tablePath);
        var tableReader = new AsyncTableReader<>(
                ytCypressService, tablePath, Range.all(), YtTableReadDriver.createYSONDriver(cls)
        ).splitInParts(10_000L).withRetry(5);
        List<T> rowsForProcess = new ArrayList<>();
        try (var iter = tableReader.read()) {
            while (iter.hasNext()) {
                rowsForProcess.add(iter.next());
                if (rowsForProcess.size() > 4096) {
                    rowProcessor.process(experiment, rowsForProcess, importDate);
                    rowsForProcess.clear();
                }
            }
            if (rowsForProcess.size() > 0) {
                rowProcessor.process(experiment, rowsForProcess, importDate);
            }
            return true;
        } catch (IOException | InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new YtException(e);
        }
    }

    private void processYtUserExperimentBatchRows(IExperiment experiment, List<YtUserExperimentRow> rows, DateTime importDate) {
        final List<Pair<Long, ExperimentInfo>> items = rows.stream().map(row -> Pair.of(row.userId, new ExperimentInfo(experiment.getName(), row.group, importDate))).collect(Collectors.toList());
        abtUserExperimentYDao.batchInsert(items);
    }

    private void processYtHostExperimentBatchRows(IExperiment experiment, List<YtHostExperimentRow> rows, DateTime importDate) {
        final List<Pair<WebmasterHostId, ExperimentInfo>> items = rows.stream()
                .filter(row -> row.hostId != null)
                .map(row -> Pair.of(row.hostId, new ExperimentInfo(experiment.getName(), row.group, importDate)))
                .collect(Collectors.toList());
        abtHostExperimentYDao.batchInsert(items);
    }

    private static class YtUserExperimentRow {
        final long userId;
        final String group;

        @JsonCreator
        public YtUserExperimentRow(@JsonProperty("user_id") long userId,
                                   @JsonProperty("group") String group) {
            this.userId = userId;
            this.group = group;
        }
    }

    private static class YtHostExperimentRow {
        final WebmasterHostId hostId;
        final String group;

        @JsonCreator
        public YtHostExperimentRow(@JsonProperty("host") String host,
                                   @JsonProperty("group") String group) {
            WebmasterHostId hostId = null;
            try {
                hostId = IdUtils.urlToHostId(host);
            } catch (Exception e) {
                log.error("Bad host: {}", host);
            }

            this.hostId = hostId;
            this.group = group;
        }
    }

    @Getter
    public static class TaskState implements PeriodicTaskState {
        private List<Pair<YtPath, IExperiment>> importedExperiments = new ArrayList<>();
    }
}
