#include "proto_table_reader.h"
#include <saas/library/rtyt/lib/util/proto_utils.h>
#include <google/protobuf/util/field_comparator.h>

namespace NRTYT {

void TProtoReader::ReadRow(NYT::Message* row) {
    if (!IsValid()) {
        ythrow yexception() << "Failed while reading from invalid reader!" << ' '
                            << "is range switched: " << RangeSwitched << ' '
                            << "is reading finished: " << (RowId >= Archive->GetDocsCount(false));
    }
    if (row->GetDescriptor() != Descriptor) {
        ythrow yexception() <<
            "Invalid row of type " << row->GetDescriptor()->full_name() <<
            ", row of type " << Descriptor->full_name();
    }
    row->CopyFrom(*CurrentRow);
}

bool TProtoReader::IsValid() const {
    return !RangeSwitched && RowId < Archive->GetDocsCount(false);
}

void TProtoReader::ReadFromArchive() {
    TBlob blob = Archive->GetDocument(RowId);
    if (!CurrentRow.Get()) {
        CurrentRow.Reset(MessageFactory->GetPrototype(Descriptor)->New());
    }

    ::google::protobuf::io::CodedInputStream protoStream(blob.AsUnsignedCharPtr(), blob.Size());
    if (!CurrentRow->ParseFromCodedStream(&protoStream)) {
        ythrow yexception() << "Failed to parse protobuf message";
    }
}

bool TProtoReader::CheckCurrentRange() const {
    if (KeyColumns.Empty()) {
        return true;
    }
    ::google::protobuf::util::DefaultFieldComparator comparator;
    for (auto fieldDescrPtr : KeyColumns.GetRef()) {
        auto comparisonResult = comparator.Compare(*CurrentRow, *CurrentRange, fieldDescrPtr,
                            /*needless arguments*/ 0, 0, nullptr);
        if (comparisonResult == ::google::protobuf::util::FieldComparator::ComparisonResult::DIFFERENT) {
            return false;
        }
    }
    return true;
}

void TProtoReader::Next() {
    if (RangeSwitched) {
        return;
    }
    RowId++;
    if (RowId >= Archive->GetDocsCount(false)) {
        CurrentRow.Destroy();
        CurrentRange.Destroy();
        return;
    }
    ReadFromArchive();

    if (!CheckCurrentRange()) {
        RangeId++;
        RangeSwitched = true;
        CurrentRange = CopyMessage(CurrentRow);
    }
}

ui32 TProtoReader::GetTableIndex() const {
    return 0;
}
ui32 TProtoReader::GetRangeIndex() const {
    return RangeId;
}
ui64 TProtoReader::GetRowIndex() const {
    return RowId;
}
void TProtoReader::NextKey() {
    while (IsValid()) {
        Next();
    }
    RangeSwitched = false;
}

} // namespace NRTYT
