#include "document.h"
#include "factors_erf.h"
#include "lables.h"
#include "text_utils.h"

#include <saas/util/types/cast.h>

#include <library/cpp/logger/global/global.h>
#include <library/cpp/charset/codepage.h>
#include <library/cpp/string_utils/base64/base64.h>
#include <util/generic/xrange.h>

using namespace NSaas;

NJson::TJsonValue TDocument::ToJson(TToJsonContext& context) const {
    if (!GetUrl())
        throw yexception() << "document must have a non-empty URL";
    if (HasBody() && SubZones.size())
        throw yexception() << "document cannot have both body and subzones";

    context.Encoding = CharsetByName(GetCharset());

    NJson::TJsonValue doc;
    for (const auto& a: Attributes)
        AddZoneValue(doc, a.first, a.second.ToJson(context));
    for (const auto& csBlock: CSBlocks)
        AddZoneValue(doc, csBlock->GetName(), csBlock->ToJson());
    for (const auto& qsBlock: QSBlocks)
        AddZoneValue(doc, qsBlock->GetName(), qsBlock->ToJson());
    for (const auto& qsBlock : IntQSBlocks)
        AddZoneValue(doc, qsBlock->GetName(), qsBlock->ToJson());

    for (const auto& embedding : InternalDocument.GetEmbeddings()) {
        NJson::TJsonValue value(NJson::JSON_MAP);
        value["value"] = Base64Encode(embedding.GetValue());
        value["version"] = embedding.GetVersion();
        if (embedding.HasTag()) {
            value["tag"] = embedding.GetTag();
        }

        NJson::TJsonValue element(NJson::JSON_MAP);
        element.InsertValue(NJsonFormat::TypeLable, "#E");
        element.InsertValue(NJsonFormat::ValueLable, std::move(value));
        AddZoneValue(doc, embedding.GetName(), element);
    }

    if (Annotations) {
        AddZoneValue(doc, NJsonFormat::AnnLable, Annotations->ToJson());
    } else if (InternalDocument.HasAnnData()) {
        AddZoneValue(doc, NJsonFormat::AnnLable, TAnnBlock(*InternalDocument.MutableAnnData()).ToJson());
    }

    if (GeoBlock) {
        for (ui32 i: xrange(GeoBlock->GetLayersCount())) {
            AddZoneValue(doc, GeoBlock->GetLayerName(i), GeoBlock->ToJson(i));
        }
    }

    NJson::TJsonValue& docOptions = doc.InsertValue(NJsonFormat::OptionsLable, NJson::JSON_MAP);
    if (!!GetMimeType())
        docOptions.InsertValue(NJsonFormat::MimeTypeLable, GetMimeType());
    if (!!GetLang())
        docOptions.InsertValue(NJsonFormat::LanguageLable, GetLang());
    if (!!GetLang2())
        docOptions.InsertValue(NJsonFormat::Language2Lable, GetLang2());
    if (!!GetLangDef())
        docOptions.InsertValue(NJsonFormat::LanguageDefaultLable, GetLangDef());
    if (!!GetLangDef2())
        docOptions.InsertValue(NJsonFormat::LanguageDefault2Lable, GetLangDef2());
    if (!!GetCharset())
        docOptions.InsertValue(NJsonFormat::CharsetLable, "utf-8");
    if (HasDeadlineMinutesUTC())
        docOptions.InsertValue(NJsonFormat::DeadlineLable, GetDeadlineMinutesUTC());
    if (HasDatabaseVersion())
        docOptions.InsertValue(NJsonFormat::DatabaseVersionLable, GetDatabaseVersion());
    if (HasVersion())
        docOptions.InsertValue(NJsonFormat::VersionLable, GetVersion());
    if (HasFilterRank())
        docOptions.InsertValue(NJsonFormat::FilterRankLable, GetFilterRank());
    if (!GetRealtime())
        docOptions.InsertValue(NJsonFormat::RealtimeLable, GetRealtime());
    if (HasTimestamp())
        docOptions.InsertValue(NJsonFormat::ModificationTimestampLable, GetTimestamp());
    if (GetCheckOnlyBeforeReply())
        docOptions.InsertValue(NJsonFormat::CheckOnlyBeforeReplyLable, GetCheckOnlyBeforeReply());
    if (HasUpdateType())
        docOptions.InsertValue(NJsonFormat::UpdateTypeLable, NRTYServer::TMessage::TUpdateType_Name(GetUpdateType()));

    doc.InsertValue(NJsonFormat::UrlLable, GetUrl());

    if (HasIndexedDoc())
        doc.InsertValue(NJsonFormat::IndexedDocLable, InsertText(Base64Encode(GetIndexedDoc().SerializeAsString()), context));

    if (SubZones.size()) {
        for (TZones::const_iterator i = SubZones.begin(); i != SubZones.end(); ++i) {
            TAbstractZone* subzone = i->second.Get();
            const TString& name = !subzone->GetName().empty() ? subzone->GetName() : NJsonFormat::NoNameZone;
            AddZoneValue(doc, name, subzone->ToJson(context));
        }
    } else if (HasBody()) {
        doc.InsertValue(NJsonFormat::BodyLable, InsertText(GetText(), context));
    }

    return doc;
}

TDocument::TDocument(NRTYServer::TMessage::TDocument& document)
    : InternalDocument(document)
{
    if (!InternalDocument.HasRealtime())
        SetRealtime(true);
}

void TDocument::OnAttribute(const TString& name, const TString& value, TAttributeValue::TAttributeValueType type) {
    NRTYServer::TAttribute* attribute = nullptr;
    NRTYServer::TMessage::TDocument::TProperty* property = nullptr;
    NRTYServer::TAttribute::TAttributeType attributeType = enum_cast<NRTYServer::TAttribute::TAttributeType>(type);
    switch(type) {
    case TAttributeValue::avtInt:
    case TAttributeValue::avtLit:
        if (SubZones.count(name))
            throw yexception() << "cannot add an attribute: document already has zone with name " << name;

        attribute = InternalDocument.AddSearchAttributes();
        break;
    case TAttributeValue::avtGrp:
    case TAttributeValue::avtGrpLit:
        attribute = InternalDocument.AddGroupAttributes();
        break;
    case TAttributeValue::avtFSProp:
        property = InternalDocument.AddFSProperties();
        Y_ASSERT(property);
        property->SetName(name);
        property->SetValue(value);
        return;
    case TAttributeValue::avtProp:
        property = InternalDocument.AddDocumentProperties();
        Y_ASSERT(property);
        property->SetName(name);
        property->SetValue(value);
        return;
    case TAttributeValue::avtSpecKey:
        property = InternalDocument.AddAdditionalKeys();
        Y_ASSERT(property);
        property->SetName(name);
        property->SetValue(value);
        return;
    case TAttributeValue::avtTIProp:
        property = InternalDocument.AddTrigramIndexedProps();
        Y_ASSERT(property);
        property->SetName(name);
        property->SetValue(value);
        return;
    case TAttributeValue::avtFactor:
        NSaas::AddSimpleFactor(name, value, * InternalDocument.MutableFactors());
        return;
    case TAttributeValue::avtIntFactor: {
        NSaas::AddSimpleIntFactor(name, value, *InternalDocument.MutableFactors());
        return;
    }
    }
    Y_ASSERT(attribute);
    attribute->SetName(name);
    attribute->SetValue(value);
    attribute->SetType(attributeType);
}

void TDocument::BuildJsonMessage(NJson::TJsonValue& result) const {
    TToJsonContext context(TToJsonContext::COMMON_JSON);
    result = ToJson(context);
}

TInstant NSaas::TDocument::GetDeadline() const {
    return TInstant::Minutes(GetDeadlineMinutesUTC());
}

TDocument& NSaas::TDocument::SetDeadline(TInstant value) {
    const ui32 minutes = ui32(value.Minutes());
    return SetDeadlineMinutesUTC(minutes);
}

TCSBlock& TDocument::AddCS(const TString& name) {
    CSBlocks.push_back(new TCSBlock(* InternalDocument.AddCSInfo()));
    CSBlocks.back()->SetName(name);
    return * CSBlocks.back();
}

TQSBlock& TDocument::AddQS(const TString& name) {
    QSBlocks.emplace_back(new TQSBlock(*InternalDocument.AddQSInfo()));
    QSBlocks.back()->SetName(name);
    return * QSBlocks.back();
}

NSaas::TIntQSBlock& NSaas::TDocument::AddIntQS(const TString& name) {
    IntQSBlocks.emplace_back(new TIntQSBlock(*InternalDocument.AddQSInfo()));
    IntQSBlocks.back()->SetName(name);
    return *IntQSBlocks.back();
}

TAnnBlock& TDocument::AddAnnotations() {
    if (!Annotations) {
        Annotations.Reset(new TAnnBlock(*InternalDocument.MutableAnnData()));
    }
    return *Annotations;
}

TGeoBlock& TDocument::AddGeo() {
    if (!GeoBlock) {
        GeoBlock = MakeHolder<TGeoBlock>(*InternalDocument.MutableGeoData());
    }
    return *GeoBlock;
}

TDocument& NSaas::TDocument::AddFactor(const TString& name, float value) {
    AddAttribute(name).AddValue(value).AddType(TAttributeValue::avtFactor);
    return *this;
}

TDocument& NSaas::TDocument::AddIntFactor(const TString& name, i32 value) {
    AddAttribute(name).AddValue(value).AddType(TAttributeValue::avtIntFactor);
    return *this;
}

TDocument& NSaas::TDocument::AddProperty(const TString& name, const TString& value) {
    AddAttribute(name).AddValue(value).AddType(TAttributeValue::avtProp);
    return *this;
}

TDocument& NSaas::TDocument::AddEmbedding(const TString& name, TStringBuf value, TStringBuf version, TStringBuf tag) {
    NRTYServer::TMessage::TEmbedding* embedding = InternalDocument.AddEmbeddings();
    Y_ASSERT(embedding);

    embedding->SetName(name);
    embedding->SetValue(value.data(), value.length());
    embedding->SetVersion(version.data(), version.length());
    if (tag.IsInited()) {
        embedding->SetTag(tag.data(), tag.length());
    }
    return *this;
}

TDocument& NSaas::TDocument::AddSpecialKey(const TString& name, const TString& value) {
    AddAttribute(name).AddValue(value).AddType(TAttributeValue::avtSpecKey);
    return *this;
}

TDocument& NSaas::TDocument::AddTrigramIndexingProp(const TString& name, const TString& value) {
    AddAttribute(name).AddValue(value).AddType(TAttributeValue::avtTIProp);
    return *this;
}

TDocument& NSaas::TDocument::AddBinaryProperty(const TString& name, const TString& value) {
    AddAttribute(name).AddValue(value).AddType(TAttributeValue::avtProp).SetBinary();
    return *this;
}

TDocument& NSaas::TDocument::AddFSProperty(const TString& name, const TString& value) {
    AddAttribute(name).AddValue(value).AddType(TAttributeValue::avtFSProp);
    return *this;
}

TDocument& NSaas::TDocument::AddGroupAttribute(const TString& name, i64 value) {
    AddAttribute(name).AddValue(value).AddType(TAttributeValue::avtGrp);
    return *this;
}

TDocument& NSaas::TDocument::AddGroupAttribute(const TString& name, const TString& value) {
    AddAttribute(name).AddValue(value).AddType(TAttributeValue::avtGrpLit);
    return *this;
}

TDocument& NSaas::TDocument::AddExtraTimestamp(ui32 stream, ui128 value) {
    auto timestamp = InternalDocument.AddTimestamps();
    timestamp->SetStream(stream);
    timestamp->SetValue(GetLow(value));
    if (GetHigh(value)) timestamp->SetValueEx(GetHigh(value));
    return *this;
}

TDocument& NSaas::TDocument::AddObjectContext(const TString& serialized) {
    CHECK_WITH_LOG(HasUrl());
    auto indexedDoc = InternalDocument.MutableIndexedDoc();
    indexedDoc->SetKiwiObject(serialized);
    indexedDoc->SetUrl(GetUrl());

    return *this;
}
