#include "data.h"

#include <saas/rtyserver/components/ann/arraydata/iterator.h>
#include <saas/rtyserver/components/ann/arraydata/writer.h>

#include <saas/rtyserver/components/ann/storage/accessor.h>
#include <saas/rtyserver/components/ann/storage/iterator.h>
#include <saas/rtyserver/components/ann/storage/writer.h>
#include <saas/rtyserver/factors/factors_config.h>
#include <saas/rtyserver/config/config.h>

namespace {
    class TBlobAccessor : public NRTYAnn::IDocBlobAccessor {
    public:
        TBlobAccessor(TArchiveOwner::TPtr source)
            : Source(source)
        {
        }
        TBlob Get(ui32 docId) const override {
            return Source->GetDocument(docId);
        }
    private:
        TArchiveOwner::TPtr Source;
    };
}

bool TAnnDataFormatNormalizer::AllRight(const NRTYServer::TNormalizerContext& context, const THolder<TFileMap>& /*indexFrq*/) const {
    return GetFormat(context) == Format;
}

void TAnnDataFormatNormalizer::Fix(const NRTYServer::TNormalizerContext& context, const THolder<TFileMap>& /*indexFrq*/) const {
    if (GetFormat(context) != Format) {
        ConvertData(context);
    }
}

NRTYAnn::EDataFormat TAnnDataFormatNormalizer::GetFormat(const NRTYServer::TNormalizerContext& context) const {
    TIndexMetadataProcessor meta(context.Dir.PathName());
    return GetFormat(meta, Suffix);
}

NRTYAnn::EDataFormat TAnnDataFormatNormalizer::GetFormat(const TIndexMetadataProcessor& meta, const TString& suffix) {
    for (const auto& header : meta->GetAnnHeaders()) {
        if (header.GetName() == suffix)
            return FromString<NRTYAnn::EDataFormat>(header.GetDataFormat());
    }
    const NRTYServer::TIndexMetadata::TAnnotationHeader* header = nullptr;

    if (suffix == "ann") {
        header = &meta->GetAnnHeader();
    }
    else if (suffix == "factorann") {
        header = &meta->GetFactorAnnHeader();
    }
    return header ? FromString<NRTYAnn::EDataFormat>(header->GetDataFormat()) : NRTYAnn::EDataFormat::COMPTRIE;
}

void TAnnDataFormatNormalizer::ConvertData(const NRTYServer::TNormalizerContext& context) const {
    const TFsPath root = context.Dir.PathName();
    {
        auto source = TArchiveOwner::Create(root / ("index" + Suffix), MultipartConfig);
        auto destination = TArchiveOwner::Create(
            root / ("normalizedindex" + Suffix),
            MultipartConfig,
            source->GetDocsCount(true)
        );

        THolder<NIndexAnn::IDocDataIterator> iterator;
        TBlobAccessor accessor(source);
        switch (GetFormat(context)) {
        case NRTYAnn::EDataFormat::COMPTRIE:
            iterator = MakeHolder<NRTYAnn::TStreamIterator>(&accessor, new NRTYAnn::TAnnDataAccessor(
                *Config.GetSearcherConfig().Factors->GetAnnStreams(Suffix)));
            break;
        case NRTYAnn::EDataFormat::PLAIN_ARRAY:
            iterator = NewArrayDocDataIterator(&accessor);
            break;
        }

        THolder<NIndexAnn::IDocDataWriter> writer;
        switch (Format) {
        case NRTYAnn::EDataFormat::COMPTRIE:
            writer = MakeHolder<NRTYAnn::TStreamDataWriter>(destination, new NRTYAnn::TAnnDataAccessor(
                *Config.GetSearcherConfig().Factors->GetAnnStreams(Suffix)));
            break;
        case NRTYAnn::EDataFormat::PLAIN_ARRAY:
            writer = NewArrayDocDataWriter(destination);
            break;
        }

        AssertCorrectConfig(!!iterator, "Cannot create Ann iterator");
        AssertCorrectConfig(!!writer, "Cannot create Ann writer");

        for (ui32 docId = 0; docId < source->GetDocsCount(true); ++docId) {
            if (source->IsRemoved(docId)) {
                continue;
            }

            writer->StartDoc(docId);
            NIndexAnn::THitMask mask(docId);
            for (iterator->Restart(mask); iterator->Valid(); iterator->Next()) {
                const NIndexAnn::THit& hit = iterator->Current();
                const ui32 value = hit.Value();
                AssertCorrectIndex(hit.DocId() == docId, "DocId inconsistency");
                writer->Add(
                    hit.Break(),
                    hit.Region(),
                    hit.Stream(),
                    TArrayRef<const char>(reinterpret_cast<const char*>(&value), sizeof(value))
                );
            }
            writer->FinishDoc();
        }
    }
    {
        NRTYServer::IIndexOwner::TGuardIndexModification g(context.Index);
        TArchiveOwner::Remove(root / ("index" + Suffix));
        NRTYArchive::HardLinkOrCopy(root / ("normalizedindex" + Suffix), root / ("index" + Suffix));
        TArchiveOwner::Remove(root / ("normalizedindex" + Suffix));

        TIndexMetadataProcessor meta(context.Dir.PathName());
        NRTYServer::TIndexMetadata::TAnnotationHeader* header = nullptr;
        for (auto& h : *meta->MutableAnnHeaders()) {
            if (h.GetName() == Suffix) {
                header = &h;
                break;
            }
        }
        if (!header) {
            header = meta->AddAnnHeaders();
        }
        header->SetName(Suffix);
        header->SetDataFormat(ToString(Format));
    }
}
