package ru.yandex.webmaster3.storage.util.ydb;

import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

import com.yandex.ydb.ValueProtos;
import com.yandex.ydb.table.query.DataQueryResult;
import com.yandex.ydb.table.result.ResultSetReader;
import com.yandex.ydb.table.settings.ReadTableSettings;
import com.yandex.ydb.table.transaction.TransactionMode;
import com.yandex.ydb.table.values.PrimitiveValue;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Pair;
import org.jetbrains.annotations.Nullable;
import org.springframework.beans.factory.annotation.Autowired;

import ru.yandex.webmaster3.core.http.RequestTrace;
import ru.yandex.webmaster3.core.http.RequestTracer;
import ru.yandex.webmaster3.core.tracer.TableNameProvider;
import ru.yandex.webmaster3.core.tracer.YdbTracer;
import ru.yandex.webmaster3.storage.util.ydb.exception.WebmasterYdbException;
import ru.yandex.webmaster3.storage.util.ydb.query.Assignment;
import ru.yandex.webmaster3.storage.util.ydb.query.DbFieldInsertAssignment;
import ru.yandex.webmaster3.storage.util.ydb.query.PreparedStatement;
import ru.yandex.webmaster3.storage.util.ydb.query.Statement;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.BatchDelete;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.BatchInsert;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.BatchUpdate;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.Delete;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.Insert;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.Select;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.Update;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.Upsert;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.typesafe.DataMapper;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.typesafe.RowMapper;
import ru.yandex.webmaster3.storage.util.ydb.querybuilder.typesafe.ValueDataMapper;

/**
 * ishalaru
 * 11.06.2020
 **/
@Slf4j
public abstract class AbstractYDao implements ExecuteQuery, SelectQuery, TableNameProvider {
    public static final int YDB_SELECT_ROWS_LIMIT = 1000;

    public static final String PREFIX_ADMIN = "/webmaster3/admin";
    public static final String PREFIX_CACHE = "/webmaster3/cache";
    public static final String PREFIX_CHECKLIST = "/webmaster3/checklist";
    public static final String PREFIX_DELEGATION = "/webmaster3/delegation";
    public static final String PREFIX_FEEDS = "/webmaster3/feeds";
    public static final String PREFIX_HOST = "/webmaster3/host";
    public static final String PREFIX_IMPORTANT_URLS = "/webmaster3/importanturls";
    public static final String PREFIX_IMPORTER = "/webmaster3/importer";
    public static final String PREFIX_INTERNAL = "/webmaster3/internal";
    public static final String PREFIX_METRIKA = "/webmaster3/metrika";
    public static final String PREFIX_MIRRORS = "/webmaster3/mirrors";
    public static final String PREFIX_MOBILE = "/webmaster3/mobile";
    public static final String PREFIX_MONITORING = "/webmaster3/monitoring";
    public static final String PREFIX_NOTIFICATION = "/webmaster3/notification";
    public static final String PREFIX_QUERIES = "/webmaster3/queries";
    public static final String PREFIX_REGION = "/webmaster3/region";
    public static final String PREFIX_ROBOTS_TXT = "/webmaster3/robotstxt";
    public static final String PREFIX_SEARCHURL = "/webmaster3/searchurl";
    public static final String PREFIX_SITEMAP = "/webmaster3/sitemap";
    public static final String PREFIX_SPRAV = "/webmaster3/sprav";
    public static final String PREFIX_TOOLS = "/webmaster3/tools";
    public static final String PREFIX_TURBO = "/webmaster3/turbo";
    public static final String PREFIX_USER = "/webmaster3/user";
    public static final String PREFIX_VERIFICATION = "/webmaster3/verification";
    public static final String PREFIX_VIDEO = "/webmaster3/video";
    public static final String PREFIX_CRAWL = "/webmaster3/crawl";
    public static final String PREFIX_LINKS = "/webmaster3/links";
    public static final String PREFIX_WEBMASTER3 = "/webmaster3/default";
    public static final String PREFIX_URLCHECK = "/webmaster3/urlcheck";
    public static final String PREFIX_URLCHECK2 = "/webmaster3/urlcheck2";
    public static final String PREFIX_NCA = "/webmaster3/nca";

    protected static final String USER_ID_INDEX = "user_id_index";

    private static final String TRACE_QUERY = "Query: %s; Parameters: %s;";
    protected String tablePrefix;
    private final String tableSpace;
    @Getter
    protected final String tableName;
    @Autowired
    protected ThreadLocalYdbTransactionManager transactionManager;

    @org.springframework.beans.factory.annotation.Value("${webmaster3.storage.ydb.dbname}")
    private String ydbDatabase;


    protected AbstractYDao(String tableSpace, String tableName) {
        this.tableSpace = tableSpace;
        this.tableName = tableName;
    }

    public void init() {
        this.tablePrefix = ydbDatabase + tableSpace;
    }

    public String getTablePath() {
        return tablePrefix + "/" + tableName;
    }

    protected Update update() {
        return new Update(tablePrefix, tableName, this);
    }

    protected Update.With update(Assignment... args) {
        final Update update = new Update(tablePrefix, tableName, this);
        final Update.With with = update.with(args[0]);
        for (int i = 1; i < args.length; i++) {
            with.and(args[i]);
        }
        return with;
    }

    /**
     * Позволяет вставлять только новые записи. Если запись с таким первичным ключом
     * уже есть в базе, то будет ошибка.
     */
    protected Insert insert(DbFieldInsertAssignment... args) {
        return new Insert(tablePrefix, tableName, Arrays.asList(args), this);
    }

    /**
     * update/insert. В отличии от insert, не кидает ошибку если запись уже есть в базе,
     * а обновляет ее.
     */
    protected Upsert upsert(DbFieldInsertAssignment... args) {
        return new Upsert(tablePrefix, tableName, Arrays.asList(args), this);
    }

    protected <T> Upsert upsert(ValueDataMapper<T> valueDataMapper, T item) {
        return new Upsert(tablePrefix, tableName, valueDataMapper.toInsert(item), this);
    }

    protected <T> BatchInsert<T> batchInsert(ValueDataMapper<T> valueDataMapper, Collection<T> items) {
        return new BatchInsert<>(tablePrefix, tableName, valueDataMapper, items, this);
    }

    protected <T> BatchDelete<T> batchDelete(ValueDataMapper<T> valueDataMapper, Collection<T> items) {
        return new BatchDelete<>(tablePrefix, tableName, valueDataMapper, items, this);
    }

    /**
     * В отличии от batchInsert, не обнуляет содержимое неуказанных колонок
     */
    protected <T> BatchUpdate<T> batchUpdate(ValueDataMapper<T> valueDataMapper, Collection<T> items) {
        return new BatchUpdate<>(tablePrefix, tableName, valueDataMapper, items, this);
    }

    protected Delete delete() {
        return new Delete(tablePrefix, tableName, this);
    }

    protected Select select() {
        return new Select(tablePrefix, tableName, null, this);
    }

    protected Select<Long> countAll() {
        return new Select<Long>(tablePrefix, tableName, DataMapper.SINGLE_COLUMN_LONG_MAPPER, this).countAll();
    }

    protected <T> Select<T> select(RowMapper<T> rowMapper) {
        return new Select<T>(tablePrefix, tableName, rowMapper, this);
    }

    protected <T> Select<T> selectDistinct(RowMapper<T> rowMapper) {
        return new Select<T>(tablePrefix, tableName, rowMapper, this).distinct();
    }

    @Override
    public DataQueryResult execute(Statement statement) {
        long startTime = System.nanoTime();
        String query = statement.toQueryString();
        try {
            return executeTM(statement);
        } finally {
            trace(statement, startTime, query);
        }
    }

    private DataQueryResult executeTM(Statement statement) {
        return transactionManager.execute(PreparedStatement.fromStatement(statement));
    }

    public List<DataQueryResult> asyncExecute(List<? extends Statement> list) {
        long startTime = System.nanoTime();
        try {
            return transactionManager.asyncExecute(list.stream().map(PreparedStatement::fromStatement).collect(Collectors.toList()));
        } finally {
            trace(startTime);
        }

    }

    public <T> T executeInTx(Supplier<T> op, TransactionMode mode) {
        return transactionManager.executeInTx(op, mode);
    }

    public <T> List<T> asyncExecute(List<? extends Statement> list, RowMapper<T> rowMapper) {
        long startTime = System.nanoTime();
        try {
            final var collect = list.stream().map(PreparedStatement::fromStatement).collect(Collectors.toList());
            List<T> result = new ArrayList<>();
            final List<DataQueryResult> dataQueryResults = transactionManager.asyncExecute(collect);
            for (DataQueryResult dataQueryResult : dataQueryResults) {
                result.addAll(mapList(dataQueryResult, rowMapper).getLeft());
            }
            return result;
        } finally {
            trace(startTime);
        }
    }


    /**
     * Использовать с осторожностью: в YDB есть ограничение на максимальное количество строк в результате запроса
     */
    @Override
    public <T> List<T> queryForList(Statement statement, RowMapper<T> rowMapper) {
        return queryForListTM(statement, rowMapper).getLeft();
    }

    /**
     * В YDB есть ограничение на максимальное количество строк в результате запроса, поэтому используем костыль
     * в виде continuationSupplier. Это функция, которая принимает усеченный результат запроса в качестве аргумента
     * и возвращает Selection для получения остальных данных запроса.
     * Один из вариантов реализации continuationSupplier - добавить в исходный запрос сортировку по первичному ключу
     * и далее добавлять к запросу условие, что первичный ключ больше чем максимальный из усеченного результата.
     */
    @Override
    public <T> List<T> queryForList(Select statement, RowMapper<T> rowMapper,
                                    Function<List<T>, Select<T>> continuationSupplier) {
        List<T> res = new ArrayList<>();
        boolean isTruncated;
        do {
            Pair<List<T>, Boolean> pair = queryForListTM(statement, rowMapper);
            res.addAll(pair.getLeft());
            isTruncated = pair.getRight() && continuationSupplier != null;
            if (isTruncated) {
                statement = continuationSupplier.apply(res);
            }
        } while (isTruncated);

        return res;
    }

    protected <T> List<T> queryForList(Function<List<T>, Select<T>> continuationSupplier) {
        return continuationSupplier.apply(null).queryForList(continuationSupplier);
    }

    @Override
    @Nullable
    public <T> T queryOne(Statement statement, RowMapper<T> rowMapper) {
        return queryForList(statement, rowMapper).stream()
                .filter(Objects::nonNull) // результат select может быть null, это ломает findAny
                .findAny()
                .orElse(null);
    }

    private void trace(long startTime) {
        RequestTracer.getCurrentTrace().stopBatch(RequestTrace.Db.YDB, System.nanoTime() - startTime);
    }

    private void trace(Statement statement, long startTime, String query) {
        if (statement.getOperationType() == Statement.OperationType.SELECT) {
            RequestTracer.getCurrentTrace().stopRead(RequestTrace.Db.YDB, String
                            .format(TRACE_QUERY, query, statement.getParameters().toString()),
                    System.nanoTime() - startTime);
        } else {
            RequestTracer.getCurrentTrace().stopWrite(RequestTrace.Db.YDB, String
                            .format(TRACE_QUERY, query, statement.getParameters().toString()),
                    System.nanoTime() - startTime);
        }
    }

    @Override
    public <T> void streamReader(RowMapper<T> rowMapper, Consumer<T> consumer) {
        streamReader(rowMapper, consumer, false);
    }

    @Override
    public <T> void streamReader(RowMapper<T> rowMapper, Consumer<T> consumer, boolean ordered) {
        streamReader(rowMapper, consumer, ordered, null, null);
    }

    public <T> void streamReader(RowMapper<T> rowMapper, Consumer<T> consumer, boolean ordered, PrimitiveValue from, PrimitiveValue to) {
        long startTime = System.nanoTime();
        long startTimeMillis = System.currentTimeMillis();
        String query = "full scan: " + tablePrefix + "/" + tableName;
        try {
            var settingsBuilder = ReadTableSettings.newBuilder()
                    .orderedRead(ordered)
                    .columns(new ArrayList<>(rowMapper.getDependencyColumns()))
                    .timeout(3, TimeUnit.HOURS);
            if (from != null) {
                settingsBuilder.fromKeyInclusive(from);
            }
            if (to != null) {
                settingsBuilder.toKeyInclusive(to);
            }
            transactionManager.readTable(tablePrefix + "/" + tableName, rowMapper, settingsBuilder.build(), consumer);

        } finally {
            long queryTimeMillis = System.currentTimeMillis() - startTimeMillis;

            YdbTracer.getCurrentTrace().addForEachStats(tableName, queryTimeMillis);
            RequestTracer.getCurrentTrace().stopRead(RequestTrace.Db.YDB, String
                            .format(TRACE_QUERY, query, ""),
                    System.nanoTime() - startTime);
        }
    }

    private <T> Pair<List<T>, Boolean> queryForListTM(Statement statement, RowMapper<T> rowMapper) {
        DataQueryResult dataQueryResult = execute(statement);
        return mapList(dataQueryResult, rowMapper);
    }

    private <T> Pair<List<T>, Boolean> mapList(DataQueryResult dataQueryResult, RowMapper<T> rowMapper) {
        try {
            ResultSetReader rs = dataQueryResult.getResultSet(0);
            return Pair.of(mapRows(rs, rowMapper), rs.isTruncated());
        } catch (Exception e) {
            throw new WebmasterYdbException(e);
        }
    }

    private <T> List<T> mapRows(ResultSetReader rs, RowMapper<T> rowMapper) throws NoSuchFieldException, IllegalAccessException {
        if (rs.getRowCount() == 0) {
            return new ArrayList<>();
        }

        Field currentRow = null;
        currentRow = rs.getClass().getDeclaredField("currentRow");
        currentRow.setAccessible(true);

        List<T> result = new ArrayList<>();

        while (rs.next()) {
            ValueProtos.Value value = (ValueProtos.Value) currentRow.get(rs);
            YdbTracer.getCurrentTrace().addBytesStats(tableName, Statement.OperationType.SELECT.name(), value.getSerializedSize());
            result.add(rowMapper.get(rs));
        }
        return result;
    }
}
