#include "parsed_entity.h"
#include "component.h"
#include "config.h"
#include "const.h"

#include <saas/rtyserver/factors/factors_config.h>
#include <saas/rtyserver/config/searcher_config.h>

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

TAnnParsedEntity::TAnnParsedEntity(TConstructParams& params)
    : TParsedDocument::TParsedEntity(params)
    , AnnStreamsInfo(params.Owner.GetConfig().GetSearcherConfig().Factors->GetAnnStreams())
{}

TAnnComponentParser::TAnnComponentParser(const TRTYServerConfig& config)
    : Config(config.ComponentsConfig.Get<TAnnComponentConfig>(NRTYServer::AnnComponentName))
{}

void TAnnComponentParser::Parse(TParsingContext& context) const {
    auto entity = context.Result.GetComponentEntity<TAnnParsedEntity>(NRTYServer::AnnComponentName);
    CHECK_WITH_LOG(entity);

    if (context.Document.HasAnnData()) {
        for (ui32 i = 0; i < context.Document.GetAnnData().SentencesSize(); ++i) {
            TAnnParsedEntity::TSentenceData sentData;
            const NRTYServer::TMessage::TSentenceData& protoData = context.Document.GetAnnData().GetSentences(i);
            if (!sentData.Parse(protoData, entity->GetAnnData(), Config->SkipUnknownStreams, Config->MapStreamIdToIndex)) {
                ythrow yexception() << "Can't parse ann data: " << protoData.DebugString();
            }
            entity->AddSentenceInfo(sentData);
        }
    }
}

bool TAnnParsedEntity::TSentenceData::Parse(const NRTYServer::TMessage::TSentenceData& protoData,
    const NRTYFactors::TAnnFactorsHashList& annStreamsData,
    bool skipUnknownStreams,
    bool mapStreamId)
{
    Sentence = UTF8ToWide(protoData.GetText());
    if (protoData.HasTextLanguage()) {
        TextLanguage = (ELanguage)(protoData.GetTextLanguage());
    }

    for (ui32 i = 0; i < protoData.StreamsByRegionSize(); ++i) {
        const auto& regionData = protoData.GetStreamsByRegion(i);
        TStreamsStorage streams;
        for (ui32 s = 0; s < regionData.StreamsSize(); ++s) {
            const NRTYServer::TMessage::TStreamData& info = regionData.GetStreams(s);
            ui32 streamNum = 0;
            ui32 streamId = 0;
            TString streamName = info.GetName();
            TString destination;
            const NRTYFactors::TSimpleFactorDescription* streamInfo = nullptr;
            for (const auto& streamsList : annStreamsData) {
                if (streamsList.second->CheckFactorName(streamName, streamNum)) {
                    destination = streamsList.first;
                    streamInfo = &streamsList.second->GetFactorsInfo()->GetFactor(streamNum);
                    streamId = streamsList.second->GetFactorsInfo()->GetFactor(streamNum).IndexGlobal;
                    break;
                }
            }
            if (!destination) {
                if (skipUnknownStreams)
                    continue;
                TString errorMessage = "unknown stream " + streamName + " for ann data";
                ERROR_LOG << errorMessage << Endl;
                ythrow yexception() << errorMessage;
            }
            TString data = Base64Decode(info.GetValue().GetValue());
            if (data.size() * 8 != streamInfo->Width) {
                TString errorMessage = "incorrect stream size " + ToString(data.size()) + " for " + streamName;
                ERROR_LOG << errorMessage << Endl;
                ythrow yexception() << errorMessage;
            }
            streams.push_back(TAnnParsedEntity::TStream(destination, data, mapStreamId ? streamNum : streamId));
            Destinations.insert(destination);
        }

        if (regionData.StreamsSize() == 0)
            return false;

        Regions[regionData.GetRegion()] = streams;
    }

    if (protoData.StreamsByRegionSize() == 0)
        return false;

    return true;
}
