#include "database.h"

#include <maps/libs/common/include/exception.h>
#include <maps/libs/json/include/value.h>
#include <maps/libs/json/include/builder.h>
#include <maps/libs/log8/include/log8.h>

namespace maps::wiki::toloka_downloader {

namespace json_keys {
const std::string COMMENT = "comment";
const std::string LAT = "lat";
const std::string LON = "lon";
const std::string NMAPS_OBJECT_ID = "nmapsObjectId";
const std::string OBJECT = "object";
const std::string PHOTO_URLS = "photoUrls";
const std::string SUBMIT_TS = "submitTs";
const std::string WORKER = "worker";
} // namespace json_keys

namespace {

const std::string TABLE_TASKS = "toloka.tasks";
const std::string TABLE_PHOTOS = "toloka.photos";


template <typename Value>
std::vector<Value> loadValues(pqxx::transaction_base& txn, const std::string& query)
{
    auto rows = txn.exec(query);

    std::vector<Value> result;
    result.reserve(rows.size());
    for (const auto& row : rows) {
        result.push_back(row[0].as<Value>());
    }
    return result;
}

std::string jsonize(const Assignment::Data& data)
{
    maps::json::Builder builder;
    builder << [&](maps::json::ObjectBuilder object) {
        if (data.comment) {
            object[json_keys::COMMENT] = *data.comment;
        }
        if (data.nmapsObjectId) {
            object[json_keys::NMAPS_OBJECT_ID] = *data.nmapsObjectId;
        }
        object[json_keys::SUBMIT_TS] = chrono::formatIsoDateTime(data.submitTs);
        if (!data.photoUrls.empty()) {
            object[json_keys::PHOTO_URLS] << [&](maps::json::ArrayBuilder array) {
                for (const auto& url : data.photoUrls) {
                    array << url;
                }
            };
        }

        auto writePoint = [&object] (const auto& key, const auto& point) {
            object[key] << [&](maps::json::ObjectBuilder object) {
                object[json_keys::LON] = point.x();
                object[json_keys::LAT] = point.y();
            };
        };
        writePoint(json_keys::OBJECT, data.objectPosition);
        writePoint(json_keys::WORKER, data.workerPosition);
    };
    return builder.str();
}

} // namespace

Database::Database(const common::ExtendedXmlDoc& config, TaskType taskType)
    : socialPool_(config, "social", "grinder")
    , taskType_(taskType)
{}

std::vector<PoolId> Database::loadPoolIds()
{
    auto txn = socialPool_.pool().slaveTransaction();

    static const auto query =
        "SELECT DISTINCT pool_id FROM " + TABLE_TASKS +
        " WHERE task_type = " + txn->quote(std::string(toString(taskType_)));
    return loadValues<PoolId>(*txn, query);
}

std::vector<AssignmentId> Database::loadAssignmentIds(PoolId poolId)
{
    auto txn = socialPool_.pool().slaveTransaction();

    auto query =
        "SELECT DISTINCT assignment_id FROM " + TABLE_TASKS +
        " WHERE pool_id = " + txn->quote(poolId) +
        " AND task_type = " + txn->quote(std::string(toString(taskType_)));

    return loadValues<AssignmentId>(*txn, query);
}

bool Database::saveAssignment(
    const Assignment& assignment,
    TaskStatus status)
{
    auto poolId = assignment.poolId();
    INFO() << "Saving " << poolId << " : " << assignment.id();
    auto jsonData = jsonize(assignment.data());

    auto txn = socialPool_.pool().masterWriteableTransaction();

    auto query =
        "INSERT INTO " + TABLE_TASKS +
        " (pool_id, assignment_id, task_type, data, status) VALUES (" +
            txn->quote(poolId) + ", " +
            txn->quote(assignment.id()) + ", " +
            txn->quote(std::string(toString(taskType_))) + ", " +
            txn->quote(jsonData) + "::jsonb, " +
            txn->quote(std::string(toString(status))) + ")"
        " ON CONFLICT DO NOTHING";

    if (!txn->exec(query).affected_rows()) {
        WARN() << "Already exists " << assignment.id();
        return false;
    }
    txn->commit();
    return true;
}

bool Database::updateAssignmentData(const Assignment& assignment)
{
    auto poolId = assignment.poolId();
    INFO() << "Saving " << poolId << " : " << assignment.id();
    auto jsonData = jsonize(assignment.data());

    auto txn = socialPool_.pool().masterWriteableTransaction();

    auto query =
        "UPDATE " + TABLE_TASKS +
        " SET data = " + txn->quote(jsonData) + "::jsonb"
        " WHERE pool_id = " + txn->quote(assignment.poolId()) +
        " AND assignment_id = " + txn->quote(assignment.id());

    if (!txn->exec(query).affected_rows()) {
        WARN() << "Assignment " << assignment.id() << " not exists";
        return false;
    }
    txn->commit();
    return true;
}

bool Database::updateStatus(
    const Assignment& assignment,
    TaskStatus status)
{
    auto txn = socialPool_.pool().masterWriteableTransaction();

    auto query =
        "UPDATE " + TABLE_TASKS +
        " SET status = " + txn->quote(std::string(toString(status))) +
        " WHERE pool_id = " + txn->quote(assignment.poolId()) +
        " AND assignment_id = " + txn->quote(assignment.id());

    if (!txn->exec(query).affected_rows()) {
        WARN() << "Assignment " << assignment.id() << " not exists";
        return false;
    }
    txn->commit();
    return true;
}

std::map<AttachmentId, PhotoId> Database::loadSavedPhotoIds(
    const Assignment& assignment)
{
    auto txn = socialPool_.pool().slaveTransaction();

    auto query =
        "SELECT attachment_id, photo_id"
        " FROM " + TABLE_PHOTOS +
        " WHERE pool_id = " + txn->quote(assignment.poolId()) +
        " AND assignment_id = " + txn->quote(assignment.id());

    std::map<AttachmentId, PhotoId> result;
    for (const auto& row : txn->exec(query)) {
        result.emplace(row[0].as<std::string>(), row[1].as<std::string>());
    }
    return result;
}

bool Database::saveAttachments(
    const Assignment& assignment,
    const std::map<AttachmentId, PhotoId>& attachmentToPhotoIds)
{
    REQUIRE(!attachmentToPhotoIds.empty(),
            "Empty photo ids, " << assignment.id());

    auto txn = socialPool_.pool().masterWriteableTransaction();

    std::string values;
    for (const auto& [attachmentId, photoId] : attachmentToPhotoIds) {
        values += values.empty() ? "" : ", ";
        values += "(" +
            txn->quote(assignment.poolId()) + "," +
            txn->quote(assignment.id()) + "," +
            txn->quote(attachmentId) + "," +
            txn->quote(photoId) +
        ")";
    }

    auto query =
        "INSERT INTO " + TABLE_PHOTOS +
        " (pool_id, assignment_id, attachment_id, photo_id) VALUES " + values;
    txn->exec(query);
    txn->commit();
    return true;
}

} // namespace maps::wiki::toloka_downloader
