package ru.yandex.webmaster3.worker.takeout;

import com.google.common.collect.Range;
import com.google.common.collect.Sets;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Pair;
import org.joda.time.Duration;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
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.takeout.YtUserDataDeleteQueueYDao;
import ru.yandex.webmaster3.storage.user.service.UserTakeoutService;
import ru.yandex.webmaster3.storage.util.yt.*;
import ru.yandex.webmaster3.storage.yql.YqlQueryBuilder;
import ru.yandex.webmaster3.storage.yql.YqlService;
import ru.yandex.webmaster3.worker.PeriodicTask;
import ru.yandex.webmaster3.worker.TaskSchedule;

import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

/**
 * @author leonidrom
 */
@Slf4j
@RequiredArgsConstructor(onConstructor_ = {@Autowired})
@Service
// Удаляет записи из Yt таблиц
public class DeleteYtUserDataPeriodicTask extends PeriodicTask<PeriodicTaskState> {
    private static final int TOTAL_THREADS = 16;
    private static final RetryUtils.RetryPolicy YDB_RETRY_POLICY = RetryUtils.linearBackoff(10, Duration.standardSeconds(30));

    private final YtService ytService;
    private final YqlService yqlService;
    private final YtUserDataDeleteQueueYDao ytUserDataDeleteQueueYDao;
    private final UserTakeoutService userTakeoutService;
    private final ExecutorService executorService = Executors.newFixedThreadPool(TOTAL_THREADS);

    @Override
    public Result run(UUID runId) throws Exception {
        // сгруппируем вся заявки в очереди по пути к таблице
        Map<YtPath, List<YtUserDataDeleteQueueYDao.Record>> toDeleteMap = new HashMap<>();
        Set<UUID> allRequestIds = new HashSet<>();
        ytUserDataDeleteQueueYDao.forEach(rec -> {
            allRequestIds.add(rec.requestId());
            if (!rec.isDone()){
                toDeleteMap.computeIfAbsent(rec.tablePath(), k -> new ArrayList<>()).add(rec);
            }
        });

        if (!toDeleteMap.isEmpty()) {
            // удалим записи о пользователях из этих таблиц
            log.info("Total tables to delete from: {}", toDeleteMap.size());
            deleteFromTables(toDeleteMap);
        } else {
            log.info("Nothing to delete");
        }

        // проверим, какие из заявок на удаление мы тем самым завершили
        // (удаление данных из Yt запускается после удаления данных из Ydb и CH)
        checkCompletedRequests(allRequestIds);

        return Result.SUCCESS;
    }

    private void checkCompletedRequests(Set<UUID> allRequestIds) throws InterruptedException {
        Set<UUID> inProgressRequestIds = new HashSet<>(ytUserDataDeleteQueueYDao.listInProgressRequestIds());
        Set<UUID> doneRequestIds = Sets.difference(allRequestIds, inProgressRequestIds);
        for (var requestId : doneRequestIds) {
            // если здесь сломается, то при следующем запуске таски попытаемся снова
            RetryUtils.execute(YDB_RETRY_POLICY, () -> {
                userTakeoutService.markRequestAsDone(requestId);
                ytUserDataDeleteQueueYDao.delete(requestId);
            });
        }
    }

    private void deleteFromTables(Map<YtPath, List<YtUserDataDeleteQueueYDao.Record>> toDeleteMap) throws Exception {
        ArrayList<Callable<Void>> callables = new ArrayList<>();
        // для каждой таблицы
        for (var e: toDeleteMap.entrySet()) {
            YtPath tablePath = e.getKey();
            List<YtUserDataDeleteQueueYDao.Record> requests = e.getValue();
            Set<Long> userIds = requests.stream().map(r -> r.userId()).collect(Collectors.toSet());
            callables.add(() -> {
                // удалим пользовательские данные из таблицы
                deleteFromTable(tablePath, userIds);

                // пометим как обработанные все заявки на удаление данных из этой таблицы
                Set<UUID> requestIds = requests.stream().map(r -> r.requestId()).collect(Collectors.toSet());

                RetryUtils.execute(YDB_RETRY_POLICY, () -> {
                    ytUserDataDeleteQueueYDao.markAsDone(requestIds, tablePath);
                });

                return null;
            });
        }

        var futures = executorService.invokeAll(callables);
        for (var f : futures) {
            f.get();
        }
    }

    private void deleteFromTable(YtPath tablePath, Collection<Long> userIds) {
        ytService.inTransaction(tablePath).withLock(tablePath, YtLockMode.EXCLUSIVE).execute(cypressService -> {
            // получим все индексы, по которым надо удалить строки из этой таблицы
            log.info("Getting indexes to delete from {}", tablePath);
            List<YtTableRange> indexesToDelete = getTableIndexesToDelete(cypressService, tablePath, userIds);
            if (indexesToDelete.isEmpty()) {
                return true;
            }

            // запустим удаление
            log.info("Starting to delete from {}", tablePath.toYtPath(indexesToDelete));
            YtOperationId operationId = cypressService.erase(tablePath, indexesToDelete);

            // дождемся результата
            log.info("Waiting for delete from {} to finish", tablePath);
            cypressService.waitFor(operationId);
            log.info("Finished deleting from {}", tablePath);

            return true;
        });
    }

    private List<YtTableRange> getTableIndexesToDelete(YtCypressService cypressService, YtPath tablePath, Collection<Long> userIds) {
        // сделаем Yql запрос для получения индексов строк с пользовательскими данными
        String query = getTableIndexesQuery(cypressService, tablePath, userIds);
        List<Long> indexesList = new ArrayList<>();
        yqlService.query(query, rs -> rs.getLong("idx"), indexesList::add);
        if (indexesList.isEmpty()) {
            return Collections.emptyList();
        }

        // разобьем полученные индексы на непрерывные диапазоны
        indexesList = indexesList.stream().sorted().distinct().toList();
        List<YtTableRange> res = new ArrayList<>();
        var curRange = Pair.of(indexesList.get(0), indexesList.get(0));
        for (int i = 1; i < indexesList.size(); i++) {
            long idx = indexesList.get(i);
            if (curRange.getRight() + 1 == idx) {
                curRange = Pair.of(curRange.getLeft(), idx);
            } else {
                res.add(YtTableRange.index(Range.closedOpen(curRange.getLeft(), curRange.getRight() + 1)));
                curRange = Pair.of(idx, idx);
            }
        }

        res.add(YtTableRange.index(Range.closedOpen(curRange.getLeft(), curRange.getRight() + 1)));

        return res;
    }

    private String getTableIndexesQuery(YtCypressService cypressService, YtPath tablePath, Collection<Long> userIds) {
        YqlQueryBuilder qb = new YqlQueryBuilder();
        qb.transaction(cypressService.getTransactionId());
        qb.cluster(tablePath);

        qb.appendText("SELECT TableRecordIndex() - 1 as idx FROM ");
        qb.appendTable(tablePath);
        String userIdsStr = userIds.stream().map(String::valueOf).collect(Collectors.joining(","));
        qb.appendText(" WHERE user_id in (" + userIdsStr + ")");

        return qb.build();
    }

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

    @Override
    public TaskSchedule getSchedule() {
        return TaskSchedule.never();
    }
}
