#include "archive.h"

#include <saas/protos/fast_archive.pb.h>

#include <library/cpp/logger/global/global.h>

#include <google/protobuf/messagext.h>

#include <util/draft/holder_vector.h>
#include <library/cpp/containers/mh_heap/mh_heap.h>
#include <util/folder/dirut.h>
#include <util/stream/file.h>
#include <util/system/fs.h>
#include <util/string/cast.h>

namespace {
    bool ReadMessageFromStream(IInputStream& input, ::google::protobuf::Message& message) {
        ui64 size = 0;
        if (input.Load(&size, sizeof(size)) != sizeof(size))
            return false;
        Y_ASSERT(size);

        TVector<ui8> buffer(size);
        ui8* ptr = &buffer[0];

        if (input.Load(ptr, size) != size)
            return false;

        return message.ParseFromArray(ptr, size);
    }

    bool WriteMessageToStream(IOutputStream& output, const ::google::protobuf::Message& message) {
        ui64 size = message.ByteSize();
        Y_ASSERT(size);

        TVector<ui8> buffer(size, 0);
        ui8* ptr = &buffer[0];

        if (!message.SerializeToArray(ptr, size))
            return false;

        output.Write(&size, sizeof(size));
        output.Write(ptr, size);
        return true;
    }
}

void WriteArchiveHeader(IOutputStream& output, ui32 version, const TSet<TString>& properties) {
    NRTYServer::TFastArchiveHeader header;
    header.SetVersion(version);
    for (const TString& p : properties)
        header.AddProperties(p);
    VERIFY_WITH_LOG(WriteMessageToStream(output, header), "cannot serialize FASTARC header");
}

bool ReadArchiveHeader(IInputStream& input, ui32& version, TSet<TString>& properties) {
    NRTYServer::TFastArchiveHeader header;
    if (!ReadMessageFromStream(input, header))
        return false;

    version = header.GetVersion();
    properties.clear();
    for (ui32 i = 0; i < header.PropertiesSize(); ++i) {
        const TString& property = header.GetProperties(i);
        if (properties.insert(property).second)
            continue;

        ERROR_LOG << "Incorrect FASTARC header: property " << property << " is defined twice" << Endl;
        return false;
    }

    return true;
}

bool ReadArchiveHeader(const TString& file, ui32& version, TSet<TString>& properties) {
    if (!NFs::Exists(file))
        return false;

    TIFStream input(file);
    return ReadArchiveHeader(input, version, properties);
}

bool TFastArchiveDoc::Parse(IInputStream& input) {
    Properties.clear();
    NRTYServer::TFastArchiveDoc doc;
    if (!ReadMessageFromStream(input, doc))
        return false;

    SetDocId(doc.GetDocId());
    for (size_t i = 0; i < doc.PropertiesSize(); ++i) {
        const TString& name = doc.GetProperties(i).GetName();
        TVector<size_t> values;
        for (auto&& v : doc.GetProperties(i).GetValues()) {
            values.push_back(AddString(v));
        }

        TProperty prop = { values, AddString(name) };
        Properties.push_back(prop);
    }

    return true;
}

bool TFastArchiveDoc::Serialize(IOutputStream& output) const {
    NRTYServer::TFastArchiveDoc doc;
    doc.SetDocId(GetDocId());
    for (auto&& p : Properties) {
        const TStringBuf name = GetString(p.Name);

        auto protobufProp = doc.AddProperties();
        protobufProp->SetName(name.data(), name.size());
        for (auto&& v : p.Values) {
            const TStringBuf value = GetString(v);
            protobufProp->AddValues(value.data(), value.size());
        }
    }
    return WriteMessageToStream(output, doc);
}

TStringBuf TFastArchiveDoc::GetPropertyName(TIndex index) const {
    Y_ASSERT(index < Properties.size());
    return GetString(Properties[index].Name);
}

TFastArchiveDoc::TIndex TFastArchiveDoc::GetPropertyIndex(const TStringBuf& name) const {
    for (TIndex i = 0; i < GetPropertiesCount(); ++i) {
        if (GetPropertyName(i) == name)
            return i;
    }

    return InvalidIndex;
}

TVector<TStringBuf> TFastArchiveDoc::GetPropertyValues(TIndex index) const {
    Y_ASSERT(index < Properties.size());
    TVector<TStringBuf> result;
    for (auto&& v : Properties[index].Values) {
        result.push_back(GetString(v));
    }
    return result;
}

void TFastArchiveDoc::FillFrom(const TFastArchiveDocWriter& writer) {
    Properties.clear();
    for (auto&& p : writer.Properties) {
        TProperty property;
        property.Name = AddString(p.first);
        for (auto&& v : p.second) {
            const size_t pos = AddString(v);
            property.Values.push_back(pos);
        }
        Properties.push_back(property);
    }
}

size_t TFastArchiveDoc::AddString(TStringBuf string) {
    Y_ASSERT(StringPool);
    return StringPool->AddString(string);
}

TStringBuf TFastArchiveDoc::GetString(size_t index) const {
    Y_ASSERT(StringPool);
    return StringPool->GetString(index);
}

bool TFastArchiveDocWriter::Add(const TString& name, const TString& value) {
    return Properties[name].insert(value).second;
}

bool TFastArchiveDocWriter::Empty() const {
    return !Properties.size();
}

void TFastArchiveDocWriter::ApplyPatch(const TFastArchiveDocWriter& patch)
{
    for (const auto& p : patch.Properties) {
        if (p.second.size() == 1 && p.second.contains("__delete__"))
            Properties.erase(p.first);
        else
            Properties[p.first] = p.second;
    }
}


bool TFastArchivePortionsMerger::Merge(bool deletePortions /*= true*/) {
    TVector<TAutoPtr<IInputStream>> inputs;
    for (const TString& portion : Portions)
        inputs.push_back(new TIFStream(portion));

    TSet<TString> properties;
    for (ui32 i = 0; i < inputs.size(); ++i) {
        ui32 version;

        if (!ReadArchiveHeader(*inputs[i], version, properties)) {
            ERROR_LOG << "Cannot read FASTARC header for portion " << Portions[i] << Endl;
            return false;
        }

        if (Version != version) {
            ERROR_LOG << "FASTARC version mismatch for portion " << Portions[i] << ": " << version << " != " << Version << Endl;
            return false;
        }
    }

    const TString& temp = Target + "_" + ToString(Now().MilliSeconds());
    TOFStream output(temp);
    WriteArchiveHeader(output, Version, properties);

    TFastArchiveStringPool StringPool;
    THolderVector<TFastArchiveIterator> iterators;
    for (ui32 i = 0; i < inputs.size(); ++i)
        iterators.PushBack(new TFastArchiveIterator(*inputs[i], StringPool));

    MultHitHeap<TFastArchiveIterator> merger(iterators.data(), iterators.size());
    for (merger.Restart(); merger.Valid(); ++merger) {
        TFastArchiveIterator* p = merger.TopIter();
        TFastArchiveDoc& doc = p->GetDoc();
        if (Remapper && doc.GetDocId() < RemapperSize) {
            const ui32 newDocId = Remapper[doc.GetDocId()];
            if (newDocId == (ui32)-1)
                continue;

            doc.SetDocId(newDocId);
        }
        doc.Serialize(output);
    }

    output.Flush();
    inputs.clear();
    if (deletePortions) {
        for (const TString& portion : Portions)
            NFs::Remove(portion);
    }

    NFs::Rename(temp, Target);
    return true;
}

size_t TFastArchiveStringPool::AddString(const TStringBuf& string) {
    if (Finalized)
        ythrow yexception() << "string pool is finalized";

    const TString& s = TString(string);
    auto p = Offsets.find(s);
    if (p != Offsets.end()) {
        return p->second;
    }

    const size_t pos = Data.Size();
    Data.Append(string.data(), string.size());
    Data.Append(0);

    TStringBuf poolStr = GetString(pos);
    if (string != poolStr) {
        ERROR_LOG << "inconsistency: original string: " << string << " added string: " << poolStr << Endl;
        Y_ASSERT(string == poolStr);
    }

    Offsets[s] = pos;
    return pos;
}

TStringBuf TFastArchiveStringPool::GetString(size_t index) const {
    if (index >= Data.Size())
        ythrow yexception() << "incorrect offset " << index << " >= " << Data.Size();

    return TStringBuf(Data.Data() + index);
}

void TFastArchiveStringPool::Finalize() {
    Finalized = true;
    Offsets.clear();
}

void TFastArchiveStringPool::Clear() {
    Finalized = false;
    Data.Reset();
    Offsets.clear();
}
