#include "yt_chunked_input.h"

#include <util/generic/yexception.h>

IYTChunkedInputProcessor::~IYTChunkedInputProcessor() {
    // SAASSUP-3255
    // We can't call Finish() in the destructor because it can throw an exception.
    // Typical scenario: buffered file write, ENOSPC when flushing the last chunk in IOutputStream::Finish().
    // Therefore we need to make sure Finish() was called prior to destructor and errors were handled if any.
    Y_VERIFY(State == EState::Finished || State == EState::Failed, "CRASH TO PREVENT SILENT DATA LOSS. Invalid class usage: missing Finish() call.");
}

void IYTChunkedInputProcessor::ProcessAll() {
    try {
        DoProcessAll();
    } catch(...) {
        State = EState::Failed;
        throw;
    }
}

void IYTChunkedInputProcessor::DoProcessAll() {
    bool firstRow = true;
    for (; TableReader->IsValid(); TableReader->Next()) {
        TYTBlob row;
        if (Format == EYTBlobFormat::Old) {
            row = TYTBlob(TableReader->GetRow<NSaas::TYTBlobBase>());
        } else {
            row = TYTBlob(TableReader->GetRow<NSaas::TYTNewBlobBase>());
        }
        const auto& baseBlob = row.GetBlobRef();
        try {
            if (firstRow) {
                ShardId = row.GetShardId();
                firstRow = false;
            }
            ProcessRow(row);
        } catch (yexception& e) {
            e << " at row " << row.GetShardId() << "#" << baseBlob.GetName() << "#" << baseBlob.GetPartIndex();
            throw;
        }
    }
    Finish();
}

void IYTChunkedInputProcessor::Finish() {
    if (State == EState::Finished) {
        return;
    }
    if (State == EState::Failed) {
        ythrow yexception() << "Trying to finish a failed IYTChunkedInputProcessor";
    }
    if (CurrentConsumer != NamedConsumersMap.end()) {
        try {
            CurrentConsumer->second.Close();
        } catch (yexception& e) {
            State = EState::Failed;
            e << "failed to close file: " << CurrentConsumer->first.Quote();
            throw;
        } catch (...) {
            State = EState::Failed;
            throw;
        }
        CurrentConsumer = NamedConsumersMap.end();
    }
    // The CurrentConsumer should've been the last one, this is a private invariant of IYTChunkedInputProcessor.
    // Using VERIFY to check this.
    for (const auto& p: NamedConsumersMap) {
        Y_VERIFY(p.second.IsClosed(), "Not closed file handle: %s", p.first.Quote().c_str());
    }
    State = EState::Finished;
}

void IYTChunkedInputProcessor::ProcessRow(const TYTBlob& row) {
    Y_ENSURE(row.GetShardId() == ShardId);
    const auto& baseBlob = row.GetBlobRef();
    Y_ENSURE(baseBlob.GetDataLength() == baseBlob.GetData().size());
    GetConsumerByName(baseBlob.GetName()).ConsumeChunk(baseBlob.GetPartIndex(), baseBlob.GetData());
}

IYTChunkedInputProcessor::TChunkConsumer& IYTChunkedInputProcessor::GetConsumerByName(const TString& name) {
    if (CurrentConsumer != NamedConsumersMap.end()) {
        if (CurrentConsumer->first == name) {
            return CurrentConsumer->second;
        }
        // Next file started, need to destroy previous writer to free the write buffer memory and avoid OOMs
        CurrentConsumer->second.Close();
    }
    CurrentConsumer = NamedConsumersMap.find(name);
    if (CurrentConsumer == NamedConsumersMap.end()) {
        auto emplaceResult = NamedConsumersMap.emplace(name, CreateOutputStream(name));
        Y_ENSURE(emplaceResult.second == true);
        CurrentConsumer = emplaceResult.first;
    } else {
        CurrentConsumer->second.Resume(ResumeOutputStream(name, CurrentConsumer->second.GetOffset()));
    }
    return CurrentConsumer->second;
}
