#pragma once

#include <mapreduce/yt/interface/operation.h>

#include <saas/library/rtyt/lib/operation/context.h>

#include <saas/library/rtyt/lib/io/node_table_reader.h>
#include <saas/library/rtyt/lib/io/node_table_writer.h>
#include <saas/library/rtyt/lib/io/proto_table_reader.h>
#include <saas/library/rtyt/lib/io/proto_table_writer.h>
#include <saas/library/rtyt/lib/io/multi_table_reader.h>

namespace NRTYT {

inline ::TIntrusivePtr<NYT::INodeReaderImpl> CreateNodeReader(const TString& path) {
    return MakeIntrusive<TNodeReader>(path);
}
inline ::TIntrusivePtr<NYT::INodeWriterImpl> CreateNodeWriter(const TString& path) {
    return MakeIntrusive<TNodeWriter>(path);
}
inline ::TIntrusivePtr<NYT::INodeWriterImpl> CreateNodeWriter(const TVector<TString>& paths) {
    return MakeIntrusive<TNodeWriter>(paths);
}

inline ::TIntrusivePtr<NYT::IProtoWriterImpl> CreateProtoWriter(const TString& path, const ::google::protobuf::Descriptor* descriptor) {
    Y_ENSURE(descriptor, "Can't create proto writer, passed nullptr as a descriptor");
    return MakeIntrusive<TProtoWriter>(path, descriptor);
}
inline ::TIntrusivePtr<NYT::IProtoWriterImpl> CreateProtoWriter(const TVector<TString>& paths, const TVector<const ::google::protobuf::Descriptor*>& descriptors) {
    return MakeIntrusive<TProtoWriter>(paths, descriptors);
}

template <class T>
inline ::TIntrusivePtr<typename NYT::TRowTraits<T>::IReaderImpl> CreateJobReaderImpl(const TInputContext& context);

template <typename TReader, typename TWriter>
void FeedJobInput(
    NYT::IMapper<TReader, TWriter>* mapper,
    typename NYT::TRowTraits<typename TReader::TRowType>::IReaderImpl* readerImpl,
    TWriter* writer) {
    THolder<TReader> reader = MakeHolder<TReader>(readerImpl);
    mapper->Do(reader.Get(), writer);
}

template <typename TReader, typename TWriter>
void FeedJobInput(
    NYT::IReducer<TReader, TWriter>* reducer,
    typename NYT::TRowTraits<typename TReader::TRowType>::IReaderImpl* readerImpl,
    TWriter* writer) {
    THolder<TReader> reader = MakeHolder<TReader>(readerImpl);
    while (reader->IsValid()) {
        while (reader->IsValid()) {
            reducer->Do(reader.Get(), writer);
        }
        readerImpl->NextKey();
    }
}

template <class T>
inline NYT::TTableWriterPtr<T> CreateJobWriter(const TOutputContext& context)
{
    static_assert(std::is_base_of<NYT::Message, T>::value, "Unknown row type");
    return MakeIntrusive<NYT::TTableWriter<T>>(CreateProtoWriter(context.Paths.Parts_, context.Descriptors->Parts_));
}

template <class T>
inline ::TIntrusivePtr<typename NYT::TRowTraits<T>::IReaderImpl> CreateJobReaderImpl(const TInputContext& context)
{

    static_assert(std::is_base_of<NYT::Message, T>::value, "Unknown row type");
    if (context.Paths.Parts_.size() == 1) {
        return MakeIntrusive<TProtoReader>(
                context.Paths.Parts_[0],
                context.Descriptors->Parts_[0],
                context.GroupedColumns.GetOrElse({})
        );
    }
    return MakeIntrusive<TMultiProtoReader>(
                context.Paths.Parts_,
                context.Descriptors->Parts_,
                context.GroupedColumns.GetOrElse({}),
                context.SortedColumns.GetOrElse({})
    );
}

template <>
inline NYT::TTableWriterPtr<NYT::TNode> CreateJobWriter<NYT::TNode>(const TOutputContext& context)
{
    return new NYT::TTableWriter<NYT::TNode>(CreateNodeWriter(context.Paths.Parts_));
}

template <>
inline ::TIntrusivePtr<NYT::INodeReaderImpl> CreateJobReaderImpl<NYT::TNode>(const TInputContext& context)
{
    return CreateNodeReader(context.Paths.Parts_[0]);
}

template <>
inline NYT::TTableWriterPtr<NYT::Message> CreateJobWriter<NYT::Message>(const TOutputContext& context)
{
    if (context.Descriptors.Empty()) {
        ythrow yexception() << "got context without descriptors, can't create proto writer";
    }
    return new NYT::TTableWriter<NYT::Message>(CreateProtoWriter(context.Paths.Parts_, context.Descriptors->Parts_));
}

template <>
inline ::TIntrusivePtr<NYT::IProtoReaderImpl> CreateJobReaderImpl<NYT::Message>(const TInputContext& context)
{
    if (context.Descriptors.Empty()) {
        ythrow yexception() << "got context without descriptors, can't create proto reader";
    }
    // TODO(derrior): if there is one table, we can create TProtoReader which is slightly faster
    if (context.Paths.Parts_.size() == 1) {
        return MakeIntrusive<TProtoReader>(
                context.Paths.Parts_[0],
                context.Descriptors->Parts_[0],
                context.GroupedColumns.GetOrElse({})
        );
    }
    return MakeIntrusive<TMultiProtoReader>(
                context.Paths.Parts_,
                context.Descriptors->Parts_,
                context.GroupedColumns.GetOrElse({}),
                context.SortedColumns.GetOrElse({})
    );
}


template <class TJob>
int RunJob(IInputStream& jobStateStream, TInputContext& inputContext, TOutputContext& outputContext) {
    using TInputRow = typename TJob::TReader::TRowType;
    using TOutputRow = typename TJob::TWriter::TRowType;

    auto reader = CreateJobReaderImpl<TInputRow>(inputContext);
    auto writer = CreateJobWriter<TOutputRow>(outputContext);

    auto job = MakeIntrusive<TJob>();
    job->Load(jobStateStream);

    job->Start(writer.Get());
    NRTYT::FeedJobInput(job.Get(), reader.Get(), writer.Get());
    job->Finish(writer.Get());

    writer->Finish();

    return 0;
}


using TCustomIOJobFunction = int (*)(IInputStream&, TInputContext&, TOutputContext&);

class TCustomIOJobFactory {
public:
    static TCustomIOJobFactory* Get()
    {
        return Singleton<TCustomIOJobFactory>();
    }

    template <class TJob>
    void RegisterJob(const char* name)
    {
        RegisterJobImpl<TJob>(name, RunJob<TJob>);
    }

    TString GetJobName(const NYT::IJob* job)
    {
        const auto typeIndex = std::type_index(typeid(*job));
        CheckJobRegistered(typeIndex);
        return JobNames[typeIndex];
    }

    TCustomIOJobFunction GetJobFunction(const char* name)
    {
        CheckNameRegistered(name);
        return JobFunctions[name];
    }

private:
    TMap<std::type_index, TString> JobNames;
    THashMap<TString, TCustomIOJobFunction> JobFunctions;

    template <typename TJob>
    void RegisterJobImpl(const char* name, TCustomIOJobFunction runner) {
        const auto typeIndex = std::type_index(typeid(TJob));
        CheckNotRegistered(typeIndex, name);
        JobNames[typeIndex] = name;
        JobFunctions[name] = runner;
    }

// Зарегистрировать в фабрику тип без подстановки не получится. Но мы берём информацию о том, какие ридеры создавать, из спеки.
// План: в фабрике не выпендриваться, сделать также, только со своими свободными функциями CreateJobReader и передачей контекстов.
// Внутри RunJob есть все нужные для кастов типы. Внутри контекста будет ещё больше типов, а также туда запихаем, откуда брать таблицы (из каких файлов)
    void CheckNotRegistered(const std::type_index& typeIndex, const char* name)
    {
        Y_ENSURE(!JobNames.contains(typeIndex),
            "type_info '" << typeIndex.name() << "'"
            "is already registered under name '" << JobNames[typeIndex] << "'");
        Y_ENSURE(!JobFunctions.contains(name),
            "job with name '" << name << "' is already registered");
    }

    void CheckJobRegistered(const std::type_index& typeIndex)
    {
        Y_ENSURE(JobNames.contains(typeIndex),
            "type_info '" << typeIndex.name() << "' is not registered, use REGISTER_RTYT_* macros");
    }

    void CheckNameRegistered(const char* name)
    {
        Y_ENSURE(JobFunctions.contains(name),
            "job with name '" << name << "' is not registered, use REGISTER_RTYT_* macros");
    }
};


template <class TMapper>
struct TMapperRegistrator
{
    TMapperRegistrator(const char* name)
    {
        static_assert(TMapper::JobType == NYT::IJob::EType::Mapper,
            "REGISTER_MAPPER is not compatible with this job class");

        NRTYT::TCustomIOJobFactory::Get()->RegisterJob<TMapper>(name);
    }
};

template <class TReducer>
struct TReducerRegistrator
{
    TReducerRegistrator(const char* name)
    {
        static_assert(TReducer::JobType == NYT::IJob::EType::Reducer,
            "REGISTER_REDUCER is not compatible with this job class");

        NRTYT::TCustomIOJobFactory::Get()->RegisterJob<TReducer>(name);
    }
};

inline TString YtRegistryTypeName(const TString& name) {
    TString res = name;
#ifdef _win_
    SubstGlobal(res, "class ", "");
#endif
    return res;
}


#define REGISTER_RTYT_MAPPER(...) \
static NRTYT::TMapperRegistrator<__VA_ARGS__> \
Y_GENERATE_UNIQUE_ID(TJobRegistrator)(NRTYT::YtRegistryTypeName(TypeName<__VA_ARGS__>()).data());

#define REGISTER_RTYT_REDUCER(...) \
static NRTYT::TReducerRegistrator<__VA_ARGS__> \
Y_GENERATE_UNIQUE_ID(TJobRegistrator)(NRTYT::YtRegistryTypeName(TypeName<__VA_ARGS__>()).data());

} // namespace NRTYT
