#include "tasks.h"
#include "file_writer.h"
#include "id_manager.h"
#include "magic_strings.h"

#include <maps/libs/log8/include/log8.h>

#include <iostream>
#include <fstream>

namespace maps {

void
TaskContext::onTaskError()
{
    failed = true;
    loadIdsPool.stop();
    loadRowsPool.stop();
    fileWritePool.stop();
}

void
TaskContext::logPoolStats() const
{
    INFO()
        << "I:" << loadIdsPool.activeTasksCount()
        << "/" << loadIdsPool.pendingTasksCount() << " "
        << "R:" << loadRowsPool.activeTasksCount()
        << "/" << loadRowsPool.pendingTasksCount() << " "
        << "W:" << fileWritePool.activeTasksCount()
        << "/" << fileWritePool.pendingTasksCount();
}

TaskBase::~TaskBase() {}

void
TaskBase::run()
{
    try {
        context_.logPoolStats();
        doRun();
    }
    catch (Exception& e) {
        context_.onTaskError();
        ERROR() << name_ << " maps exception " << e;
    }
    catch (std::exception& e) {
        context_.onTaskError();
        ERROR() << name_ << " exception " << e.what();
    }
    catch (...) {
        context_.onTaskError();
        ERROR() << name_ << " unknown exception";
    }
}

void
LoadIdsTask::doRun()
{
    INFO() << ">" << name_ << " " << category_.name();
    const size_t batchSize = context_.params.objectsPerDbRead;
    auto tr = context_.pgPool.slaveTransaction();
    setSearchPath(*tr, context_.params.db.schema);
    auto result(workExec(*tr, category_.loadIdsSql()));

    std::vector<DBID> ids;
    ids.reserve(batchSize);
    for (const auto& row: result) {
        ids.push_back(row[0].as<DBID>());
        if (ids.size() == batchSize) {
            runBatch(ids);
            ids.clear();
            ids.reserve(batchSize);
        }
    }
    runBatch(ids);
    INFO() << "<" << name_ << " " << category_.name() << " " << result.size();
}

void
LoadIdsTask::runBatch(std::vector<DBID>& ids) const
{
    if (ids.empty()) {
        return;
    }
    registerIds(category_.table(), ids);
    std::vector<DBID> tmpIds;
    tmpIds.swap(ids);
    context_.loadRowsPool.push(
        makeTask<LoadRowsTask>(context_, category_, std::move(tmpIds)));
}

void
LoadRowsTask::doRun()
{
    INFO() << ">" << name_ << " " << category_.name();
    auto tr = context_.pgPool.slaveTransaction();
    const std::string query =
        expandSqlTemplate(category_.loadRowsSqlTemplate(), ids_);
    setSearchPath(*tr, context_.params.db.schema);
    auto rows = workExec(*tr, query);

    JsonChunk jsonChunk(category_);
    jsonChunk.objects.reserve(rows.size());
    for (const auto& tuple: rows) {
        try{
            std::stringstream stream;
            json::Builder builder{stream};
            builder << [&](json::ObjectBuilder objectBuilder) {
                category_.tupleToJson(objectBuilder, tuple);
            };
            auto str = stream.str();
            if (str.empty() || str == "{}") {
                continue;
            }
            auto id = convertId(category_.table(), tuple.at(ID).as<DBID>());
            jsonChunk.objects.push_back(JsonObject{id, str});
        } catch (Exception& e) {
            throw e << ", oid=" << tuple.at(ID).as<std::string>()
                << ", cat=" << category_.name();
        }
    }
    context_.writer->put(context_, std::move(jsonChunk));
    INFO() << "<" << name_ << " " << category_.name() << " " << rows.size();
}

void WriteFileTask::doRun()
{
    INFO() << ">" << name_ << " " << filePath_;
    std::ofstream stream(filePath_.string());
    REQUIRE(stream.good(), "error opening '" << filePath_ << "' for writing");
    json::Builder builder(stream);
    builder << [&](json::ObjectBuilder objectBuilder) {
        objectBuilder[jkey::ATTRIBUTES] = [&](json::ObjectBuilder objectBuilder) {
            context_.jsonFileAttributes.write(objectBuilder);
        };
        objectBuilder[jkey::OBJECTS] = [&](json::ObjectBuilder objectBuilder) {
            for (const auto& object: objects_) {
                objectBuilder[object.id] = json::Verbatim(object.json);
                REQUIRE(stream.good(), "error writing to '" << filePath_ << "'");
            }
        };
    };
    INFO() << "<" << name_ << " " << filePath_ << " " << objects_.size();
}

void WriteSingleFileTask::doRun()
{
    INFO() << ">" << name_;
    std::unique_ptr<std::ostream> fileStream;
    if (!context_.params.outputFile.empty()) {
        fileStream.reset(new std::ofstream(context_.params.outputFile));
    }
    std::ostream& stream = fileStream ? *fileStream : std::cout;
    std::string filePath = fileStream ? context_.params.outputFile : "stdout";
    REQUIRE(stream.good(), "error opening '" << filePath << "' for writing");
    json::Builder builder(stream);
    builder << [&](json::ObjectBuilder objectBuilder) {
        objectBuilder[jkey::ATTRIBUTES] = [&](json::ObjectBuilder objectBuilder) {
            context_.jsonFileAttributes.write(objectBuilder);
        };
        objectBuilder[jkey::OBJECTS] = [&](json::ObjectBuilder objectBuilder) {
            bool finished = false;
            while (!finished) {
                finished = queue_.finished();
                std::list<JsonChunk> chunks;
                queue_.popAll(chunks);
                for (const auto& chunk: chunks) {
                    for (const auto& object: chunk.objects) {
                        objectBuilder[object.id] = json::Verbatim(object.json);
                        REQUIRE(stream.good(),
                            "error writing to '" << filePath << "'");
                    }
                }
            }
        };
    };
    INFO() << "<" << name_;
}

void RegisterIdsTask::doRun()
{
    auto writeTr = context_.pgPool.masterWriteableTransaction();
    setSearchPath(*writeTr, context_.params.db.schema);

    completeIdRegistration(*writeTr, context_.params.objectsPerDbRead);
    writeTr->commit();
}

} // namespace maps
