#include "multi_table_reader.h"

namespace NRTYT {

void TMultiProtoReader::ReadRow(NYT::Message* row) {
    if (!IsValid()) {
        ythrow yexception() << "Failed while reading from invalid reader!" << ' '
                            << "is range switched: " << RangeSwitched << ' '
                            << "is reading finished: " << (TableIndex == Max<ui32>()) << ". "
                            << "Debug: rowId = " << RowId << " rangeId = " << RangeId;
    }
    Readers[TableIndex].ReadRow(row);
}

void TMultiProtoReader::UpdateTableIndex() {
   TableIndex = Max<ui32>();
   for (ui32 i = 0; i < MessageBuffer.size(); i++) {
       if (Readers[i].IsValid()) {
           if (TableIndex == Max<ui32>() || 
                RowComparator(MessageBuffer[i].Get(), i,
                            MessageBuffer[TableIndex].Get(), TableIndex)) {
               TableIndex = i;
           }
       }
   }
}

void TMultiProtoReader::Next() {
    if (!IsValid()) {
        // Nothing left to read in current range
        return;
    }
    TRangeItem current({CopyMessage(MessageBuffer[TableIndex]), TableIndex});
    Readers[TableIndex].Next();
    RowId++;
    if (Readers[TableIndex].IsValid()) {
        Readers[TableIndex].ReadRow(MessageBuffer[TableIndex].Get());
        UpdateTableIndex();
        // We can remain in the same range, but need to update tableIndex
        return;
    }

    // next range may be in same table, so we should put new range before getting next range
    Readers[TableIndex].NextKey();
    if (Readers[TableIndex].IsValid()) {
        Readers[TableIndex].ReadRow(MessageBuffer[TableIndex].Get());
    }
    UpdateTableIndex();

    if (TableIndex == Max<ui32>()) {
        return;
    }

    TRangeItem next({MessageBuffer[TableIndex], TableIndex});

    if (RangesComparator(current, next)) { // (current < next)
        // current range ended, so we need to stop
        RangeSwitched = true;
    }
}

bool TMultiProtoReader::IsValid() const {
    return !RangeSwitched && TableIndex != Max<ui32>();
}

ui32 TMultiProtoReader::GetTableIndex() const {
    return TableIndex;
}

ui32 TMultiProtoReader::GetRangeIndex() const {
    return RangeId;
}

ui64 TMultiProtoReader::GetRowIndex() const {
    return RowId;
}

void TMultiProtoReader::NextKey() {
    while (IsValid()) {
        Next();
    }
    if (RangeSwitched) {
        RangeSwitched = false;
        RangeId++;
    }
}

} // namespace NRTYT

