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

import com.google.common.base.Utf8;
import com.google.common.collect.Iterables;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.commons.lang3.tuple.Triple;
import org.jetbrains.annotations.Nullable;
import org.springframework.stereotype.Repository;
import ru.yandex.webmaster3.core.data.WebmasterHostId;
import ru.yandex.webmaster3.core.searchquery.Query;
import ru.yandex.webmaster3.core.searchquery.QueryId;
import ru.yandex.webmaster3.storage.util.ydb.AbstractYDao;
import ru.yandex.webmaster3.storage.util.ydb.query.Statement;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.QueryBuilder;
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;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * @author leonidrom
 */
@Repository
public class QueriesYDao extends AbstractYDao {
    private static final long MAX_YDB_BATCH_SIZE_IN_BYTES = 32_000_000;

    public QueriesYDao() {
        super(PREFIX_QUERIES, "queries");
    }

    public List<Query> get(WebmasterHostId hostId) {
        return select(QUERY_DATA_MAPPER)
                .where(F.HOST_ID.eq(hostId))
                .queryForList(Pair.of(F.QUERY_ID, q -> q.getQueryId().getQueryId()));
    }

    @Nullable
    public Query get(WebmasterHostId hostId, QueryId queryId) {
        return select(QUERY_DATA_MAPPER)
                .where(F.HOST_ID.eq(hostId))
                .and(F.QUERY_ID.eq(queryId.getQueryId()))
                .queryOne();
    }

    public Set<Query> get(WebmasterHostId hostId, Set<QueryId> queryIds) {
        return getForQueryIds(hostId, queryIds, QUERY_DATA_MAPPER);
    }

    public List<Query> getQueries(List<WebmasterHostId> hostIds) {
        var st = select(QUERY_DATA_MAPPER).where(F.HOST_ID.in(hostIds))
                .order(F.HOST_ID.asc())
                .order(F.QUERY_ID.asc());

        return queryForList(st, QUERY_DATA_MAPPER, list -> {
            var last = Iterables.getLast(list);
            WebmasterHostId lastHostId = last.getHostId();
            long lastQueryId = last.getQueryId().getQueryId();

            List<WebmasterHostId> contHostIds = hostIds.stream().filter(h -> h.compareTo(lastHostId) >= 0).toList();
            var newSt = select(QUERY_DATA_MAPPER).where(F.HOST_ID.in(contHostIds))
                    .order(F.HOST_ID.asc())
                    .order(F.QUERY_ID.asc());

            var or1 = F.HOST_ID.gt(lastHostId);

            var or2 = QueryBuilder.and(List.of(
                    F.HOST_ID.eq(lastHostId),
                    F.QUERY_ID.gt(lastQueryId)
            ));

            return newSt.cont(QueryBuilder.or(List.of(or1, or2))).getStatement();
        });

    }

    public void addBatchIfNotExist(WebmasterHostId hostId, Set<String> queries) {
        Set<QueryId> queryIds = queries.stream().map(QueryId::textToId).collect(Collectors.toSet());
        Set<QueryId> existent = getForQueryIds(hostId, queryIds, QUERY_ID_DATA_MAPPER);
        List<String> queriesToAdd = queries.stream()
                .filter(q -> !existent.contains(QueryId.textToId(q)))
                .toList();

        // Нарезаем на батчи
        List<List<String>> batches = new ArrayList<>();
        List<String> curBatch = new ArrayList<>();
        int curBatchSize = 0;
        for (String q : queriesToAdd) {
            int utf8Length = Utf8.encodedLength(q);
            if (curBatchSize + utf8Length > MAX_YDB_BATCH_SIZE_IN_BYTES
                    || curBatch.size() >= YDB_SELECT_ROWS_LIMIT) {
                batches.add(curBatch);
                curBatch = new ArrayList<>();
                curBatchSize = 0;
            }

            curBatchSize += utf8Length;
            curBatch.add(q);
        }

        if (!curBatch.isEmpty()) {
            batches.add(curBatch);
        }

        var statements = batches.stream()
                .map(batch -> batchUpdate(
                        VALUE_DATA_MAPPER,
                        batch.stream().map(q -> Triple.of(hostId, QueryId.textToId(q).getQueryId(), q)).toList()
                )).toList();

        asyncExecute(statements);
    }

    public void delete(WebmasterHostId hostId, QueryId queryId) {
        delete()
                .where(F.HOST_ID.eq(hostId))
                .and(F.QUERY_ID.eq(queryId.getQueryId()))
                .execute();
    }

    public void delete(WebmasterHostId hostId, Set<QueryId> queryIds) {
        if (queryIds.isEmpty()) {
            return;
        }

        List<Statement> statements = new ArrayList<>();
        for (List<QueryId> queryIdsBatch : Iterables.partition(queryIds, YDB_SELECT_ROWS_LIMIT)) {
            var st = delete()
                    .where(F.HOST_ID.eq(hostId))
                    .and(F.QUERY_ID.in(queryIdsBatch.stream().map(QueryId::getQueryId).toList()))
                    .getStatement();
            statements.add(st);
        }

        asyncExecute(statements);
    }

    private <T> Set<T> getForQueryIds(WebmasterHostId hostId, Set<QueryId> queryIds, DataMapper<T> mapper) {
        List<Statement> statements = new ArrayList<>();
        for (List<QueryId> queryIdsBatch : Iterables.partition(queryIds, YDB_SELECT_ROWS_LIMIT)) {
            var st = select(mapper)
                    .where(F.HOST_ID.eq(hostId))
                    .and(F.QUERY_ID.in(queryIdsBatch.stream().map(QueryId::getQueryId).toList()))
                    .getStatement();
            statements.add(st);
        }

        var res = asyncExecute(statements, mapper);
        return new HashSet<>(res);
    }

    public void addBatch(List<Triple<WebmasterHostId, Long, String>> batch) {
        batchUpdate(VALUE_DATA_MAPPER, batch).execute();
    }

    private static final ValueDataMapper<Triple<WebmasterHostId, Long, String>> VALUE_DATA_MAPPER = ValueDataMapper.create2(
        Pair.of(F.HOST_ID, Triple::getLeft),
        Pair.of(F.QUERY_ID, Triple::getMiddle),
        Pair.of(F.QUERY, Triple::getRight)
    );

    private static final DataMapper<Query> QUERY_DATA_MAPPER = DataMapper.create(
            F.HOST_ID, F.QUERY_ID, F.QUERY, (h, id, q) -> new Query(h, new QueryId(id), q));

    public static final DataMapper<QueryId> QUERY_ID_DATA_MAPPER = DataMapper.create(F.QUERY_ID, QueryId::new);

    private static final class F {
        static final Field<WebmasterHostId> HOST_ID = Fields.hostIdField("host_id");
        static final Field<Long> QUERY_ID = Fields.longField("query_id");
        static final Field<String> QUERY = Fields.stringField("query");
    }
}
