#pragma once

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

#include <mapreduce/yt/interface/client.h>
#include <maps/libs/json/include/value.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/common.h>

#include <string>
#include <memory>
#include <iostream>
#include <fstream>
#include <map>
#include <tuple>

namespace maps::mrc::eye::datasets {

enum class DatasetType {
    Images,
    Objects,
    Matches,
    Clusters,
};

DECLARE_ENUM_IO(DatasetType);

class BaseDataset {
public:
    virtual void readFromJson(std::istream& is) = 0;
    virtual void readFromJson(const std::string& path) = 0;

    virtual void writeToJson(std::ostream& os) const = 0;
    virtual void writeToJson(const std::string& path) const = 0;

    virtual void readFromYT(NYT::IClientPtr client, const std::string& path) = 0;
    virtual void writeToYT(NYT::IClientPtr client, const std::string& path) const = 0;

    virtual ~BaseDataset() = default;
};

using BaseDatasetPtr = std::unique_ptr<BaseDataset>;

BaseDatasetPtr createDataset(DatasetType type);

template <typename K, typename V>
std::tuple<K, V> fromJson(const json::Value& item);
template <typename K, typename V>
json::Value toJson(const K& key, const V& value);

template <typename K, typename V>
std::tuple<K, V> fromNode(const NYT::TNode& node);
template <typename K, typename V>
NYT::TNode toNode(const K& key, const V& value);

// Шаблонный класс датасета для signs_map
//
// K - ключ по которому хранятся данные.
//     Например: для датасета картинок ключ (K) - feature id
//               для датасета матчей ключ (K) - пара feature id
// V - значение в датасете, которое хранится по какому-то ключу
//     Например: для датасета картинок значение (V) - std::string с закодированным изображением
//               для датасета матчей значение (V) - массив матчей на двух изображениях
// HeaderWrap - специальная обертка (структура и т.п.) над строкой, которая является начальным
//              ключом в json формате датасета. Обертка должна иметь поле HEADER, содержащее
//              необходимое значение.
//              Например: в датасете картинок HEADER - "features_image"
// YTSchemaWrap - специальная оберка для схемы YT таблицы, в которой хранится датасет.
//                Обертка должна иметь поле SCHEMA.
//
template <typename K, typename V, typename HeaderWrap, typename YTSchemaWrap>
class Dataset : public BaseDataset {
public:
    Dataset() = default;
    Dataset(std::map<K, V> data)
        : data_(std::move(data))
    {}

    const std::map<K, V>& data() const { return data_; }

    auto begin() { return data_.begin(); }
    auto begin() const { return data_.cbegin(); }

    auto end() { return data_.end(); }
    auto end() const { return data_.cend(); }

    V& operator[](const K& key) { return data_[key]; }
    const V& at(const K& key) const { return data_.at(key); }

    static void uploadToYT(
        NYT::IClientPtr client,
        std::istream& is, const std::string& tablePath)
    {
        auto datasetJson = json::Value::fromStream(is);
        createYTTableIfNotExists(client, NYT::TYPath(tablePath));
        auto writer = client->CreateTableWriter<NYT::TNode>(NYT::TYPath(tablePath));

        for (const json::Value& itemJson : datasetJson[HeaderWrap::HEADER]) {
            const auto& [key, value] = fromJson<K, V>(itemJson);
            writer->AddRow(toNode(key, value));
        }

        writer->Finish();
    }

    static void uploadToYT(
        NYT::IClientPtr client,
        const std::string& jsonPath, const std::string& tablePath)
    {
        std::ifstream ifs(jsonPath);
        REQUIRE(ifs.is_open(), "Failed to open file: " << jsonPath);
        uploadToYT(client, ifs, tablePath);
    }


    static void downloadFromYT(
        NYT::IClientPtr client,
        std::ostream& os, const std::string& tablePath)
    {
        auto reader = client->CreateTableReader<NYT::TNode>(NYT::TYPath(tablePath));

        json::Builder builder(os);
        builder << [&](json::ObjectBuilder b) {
            b[HeaderWrap::HEADER] << [&](json::ArrayBuilder b) {
                for (; reader->IsValid(); reader->Next()) {
                    const auto& [key, value] = fromNode<K, V>(reader->GetRow());
                    b << toJson(key, value);
                }
            };
        };
    }

    static void downloadFromYT(
        NYT::IClientPtr client,
        const std::string& jsonPath, const std::string& tablePath)
    {
        std::ofstream ofs(jsonPath);
        REQUIRE(ofs.is_open(), "Failed to create file: " << jsonPath);
        downloadFromYT(client, ofs, tablePath);
    }


    void readFromJson(std::istream& is) {
        auto datasetJson = json::Value::fromStream(is);

        for (const json::Value& itemJson : datasetJson[HeaderWrap::HEADER]) {
            const auto& [key, value] = fromJson<K, V>(itemJson);
            data_.emplace(key, value);
        }
    }

    void readFromJson(const std::string& path) {
        std::ifstream ifs(path);
        REQUIRE(ifs.is_open(), "Failed to open file: " << path);
        readFromJson(ifs);
    }


    void writeToJson(std::ostream& os) const {
        json::Builder builder(os);
        builder << [&](json::ObjectBuilder b) {
            b[HeaderWrap::HEADER] << [&](json::ArrayBuilder b) {
                for (const auto& [key, value] : data_) {
                    b << toJson(key, value);
                }
            };
        };
    }

    void writeToJson(const std::string& path) const {
        std::ofstream ofs(path);
        REQUIRE(ofs.is_open(), "Failed to create file: " << path);
        writeToJson(ofs);
    }


    void readFromYT(NYT::IClientPtr client, const std::string& path) {
        auto reader = client->CreateTableReader<NYT::TNode>(NYT::TYPath(path));

        for (; reader->IsValid(); reader->Next()) {
            const auto& [key, value] = fromNode<K, V>(reader->GetRow());
            data_.emplace(key, value);
        }
    }

    void writeToYT(NYT::IClientPtr client, const std::string& path) const {
        createYTTableIfNotExists(client, NYT::TYPath(path));
        auto writer = client->CreateTableWriter<NYT::TNode>(NYT::TYPath(path));

        for (const auto& [key, value] : data_) {
            writer->AddRow(toNode(key, value));
        }

        writer->Finish();
    }

private:
    static void createYTTableIfNotExists(
        NYT::IClientPtr client, const NYT::TYPath& tablePath)
    {
        if (!client->Exists(tablePath)) {
            client->Create(tablePath, NYT::NT_TABLE,
                NYT::TCreateOptions()
                    .Recursive(true)
                    .Attributes(NYT::TNode()("schema", YTSchemaWrap::SCHEMA))
            );
        }
    }

    std::map<K, V> data_;
};

} // namespace maps::mrc::eye::datasets
