#include "transform.h"
#include "incomplete_manager.h"
#include "files_queue.h"
#include "id_manager.h"

#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/transformers/configuration.h>

#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/transformers/fix_incomplete.h>
#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/transformers/transformer.h>


#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/common/helpers.h>
#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/common/params.h>
#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/common/config.h>
#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/common/data_error.h>

#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/ymapsdf/schema/pqxx.h>
#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/ymapsdf/db.h>
#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/ymapsdf/record.h>
#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/ymapsdf/schema.h>

#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/tds/schema.h>
#include <maps/wikimap/mapspro/tools/ymapsdf-conversion/json2ymapsdf/lib/tds/json_helper.h>

#include <maps/libs/concurrent/include/scoped_guard.h>
#include <maps/libs/common/include/exception.h>
#include <yandex/maps/wiki/common/retry_duration.h>
#include <yandex/maps/wiki/threadutils/scheduler.h>

#include <boost/filesystem.hpp>

#include <algorithm>
#include <chrono>
#include <atomic>

namespace maps::wiki::json2ymapsdf {

namespace {
const size_t BUCKET_SIZE = 200;
const size_t DEFAULT_QUEUE_SIZE = 100;

class WorkManager {
public:
    WorkManager(
        pgpool3::Pool& pool,
        const ymapsdf::schema::Schema& ymapsdfSchema,
        const Params& params);

    void run();
    bool failed() const { return failed_; }

private:
    void progress();
    void printThreadStat();

    void filesLister();
    void reader();
    void transformer();
    void writer();

    void completeReading();
    void completeTransforming();
    void completeWriting();

    std::function<void()> safeRunner(const std::string& name, void (WorkManager::*function)());

    void fail();
    void abort();

    void transformObject(const std::string& objectId, const json::Value& jsonObject);
    void writeBucket(const ymapsdf::Bucket& bucket);
    void writeRecord(const ymapsdf::Record& record);

    void storeCompleted(ymapsdf::Records& records);
    void completeInvalid(ymapsdf::Records& records);

    pgpool3::Pool& pool_;
    const ymapsdf::schema::Schema& ymapsdfSchema_;
    const Params& params_;

    ThreadedQueue<std::string> jsonFilesQueue_;
    ThreadedQueue<json::Value> jsonQueue_;
    ThreadedQueue<ymapsdf::Bucket> bucketfuls_;

    using BucketMap = std::map<const ymapsdf::schema::Table *, ymapsdf::Bucket, PtrLess<ymapsdf::schema::Table>>;
    BucketMap buckets_;
    std::atomic<size_t> bucketsSize_;

    IncompleteManager incompleteManager_;

    std::atomic<size_t> activeReaders_;
    std::atomic<size_t> activeTransformers_;
    std::atomic<size_t> activeWriters_;

    std::atomic<bool> failed_;
    std::atomic<bool> finished_;
    std::condition_variable finishedCondition_;
    std::mutex finishedMutex_;
    std::mutex bucketsMutex_;
};

std::vector<Scheduler::TTaskId>
addTasks(
    Scheduler& scheduler,
    const Scheduler::Runner& runner,
    const Scheduler::Executor& executor,
    const std::vector<Scheduler::TTaskId>& dependencies,
    size_t threadCount)
{
    std::vector<Scheduler::TTaskId> taskIds;

    for (size_t i = 0; i < threadCount; ++i) {
        taskIds.push_back(scheduler.addTask(runner, executor, dependencies));
    }

    return taskIds;
}

void
WorkManager::run()
{
    Scheduler scheduler;
    auto executor = [&](Scheduler::Runner runner) { std::thread t(std::move(runner)); t.detach(); };

    scheduler.addTask(safeRunner("Progress", &WorkManager::progress), executor, {});
    scheduler.addTask(safeRunner("Init files list", &WorkManager::filesLister), executor, {});

    auto readerTasks = addTasks(scheduler, safeRunner("Reader", &WorkManager::reader), executor, {}, params_.threadCount);
    scheduler.addTask(safeRunner("Complete reading", &WorkManager::completeReading), executor, readerTasks);

    auto transformerTasks = addTasks(scheduler, safeRunner("Transformer", &WorkManager::transformer), executor, {}, params_.threadCount);
    scheduler.addTask(safeRunner("Complete transforming", &WorkManager::completeTransforming), executor, transformerTasks);

    auto writerTasks = addTasks(scheduler, safeRunner("Writer", &WorkManager::writer), executor, {}, params_.threadCount);
    scheduler.addTask(safeRunner("Complete writing", &WorkManager::completeWriting), executor, writerTasks);

    scheduler.executeAll();

    if (failed()) {
        PROGRESS() << "Export finished with errors";
        ERROR() << "Export finished with errors";
    } else {
        PROGRESS() << "Export finished successfully";
    }
}

WorkManager::WorkManager(
        pgpool3::Pool& pool,
        const ymapsdf::schema::Schema& ymapsdfSchema,
        const Params& params)
    : pool_(pool)
    , ymapsdfSchema_(ymapsdfSchema)
    , params_(params)
    , jsonFilesQueue_(DEFAULT_QUEUE_SIZE)
    , jsonQueue_(DEFAULT_QUEUE_SIZE)
    , bucketfuls_(DEFAULT_QUEUE_SIZE * 4)
    , bucketsSize_(0)
    , activeReaders_(0)
    , activeTransformers_(0)
    , activeWriters_(0)
    , failed_(false)
    , finished_(false)
{ }

void
validateJsonAttributes(const json::Value& json, const std::string& cfg)
{
    static std::atomic<bool> firstRun(true);

    if (!firstRun) {
        return;
    }

    static std::mutex mtx;
    std::lock_guard<std::mutex> lock(mtx);

    if (!firstRun) {
        return;
    }

    auto nextFreeObjectId = tds::nextFreeObjectId(json);
    idManager().advanceTo(nextFreeObjectId);
    idManager().roundUp();

    auto cfgRelationsMode = config::relationsMode(cfg);
    auto tdsRelationsMode = tds::relationsMode(json);
    if (!cfgRelationsMode.unknown() && cfgRelationsMode != tdsRelationsMode) {
        tds::schema::discardRelationsDirection();
        WARN() << "Relation-mode specified in json mismatch config. "
            << "Ignoring relation direction. "
            << "Beware of increased memory consumption."
            << "cfg:{" << cfgRelationsMode.info() << "}, "
            << "tds:{" << tdsRelationsMode.info() << "}";
    }

    firstRun = false;
}

void
WorkManager::reader()
{
    std::string jsonFile;
    while (jsonFilesQueue_.pop(jsonFile)) {
        json::Value json{json::null};
        {
            ++activeReaders_;
            concurrent::ScopedGuard guard([this]() { --activeReaders_; });
            json = tds::readJson(jsonFile);
            validateJsonAttributes(json, params_.transformCfg);
            if (params_.deleteJson) {
                boost::filesystem::remove(jsonFile);
            }
        }
        jsonQueue_.push(std::move(json));
    }
}

void
WorkManager::transformer()
{
    json::Value json{json::null};
    while (jsonQueue_.pop(json)) {
        ++activeTransformers_;
        concurrent::ScopedGuard guard([this]() { --activeTransformers_; });
        const auto& jsonObjects = tds::getObjects(json);
        for (const auto& objectId : jsonObjects.fields()) {
            transformObject(objectId, jsonObjects[objectId]);
        }
    }
}

void
WorkManager::transformObject(const std::string& objectId, const json::Value& object)
{
    try {
        auto records = transformers::transform(tds::Object(object, objectId));
        incompleteManager_.mergeIncomplete(records);
        completeInvalid(records);
        storeCompleted(records);
    } catch (const TdsDataError &dataErr) {
        DATA_ERROR() << "Object " << objectId << " failed: " << dataErr.what();
        fail();
    } catch (const maps::RuntimeError &rtErr) {
        DATA_ERROR() << "Object " << objectId << " failed: " << rtErr.what();
        ERROR() << "Object " << objectId << " failed: " << rtErr;
        fail();
    }
}

void
WorkManager::storeCompleted(ymapsdf::Records& records)
{
    for (auto& record : records) {
        if (!record.valid()) {
            ERROR() << "invalid: " << record.table().insertClause() << " " << record;
            return; //skip invalid
        }
        --activeTransformers_;
        std::lock_guard<std::mutex> lock(bucketsMutex_);
        ++activeTransformers_;
        ymapsdf::Bucket& bucket = buckets_[&record.table()];
        if (bucket.size() >= BUCKET_SIZE) {
            ymapsdf::Bucket bucketful;
            bucketsSize_ -= bucket.size();
            std::swap(bucketful, bucket);
            --activeTransformers_;
            bucketfuls_.push(std::move(bucketful));
            ++activeTransformers_;
        }
        bucket.push(std::move(record));
        ++bucketsSize_;
    }
}

void
WorkManager::completeInvalid(ymapsdf::Records& records)
{
    auto it = records.begin();
    while (it != records.end()) {
        if (it->valid()) {
            ++it;
        } else {
            ymapsdf::Record invalid(std::move(*it));
            it = records.erase(it);
            try {
                records.splice(it, transformers::fix(ymapsdfSchema_, std::move(invalid)));
            } catch (const TdsDataError &dataErr) {
                DATA_ERROR() << "Fix failed: " << dataErr.what();
                fail();
            }
        }
    }
}

void
WorkManager::filesLister()
{
    initFilesQueue(jsonFilesQueue_, params_);
}

void
WorkManager::writer()
{
    ymapsdf::Bucket bucket;
    while(bucketfuls_.pop(bucket)) {
        ++activeWriters_;
        concurrent::ScopedGuard guard([this]() { --activeWriters_; });
        try {
            writeBucket(bucket);
        } catch (...) {
            for (const auto& record: bucket.records()) {
                writeRecord(record);
            }
        }
    }
}

void
WorkManager::writeBucket(const ymapsdf::Bucket& bucket)
{
    auto query = printBucketInsert(bucket);

    auto work = pool_.masterWriteableTransaction();
    work->exec(query);
    work->commit();
}

void
WorkManager::writeRecord(const ymapsdf::Record& record)
{
    auto query = printRecordInsert(record);
    auto name = record.table().name() + "." + record.key();
    try {
        execCommitWithRetries(pool_, name, query);
    } catch (const common::RetryDurationExpired& ex) {
        ERROR() << "Writing to table " << record.table().name()
            << " failed on record " << record.key() << ": " << ex.what();
        throw; // exec WorkManager::abort() on catch exception
    } catch (const std::exception& ex) {
        DATA_ERROR() << "Writing to table " << record.table().name()
            << " failed on record " << record.key() << ": " << ex.what();
        fail();
    }
}

void
WorkManager::progress()
{
    std::unique_lock<std::mutex> lock(finishedMutex_);
    std::chrono::seconds delay(log8::isEnabled(log8::Level::DEBUG) ? 120 : 600);
    while(!finished_) {
        if (!bucketfuls_.finished()) {
            PROGRESS()
                << jsonQueue_.pushedItemsCount() << " files, "
                << incompleteManager_.size() << " incomplete, "
                << bucketsSize_ + bucketfuls_.pendingItemsCount() * BUCKET_SIZE << " pending, "
                << bucketfuls_.poppedItemsCount() * BUCKET_SIZE << " inserted records.";
            INFO() << "Incompleted: " << incompleteManager_.stat();
        }
        printThreadStat();

        finishedCondition_.wait_for(lock, delay);
    }
}

void
WorkManager::printThreadStat()
{
    INFO()
        << " R:" << activeReaders_
        << " T:" << activeTransformers_
        << " W:" << activeWriters_
        << " i:" << incompleteManager_.waiters()
        << " f:" << jsonFilesQueue_.pendingItemsCount() << "+" << jsonFilesQueue_.poppedItemsCount()
        << " j:" << jsonQueue_.pendingItemsCount() << "+" << jsonQueue_.poppedItemsCount()
        << " b:" << bucketfuls_.pendingItemsCount() << "+" << bucketfuls_.poppedItemsCount();
}

void
WorkManager::completeReading()
{
    jsonQueue_.finish();
}

void
WorkManager::completeTransforming()
{
    INFO() << "Fixing invalid records";

    PROGRESS() << "Fixing invalid records.";
    INFO() << "Incompleted: " << incompleteManager_.stat();

    while (true) {
        auto records = incompleteManager_.popRecords();
        if (records.empty()) {
            break;
        }
        completeInvalid(records);
        storeCompleted(records);
    }

    BucketMap buckets;
    {
        std::lock_guard<std::mutex> lock(bucketsMutex_);
        buckets_.swap(buckets);
        bucketsSize_ = 0;
    }

    for (auto& kv: buckets) {
        bucketfuls_.push(std::move(kv.second));
    }
    bucketfuls_.finish();
}

void
WorkManager::completeWriting()
{
    finished_ = true;
    finishedCondition_.notify_all();
}

void
WorkManager::fail()
{
    failed_ = true;
}

void
WorkManager::abort()
{
    if (finished_) {
        return;
    }
    ERROR() << "Aborting json2ymapsdf";
    fail();
    jsonFilesQueue_.finish();
    jsonQueue_.finish();
    bucketfuls_.finish();
}

std::function<void()>
WorkManager::safeRunner(const std::string& name, void (WorkManager::*function)())
{
    return ::maps::wiki::json2ymapsdf::safeRunner(
        name,
        std::bind(&WorkManager::abort, this),
        std::bind(function, this));
}

} // namespace

void
loadAndTransformToYmapsdf(
    pgpool3::Pool& pool,
    const ymapsdf::schema::Schema& ymapsdfSchema,
    const Params& params)
{
    WorkManager t2y(pool, ymapsdfSchema, params);
    t2y.run();
    DATA_REQUIRE(!t2y.failed(), "Transformation failed");
}

} // namespace maps::wiki::json2ymapsdf
