package ru.yandex.webmaster3.storage.abt.dao;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import com.google.common.collect.Iterables;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Pair;
import org.joda.time.DateTime;
import org.springframework.stereotype.Repository;

import ru.yandex.webmaster3.core.data.WebmasterHostId;
import ru.yandex.webmaster3.storage.abt.model.ExperimentInfo;
import ru.yandex.webmaster3.storage.util.ydb.AbstractYDao;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.Delete;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.Select;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.typesafe.DataMapper;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.typesafe.Field;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.typesafe.Fields;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.typesafe.ValueDataMapper;


/**
 * @author ktravchenko99 13.01.2021
 * expiration_time для этой таблицы 14 дней по полю update_date
 */
@Slf4j
@Repository("abtHostExperimentYDao")
public class AbtHostExperimentYDao extends AbstractYDao {

    private static final String TABLE_NAME = "abt_host_experiment";

    protected AbtHostExperimentYDao() {
        super(PREFIX_INTERNAL, TABLE_NAME);
    }

    public List<ExperimentInfo> getHostExperiments(WebmasterHostId hostId) {
        return select(MAPPER)
                .where(F.HOST_ID.eq(hostId))
                .queryForList();
    }

    public Map<WebmasterHostId, List<ExperimentInfo>> getHostsExperiments(Collection<WebmasterHostId> hostIds) {
        List<Select.Where<Pair<WebmasterHostId, ExperimentInfo>>> statements = new ArrayList<>();
        for (List<WebmasterHostId> hosts : Iterables.partition(hostIds, 512)) {
            statements.add(select(HOST_EXPERIMENT_INFO_MAPPER).where(F.HOST_ID.in(hosts)));
        }

        Map<WebmasterHostId, List<ExperimentInfo>> res = new HashMap<>();
        statements.parallelStream()
                .map(st -> st.queryForList(Pair.of(F.HOST_ID, Pair::getKey), Pair.of(F.EXPERIMENT_TYPE, x -> x.getValue().getExperiment())))
                .flatMap(Collection::stream)
                .sequential()
                .forEach(p -> res.computeIfAbsent(p.getLeft(), k -> new ArrayList<>()).add(p.getRight()));

        return res;
    }

    public void delete(WebmasterHostId hostId, ExperimentInfo info) {
        getDeleteStatement(hostId, info).execute();
    }

    public void delete(Map<WebmasterHostId, ExperimentInfo> batch) {
        batch.entrySet().parallelStream()
                .map(e -> getDeleteStatement(e.getKey(), e.getValue()))
                .forEach(Delete.Where::execute);
    }

    private Delete.Where getDeleteStatement(WebmasterHostId hostId, ExperimentInfo info) {
        return delete()
                .where(F.HOST_ID.eq(hostId))
                .and(F.EXPERIMENT_TYPE.eq(info.experiment))
                .and(F.GROUP_TYPE.eq(info.group));
    }

    public void insert(Map<WebmasterHostId, ExperimentInfo> batch) {
        final List<Pair<WebmasterHostId, ExperimentInfo>> items = batch.entrySet().stream()
                .map(x -> Pair.of(x.getKey(), x.getValue())).collect(Collectors.toList());

        batchInsert(VALUE_DATA_MAPPER, items).execute();
    }


    public void batchInsert(List<Pair<WebmasterHostId, ExperimentInfo>> items) {
        batchInsert(VALUE_DATA_MAPPER, items).execute();
    }

    public void insert(WebmasterHostId hostId, ExperimentInfo info) {
        upsert(
                F.HOST_ID.value(hostId),
                F.EXPERIMENT_TYPE.value(info.experiment),
                F.GROUP_TYPE.value(info.group),
                F.UPDATE_DATE.value(info.updateDate)
        ).execute();
    }


    public void forEach(Consumer<Pair<WebmasterHostId, ExperimentInfo>> consumer) {
        streamReader(HOST_EXPERIMENT_INFO_MAPPER, consumer);
    }

    private static final DataMapper<ExperimentInfo> MAPPER = DataMapper.create(
            F.EXPERIMENT_TYPE, F.GROUP_TYPE, F.UPDATE_DATE, ExperimentInfo::new
    );

    static final DataMapper<Pair<WebmasterHostId, ExperimentInfo>> HOST_EXPERIMENT_INFO_MAPPER = DataMapper.create(
            F.HOST_ID, MAPPER, Pair::of
    );

    private static final ValueDataMapper<Pair<WebmasterHostId, ExperimentInfo>> VALUE_DATA_MAPPER = ValueDataMapper.create2(
            Pair.of(F.HOST_ID, Pair::getKey),
            Pair.of(F.EXPERIMENT_TYPE, p -> p.getValue().getExperiment()),
            Pair.of(F.GROUP_TYPE, p -> p.getValue().getGroup()),
            Pair.of(F.UPDATE_DATE, p -> p.getValue().getUpdateDate())
    );

    private interface F {
        Field<WebmasterHostId> HOST_ID = Fields.hostIdField("host_id");
        Field<String> EXPERIMENT_TYPE =
                Fields.stringField("experiment_type");

        Field<String> GROUP_TYPE = Fields.stringField("group_type");
        // Необязательное поле. Используется только для автоматически импортируемых экспериментов
        // через таску ImportExperimentsPeriodicTask
        Field<DateTime> UPDATE_DATE = Fields.jodaDateTimeField("update_date").makeOptional();
    }
}
