#include "export_domain.h"
#include "export_ymapsdf.h"
#include "pg_helpers.h"

#include <maps/libs/common/include/file_utils.h>
#include <yandex/maps/wiki/common/retry_duration.h>
#include <yandex/maps/wiki/tasks/ymapsdf.h>
#include <yandex/maps/wiki/threadutils/executor.h>

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

#include <exception>
#include <fstream>
#include <map>
#include <mutex>
#include <unordered_set>

namespace maps::wiki::exporter {

namespace {

const std::string TABLES_WITH_OBJECT_IDS_QUERY_FILEPATH = "tables_with_object_ids.sql";
const int TABLE_NAME = 0;
const int PRIMARY_KEY_NAME = 1;

std::map<std::string, std::string> findTablesWithObjectIds(pgpool3::Pool& pool)
{
    auto query = maps::common::readFileToString(
        tasks::ymapsdf::getScriptPath(TABLES_WITH_OBJECT_IDS_QUERY_FILEPATH));

    return common::retryDuration([&] {
        auto txn = pool.masterReadOnlyTransaction();

        std::map<std::string, std::string> tables;
        for (const auto& row : txn->exec(query)) {
            tables.emplace(
                row[TABLE_NAME].as<std::string>(),
                row[PRIMARY_KEY_NAME].as<std::string>());
        }
        return tables;
    });
}

class ExportDomain: public ExportViaYmapsdf {
public:
    ExportDomain(
            const ExportConfig& exportCfg,
            const ExportFiles& exportFiles)
        : ExportViaYmapsdf(exportCfg, exportFiles)
    {}

private:
    void ymapsdf2Dump(const std::string& schema) override
    {
        auto objectIdsDumpFilepath = exportFiles_(TmpFile::OBJECT_IDS_DUMP_FILE);
        std::ofstream objectIdsFile(objectIdsDumpFilepath);
        REQUIRE(objectIdsFile, "Can't create dump file " << objectIdsDumpFilepath);

        for (const auto& id: dumpSchema(schema)) {
            objectIdsFile << id << "\n";
        }
    }

    std::unordered_set<revision::DBID> dumpSchema(const std::string& schema) const
    {
        std::unordered_set<revision::DBID> result;
        std::mutex mutex;

        auto tablesWithObjectIds = findTablesWithObjectIds(exportCfg().ymapsdfPool());
        INFO() << "Dumping table " << schema << " tables: " << tablesWithObjectIds.size();

        Executor executor;

        const auto tables = getSchemaTables(exportCfg().ymapsdfPool(), schema);
        for (const auto& table: tables) {
            executor.addTask(
                [&] {
                    std::string primaryKeyToDump = tablesWithObjectIds.count(table) > 0
                        ? tablesWithObjectIds.at(table)
                        : "";
                    auto ids = dumpTable(schema, table, primaryKeyToDump);
                    if (!ids.empty()) {
                        std::lock_guard<std::mutex> lock(mutex);
                        result.insert(ids.cbegin(), ids.cend());
                    }
                });
        }

        const auto maxTxnNumber = exportCfg().ymapsdfPool().state().constants.masterMaxSize;
        ThreadPool threadPool(maxTxnNumber);
        executor.executeAllInThreads(threadPool);

        INFO() << "Dumping table " << schema << " result ids: " << result.size();
        return result;
    }

    std::vector<revision::DBID> dumpTable(
        const std::string& schema,
        const std::string& table,
        const std::string& primaryKey) const
    {
        INFO() << "Dumping table " << schema << "." << table;

        auto filePath = fs::path(exportFiles_(TmpFile::DUMP_DIR)) / table;
        std::ofstream file(filePath.string());
        REQUIRE(file, "Can't create dump file " + filePath.string());
        csv::OutputStream csvOut(file);

        auto rows = common::retryDuration([&] {
            INFO() << "Loading data from table " << schema << "." << table;
            auto txn = exportCfg().ymapsdfPool().masterReadOnlyTransaction();
            return txn->exec("SELECT * FROM " + schema + "." + table);
        });

        std::vector<revision::DBID> result;
        if (!primaryKey.empty()) {
            result.reserve(rows.size());
        }

        INFO() << "Dumping table " << schema << "." << table << " rows: " << rows.size();
        for (const auto& row: rows) {
            for (const auto& col: row) {
                // To emulate postgres COPY behavior and allow the reader
                // to distinguish between null fields and empty fields,
                // we must put empty fields with double quotes
                if (col.is_null()) {
                    csvOut << "";
                } else {
                    const char* value = col.c_str();
                    if (!value[0]) {
                        csvOut.forceNextQuote();
                    }
                    csvOut << value;
                }
            }
            csvOut << csv::endl;

            if (!primaryKey.empty()) {
                result.push_back(row[primaryKey].as<revision::DBID>());
            }
        }

        INFO() << "Dumping table " << schema << "." << table << " ids: " << result.size();
        return result;
    }
};

} // namespace


Datasets exportDomainSubset(
    const ExportConfig& exportCfg,
    const ExportFiles& exportFiles)
{
    return ExportDomain(exportCfg, exportFiles).run(ResultFile::DUMP_GZ_TAR);
}


} // namespace maps::wiki::exporter
