#include "client.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/operation/factory.h>
#include <saas/library/rtyt/lib/util/proto_utils.h>

#include <util/stream/str.h>
#include <google/protobuf/descriptor.h>
#include <google/protobuf/descriptor.pb.h>

namespace NRTYT {

    TClientBase::TClientBase(TMaybe<TFsPath> storageRoot, bool failIfNotExist) {
        CypressClient = MakeIntrusive<TCypressClient>(std::move(storageRoot), failIfNotExist);
    }

    NYT::ITransactionPtr TClientBase::StartTransaction(
        const NYT::TStartTransactionOptions&) {
        return MakeIntrusive<TTransaction>(*this);
    }

    NYT::TNodeId TClientBase::Create(
        const NYT::TYPath& path,
        NYT::ENodeType type,
        const NYT::TCreateOptions& options) {
        return CypressClient->Create(path, type, options);
    }

    void TClientBase::Remove(
        const NYT::TYPath& path,
        const NYT::TRemoveOptions& options) {
        return CypressClient->Remove(path, options);
    }

    bool TClientBase::Exists(
        const NYT::TYPath& path,
        const NYT::TExistsOptions& options) {
        return CypressClient->Exists(path, options);
    }

    NYT::TNode TClientBase::Get(
        const NYT::TYPath& path,
        const NYT::TGetOptions& options) {
        return CypressClient->Get(path, options);
    }

    void TClientBase::Set(
        const NYT::TYPath& path,
        const NYT::TNode& value,
        const NYT::TSetOptions& options) {
        return CypressClient->Set(path, value, options);
    }
    void TClientBase::MultisetAttributes(
        const NYT::TYPath& path,
        const NYT::TNode::TMapType& value,
        const NYT::TMultisetAttributesOptions& options) {
        return CypressClient->MultisetAttributes(path, value, options);
    }
    NYT::TNode::TListType TClientBase::List(
        const NYT::TYPath& path,
        const NYT::TListOptions& options) {
        return CypressClient->List(path, options);
    }
    NYT::TNodeId TClientBase::Copy(
        const NYT::TYPath& sourcePath,
        const NYT::TYPath& destinationPath,
        const NYT::TCopyOptions& options) {
        return CypressClient->Copy(sourcePath, destinationPath, options);
    }
    NYT::TNodeId TClientBase::Move(
        const NYT::TYPath& sourcePath,
        const NYT::TYPath& destinationPath,
        const NYT::TMoveOptions& options) {
        return CypressClient->Move(sourcePath, destinationPath, options);
    }
    NYT::TNodeId TClientBase::Link(
        const NYT::TYPath& targetPath,
        const NYT::TYPath& linkPath,
        const NYT::TLinkOptions& options) {
        return CypressClient->Link(targetPath, linkPath, options);
    }
    void TClientBase::Concatenate(
        const TVector<NYT::TRichYPath>& sourcePaths,
        const NYT::TRichYPath& destinationPath,
        const NYT::TConcatenateOptions& options) {
        return CypressClient->Concatenate(sourcePaths, destinationPath, options);
    }
    NYT::TRichYPath TClientBase::CanonizeYPath(const NYT::TRichYPath& path) {
        return CypressClient->CanonizeYPath(path);
    }
    TVector<NYT::TTableColumnarStatistics> TClientBase::GetTableColumnarStatistics(
        const TVector<NYT::TRichYPath>& paths,
        const NYT::TGetTableColumnarStatisticsOptions& options) {
        return CypressClient->GetTableColumnarStatistics(paths, options);
    }
    TMaybe<NYT::TYPath> TClientBase::GetFileFromCache(
        const TString& md5Signature,
        const NYT::TYPath& cachePath,
        const NYT::TGetFileFromCacheOptions& options) {
        return CypressClient->GetFileFromCache(md5Signature, cachePath, options);
    }
    NYT::TYPath TClientBase::PutFileToCache(
        const NYT::TYPath& filePath,
        const TString& md5Signature,
        const NYT::TYPath& cachePath,
        const NYT::TPutFileToCacheOptions& options) {
        return CypressClient->PutFileToCache(filePath, md5Signature, cachePath, options);
    }

    NYT::TTableWriterPtr<NYT::TNode> TClientBase::CreateTableWriter(
        const NYT::TRichYPath& path,
        const NYT::TTableWriterOptions& options) {
        NYT::TCreateOptions opts;
        opts.Force(!path.Append_.GetOrElse(false)).IgnoreExisting(path.Append_.GetOrElse(false));
        CypressClient->Create(path.Path_, ENodeType::NT_TABLE, opts);
        return MakeIntrusive<NYT::TTableWriter<NYT::TNode>>(CreateNodeWriter(path, options));
    }

    NYT::TTableWriterPtr< ::google::protobuf::Message> TClientBase::CreateTableWriter(
        const NYT::TRichYPath& path,
        const ::google::protobuf::Descriptor& descriptor,
        const NYT::TTableWriterOptions& options) {
        ::google::protobuf::DynamicMessageFactory descriptorFactory;
        const auto* prototype = descriptorFactory.GetPrototype(&descriptor);
        CreateOperationOutput(path, &descriptor);
        return MakeIntrusive<NYT::TTableWriter< ::google::protobuf::Message>>(CreateProtoWriter(path, options, prototype));
    }

    NYT::IOperationPtr TClientBase::DoMap(
        const NYT::TMapOperationSpec& spec,
        ::TIntrusivePtr<NYT::IStructuredJob> mapper,
        const NYT::TOperationOptions&) {
        CreateOperationOutput(spec);
        TInputContext inputContext = GetMapInputContext(spec);
        TOutputContext outputContext = GetMapOutputContext(spec);

        TString state;
        TStringOutput stateSave(state);
        mapper->Save(stateSave);
        TStringInput stateLoad(state);
        auto jobName = TCustomIOJobFactory::Get()->GetJobName(mapper.Get());
        auto jobRunner = TCustomIOJobFactory::Get()->GetJobFunction(jobName.data());
        jobRunner(stateLoad, inputContext, outputContext);
        
        return nullptr;
    }

    NYT::IOperationPtr TClientBase::DoReduce(
        const NYT::TReduceOperationSpec& spec,
        ::TIntrusivePtr<NYT::IStructuredJob> reducer,
        const NYT::TOperationOptions&) {
        CheckReducerInputTables(spec);
        CreateOperationOutput(spec);
        TInputContext inputContext = GetReduceInputContext(spec);
        TOutputContext outputContext = GetReduceOutputContext(spec);

        TString state;
        TStringOutput stateSave(state);
        reducer->Save(stateSave);
        TStringInput stateLoad(state);
        auto jobName = TCustomIOJobFactory::Get()->GetJobName(reducer.Get());
        auto jobRunner = TCustomIOJobFactory::Get()->GetJobFunction(jobName.data());
        jobRunner(stateLoad, inputContext, outputContext);
        return nullptr;
    }

    NYT::IOperationPtr TClientBase::Sort(
        const NYT::TSortOperationSpec& spec,
        const NYT::TOperationOptions&) {
        TInputContext inputContext = GetSortInputContext(spec);

        auto descriptor = inputContext.Descriptors->Parts_[0];
        CreateOperationOutput(spec.Output_.Path_, descriptor);
        Y_ENSURE(descriptor, "While preparing sort operation, found nullptr instead of descriptor");
        TOutputContext outputContext = GetSortOutputContext(spec);
        TVector<const ::google::protobuf::FieldDescriptor*> keyDescriptors = 
            GetFieldDescriptorsByName(
                descriptor,
                spec.SortBy_.EnsureAscending().GetNames(),
                /*ignoreUnknownColumns = */ false);

        /*
         * We need to call destructors for various protobuf objects in the following order
         * 1. destruct messages, generated with dynamicFactory (data)
         * 2. destruct factory, all messages should be destroyed. (descriptorFactory)
         *    This should be done before destructing any of descriptors
         * 3. destruct descriptors which were used with factory (descriptor)
        */
        ::google::protobuf::DynamicMessageFactory descriptorFactory;
        TVector<THolder<::google::protobuf::Message>> data;
        const ::google::protobuf::Message* prototype = descriptorFactory.GetPrototype(descriptor);
        auto reader = CreateProtoReader(spec.Inputs_[0], NYT::TTableReaderOptions(), prototype);
        THolder<::google::protobuf::Message> tmpMessage = THolder(prototype->New());
        for (; reader->IsValid(); reader->Next()) {
            reader->ReadRow(tmpMessage.Get());
            data.push_back(std::move(tmpMessage));
            tmpMessage.Reset(prototype->New());
        }

        SortProtoVector(data, keyDescriptors);

        auto writer = NRTYT::CreateProtoWriter(outputContext.Paths.Parts_[0], descriptor);
        for (const auto& row : data) {
            writer->AddRow(*row.Get(), 0);
        }
        return nullptr;
    }

    NYT::IClientPtr TClientBase::GetParentClient() {
        return MakeIntrusive<TClient>(*this);
    }

// private functions
    void TClientBase::CreateOperationOutput(const NYT::TRichYPath& path, const ::google::protobuf::Descriptor* descriptor) {
        NYT::TCreateOptions opts;
        opts.Force(!path.Append_.GetOrElse(false)).IgnoreExisting(path.Append_.GetOrElse(false));
        if (descriptor) {
            auto fileDescriptor = descriptor->file();
            auto messageName = descriptor->full_name();
            opts.Attributes_ = NYT::TNode()
                            ("@_rtyt_file_descriptor", SerializeFileDescriptor(fileDescriptor))
                            ("@_rtyt_message_name", messageName);
        }
        CypressClient->Create(path.Path_, ENodeType::NT_TABLE, opts);
    }

    void TClientBase::CreateOperationOutput(const NYT::TOperationOutputSpecBase& operationSpec) {
        for (const auto& output : operationSpec.GetStructuredOutputs()) {
            const ::google::protobuf::Descriptor* currentDescriptor = nullptr;
            if (const auto* protoStruct = std::get_if<NYT::TProtobufTableStructure>(&output.Description)) {
                currentDescriptor = protoStruct->Descriptor;
            }
            CreateOperationOutput(output.RichYPath, currentDescriptor);
        }
    }

    void TClientBase::FillOperationInputs(const NYT::TOperationInputSpecBase& spec, TInputContext& context) {
        auto input = spec.GetStructuredInputs()[0];
        if (std::holds_alternative<NYT::TProtobufTableStructure>(input.Description)) {
            context.Descriptors.ConstructInPlace();
        }
        for (const auto& currentInput : spec.GetStructuredInputs()) {
            try {
                context.Paths.Add(CypressClient->GetTableStorage(currentInput.RichYPath.Path_));
            } catch (...) {
                WARNING_LOG << "Skip non-existent operation input: "
                            << currentInput.RichYPath.Path_ << "; error: "
                            << CurrentExceptionMessage() << Endl; 
                continue;
            }
            if (context.Descriptors.Defined()) {
                Y_ENSURE(std::holds_alternative<NYT::TProtobufTableStructure>(currentInput.Description),
                        "There is non-protobuf table along with protobuf in one operation, which is currently not supported");
                context.Descriptors->Add(std::get<NYT::TProtobufTableStructure>(currentInput.Description).Descriptor); // TODO(derrior): check if two descriptors are equal (from attributes and from spec)
            }
        }
        Y_ENSURE(!context.Paths.Parts_.empty(), "Operation must have at least one input table");
    }

    void TClientBase::FillOperationOutputs(const NYT::TOperationOutputSpecBase& spec, TOutputContext& context) {
        auto output = spec.GetStructuredOutputs()[0];
        if (std::holds_alternative<NYT::TProtobufTableStructure>(output.Description)) {
            context.Descriptors.ConstructInPlace();
        }
        for (const auto& currentOutput : spec.GetStructuredOutputs()) {
            context.Paths.Add(CypressClient->GetTableStorage(currentOutput.RichYPath.Path_));
            if (context.Descriptors.Defined()) {
                Y_ENSURE(std::holds_alternative<NYT::TProtobufTableStructure>(output.Description),
                        "There is non-protobuf table along with protobuf in one operation, which is currently not supported");
                context.Descriptors->Add(std::get<NYT::TProtobufTableStructure>(currentOutput.Description).Descriptor); // TODO(derrior): check if two descriptors are equal (from attributes and from spec)
            }
        }
    }

    const ::google::protobuf::Descriptor* TClientBase::RestoreProtoFromAttrs(const NYT::TRichYPath& path) {
        return CypressClient->RestoreProtoFromAttrs(path.Path_);
    }


    TInputContext TClientBase::GetMapInputContext(const NYT::TMapOperationSpec& spec) {
        TInputContext inputContext;
        FillOperationInputs(spec, inputContext);
        return inputContext;
    }
    TOutputContext TClientBase::GetMapOutputContext(const NYT::TMapOperationSpec& spec) {
        TOutputContext outputContext;
        FillOperationOutputs(spec, outputContext);
        return outputContext;
    }

    TInputContext TClientBase::GetReduceInputContext(const NYT::TReduceOperationSpec& spec) {
        TInputContext inputContext;
        FillOperationInputs(spec, inputContext);
        inputContext.GroupedColumns = spec.ReduceBy_;
        inputContext.SortedColumns = spec.SortBy_;
        if (spec.SortBy_.Parts_.empty()) {
            inputContext.SortedColumns = spec.ReduceBy_;
        } else {
            Y_ENSURE(spec.SortBy_.Parts_.size() >= spec.ReduceBy_.Parts_.size(), "ReduceBy parameter of spec must be prefix of SortBy parameter, but ReduceBy is longer");
            for (size_t i = 0; i < spec.ReduceBy_.Parts_.size(); i++) {
                Y_ENSURE(spec.SortBy_.Parts_[i] == spec.ReduceBy_.Parts_[i], "ReduceBy parameter of spec must be prefix of SortBy parameter, but they differ in position " << i <<": SortBy: " << spec.SortBy_.Parts_[i] << ", ReduceBy: " << spec.ReduceBy_.Parts_[i]);
            }
        }


        return inputContext;
    }
    TOutputContext TClientBase::GetReduceOutputContext(const NYT::TReduceOperationSpec& spec) {
        TOutputContext outputContext;
        FillOperationOutputs(spec, outputContext);
        return outputContext;
    }

    TInputContext TClientBase::GetSortInputContext(const NYT::TSortOperationSpec& spec) {
        TInputContext inputContext;
        inputContext.Descriptors.ConstructInPlace();
        for (const auto& currentInput : spec.Inputs_) {
            try {
                inputContext.Paths.Add(CypressClient->GetTableStorage(currentInput.Path_));
            } catch (...) {
                ythrow yexception() <<"Non-existent operation input: "
                            << currentInput.Path_ << "; error: "
                            << CurrentExceptionMessage();
            }
            const auto currentDescriptor = RestoreProtoFromAttrs(currentInput);
            inputContext.Descriptors->Add(currentDescriptor);
        }
        Y_ENSURE(!inputContext.Paths.Parts_.empty(), "Operation must have at least one input table");
        return inputContext;
    }
    TOutputContext TClientBase::GetSortOutputContext(const NYT::TSortOperationSpec& spec) {
        auto output = spec.Output_;
        auto path = CypressClient->GetTableStorage(output.Path_);

        const auto descriptor = RestoreProtoFromAttrs(output);
        TOutputContext outputContext({path, descriptor});
        return outputContext;
    }

    ::TIntrusivePtr<NYT::INodeReaderImpl> TClientBase::CreateNodeReader(
        const NYT::TRichYPath& path,
        const NYT::TTableReaderOptions&) {
        return MakeIntrusive<TNodeReader>(CypressClient->GetTableStorage(path.Path_));
    }

    ::TIntrusivePtr<NYT::IProtoReaderImpl> TClientBase::CreateProtoReader(
        const NYT::TRichYPath& path,
        const NYT::TTableReaderOptions& options,
        const ::google::protobuf::Message* prototype) {
        Y_UNUSED(options);
        return MakeIntrusive<TProtoReader>(
            CypressClient->GetTableStorage(path.Path_), 
            prototype->GetDescriptor());
    }

    ::TIntrusivePtr<NYT::ISkiffRowReaderImpl> TClientBase::CreateSkiffRowReader(
        const NYT::TRichYPath& /*path*/,
        const NYT::TTableReaderOptions& /*options*/,
        const NYT::ISkiffRowSkipperPtr& /*skipper*/,
        const NSkiff::TSkiffSchemaPtr& /*schema*/) {
        Y_ENSURE("CreateSkiffRowReader not implemented");
        return nullptr;
    }

    ::TIntrusivePtr<NYT::INodeWriterImpl> TClientBase::CreateNodeWriter(
        const NYT::TRichYPath& path,
        const NYT::TTableWriterOptions&) {
        return MakeIntrusive<TNodeWriter>(CypressClient->GetTableStorage(path.Path_));
    }

    ::TIntrusivePtr<NYT::IProtoWriterImpl> TClientBase::CreateProtoWriter(
        const NYT::TRichYPath& path,
        const NYT::TTableWriterOptions& options,
        const ::google::protobuf::Message* prototype) {
        Y_UNUSED(options);
        return MakeIntrusive<TProtoWriter>(
            CypressClient->GetTableStorage(path.Path_), 
            prototype->GetDescriptor());
    }

    NYT::TNodeId TClientBase::GetNodeId(const NYT::TYPath& path) {
        return CypressClient->GetNodeId(path);
    }

    void TClientBase::Sync(TMaybe<NYT::TNodeId> id, bool needFlush) {
        CypressClient->Sync(id, needFlush);
    }

    void TClientBase::Mount(const TFsPath& storagePath,
                            const NYT::TYPath& mountPoint,
                            bool failIfNotExist) {
        CypressClient->Mount(storagePath, mountPoint, failIfNotExist);
    }
    void TClientBase::Umount(NYT::TNodeId mountId) {
        CypressClient->Umount(mountId);
    }

    ::TIntrusivePtr<TClient> MakeClient(TMaybe<TFsPath> storageRoot, bool failIfNotExist) {
        return new TClient(TClientBase(storageRoot, failIfNotExist));
    }
}
