#include <maps/wikimap/mapspro/services/mrc/long_tasks/toloka_manager_cron_jobs/lib/include/delete_free_tasks.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/toloka_manager_cron_jobs/lib/include/common.h>

#include <maps/libs/log8/include/log8.h>
#include <maps/libs/common/include/make_batches.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/toloka/mds_file_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/toloka/task_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/toloka/toloka_task_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/toloka/toloka_task_suite_gateway.h>

#include <set>
#include <unordered_map>
#include <unordered_set>

namespace maps {
namespace mrc {
namespace toloka {

namespace {

// Typical interval during which Toloka users may inspect solved tasks
// and send feedback is usually under 1 week. Use 2 times this interval.
constexpr std::chrono::hours TASK_SUITE_EXPIRY_INTERVAL{2 * 7 * 24};
constexpr size_t BATCH_SIZE = 1000;

std::unordered_set<db::TId>
getExpiredTaskSuiteIds(pqxx::transaction_base& txn,
                       const db::toloka::TolokaTasks& tolokaTasks)
{
    db::TIds taskSuiteIds;
    taskSuiteIds.reserve(tolokaTasks.size());
    for (const auto& tolokaTask : tolokaTasks) {
        taskSuiteIds.push_back(tolokaTask.taskSuiteId());
    }
    sortUnique(taskSuiteIds);
    auto taskSuites = db::toloka::TolokaTaskSuiteGateway{txn}.loadByIds(std::move(taskSuiteIds));

    auto now = chrono::TimePoint(std::chrono::system_clock::now());
    auto end = std::partition(taskSuites.begin(), taskSuites.end(),
        [now](const db::toloka::TolokaTaskSuite& ts) {
            auto solvedAt = ts.solvedAt();
            if (solvedAt) {
                return now - *solvedAt > TASK_SUITE_EXPIRY_INTERVAL;
            }

            // Possibly task has been deliberately stopped in toloka
            return now - ts.createdAt() > TASK_SUITE_EXPIRY_INTERVAL;
        });

    std::unordered_set<db::TId> result;
    for (auto itr = taskSuites.begin(); itr != end; ++itr) {
        result.insert(itr->id());
    }

    return result;
}

void deleteFilesFromMds(mds::Mds& mdsClient,
                        const db::toloka::MdsFiles& files)
{
    for (const auto& file : files) {
        const mds::Key key(file.mdsGroupId(), file.mdsPath());
        try {
            mdsClient.del(key);
        }
        catch (const maps::Exception& e) {
            WARN() << "Failed to delete from MDS, key: " << key.groupId << "/"
                   << key.path << ", " << e;
        }
    }
}

void deleteFreeTasksBatch(pqxx::transaction_base& txn,
                          mds::Mds& mdsClient,
                          const db::TIds& taskIds)
{
    const auto tolokaTasks = db::toloka::TolokaTaskGateway{txn}.loadByTaskIds(taskIds);

    std::unordered_map<db::TId, db::TIds> taskIdToTaskSuiteIds;
    std::set<db::TId> postedTaskIds; // those which have related Toloka tasks
    for (const auto& tt : tolokaTasks) {
        taskIdToTaskSuiteIds[tt.taskId()].push_back(tt.taskSuiteId());
        postedTaskIds.insert(tt.taskId());
    }

    const auto expiredTsIds = getExpiredTaskSuiteIds(txn, tolokaTasks);
    auto isExpired = [&](db::TId id) { return expiredTsIds.count(id) > 0; };

    db::TIds taskIdsToDelete;
    for (const auto& pair : taskIdToTaskSuiteIds) {
        if (std::all_of(pair.second.begin(), pair.second.end(), isExpired)) {
            taskIdsToDelete.push_back(pair.first);
        }
    }
    // Add IDs of free tasks which have not been posted to Toloka
    std::set_difference(taskIds.begin(), taskIds.end(),
                        postedTaskIds.begin(), postedTaskIds.end(),
                        std::back_inserter(taskIdsToDelete));

    if (taskIdsToDelete.empty()) return;
    INFO() << "Collected batch of " << taskIdsToDelete.size()
           << " tasks to delete";

    deleteFilesFromMds(mdsClient, db::toloka::MdsFileGateway{txn}.loadByTaskIds(taskIdsToDelete));
    db::toloka::MdsFileGateway{txn}.removeByTaskIds(taskIdsToDelete);
    db::toloka::TolokaTaskGateway{txn}.removeByTaskIds(taskIdsToDelete);
    // TODO: check deleted > 0 ?
    auto deleted = db::toloka::TaskGateway{txn}.removeByIds(taskIdsToDelete);
    INFO() << "Delete " << deleted << " tasks";
}

} // anonymous namespace

void deleteFreeTasks(db::toloka::Platform platform, pgpool3::Pool& pool, mds::Mds& mdsClient)
{
    auto taskIds = db::toloka::TaskGateway{*pool.masterReadOnlyTransaction()}
                       .loadIdsByStatus(platform, db::toloka::TaskStatus::Free);
    std::sort(taskIds.begin(), taskIds.end());
    INFO() << "Number of free tasks loaded: " << taskIds.size();

    for (const auto& batch : maps::common::makeBatches(taskIds, BATCH_SIZE)) {
        auto txn = pool.masterWriteableTransaction();
        const db::TIds range{batch.begin(), batch.end()};
        deleteFreeTasksBatch(*txn, mdsClient, range);
        txn->commit();
    }
}

} // namespace toloka
} // namespace mrc
} // namespace maps
