#include <util/memory/blob.h>
#include <util/string/strip.h>
#include <library/cpp/archive/yarchive.h>
#include <library/cpp/json/json_reader.h>

#include "threats-config.h"

namespace NWebmaster {

TThreatsConfig::TThreatsConfig() {
    static const unsigned char THREATS_CONFIG_DATA[] = {
        #include "threats.zlib"
    };

    static const unsigned char DISPLAYABLE_THREATS_DATA[] = {
        #include "displayable_threats.zlib"
    };

    TArchiveReader archive(TBlob::NoCopy(THREATS_CONFIG_DATA, Y_ARRAY_SIZE(THREATS_CONFIG_DATA)));
    for (size_t i = 0; i < archive.Count(); ++i) {
        const TString key = archive.KeyByIndex(i);
        const TBlob blob = archive.ObjectBlobByKey(key);
        AddThreatsJson(blob.Data(), blob.Length());
    }
    TArchiveReader dispThreatsArchive(TBlob::NoCopy(DISPLAYABLE_THREATS_DATA, Y_ARRAY_SIZE(DISPLAYABLE_THREATS_DATA)));
    for (size_t i = 0; i < dispThreatsArchive.Count(); ++i) {
        const TString key = dispThreatsArchive.KeyByIndex(i);
        const TBlob blob = dispThreatsArchive.ObjectBlobByKey(key);
        TMemoryInput in(blob.Data(), blob.Length());
        TString line;
        while (in.ReadLine(line)) {
            DisplayableThreats.insert(StripString(line));
        }
    }
    SaveThreats();
}

void TThreatsConfig::AddThreatsJson(const void* data, size_t size) {

    TMemoryInput in(data, size);
    NJson::TJsonValue configJson;
    NJson::ReadJsonTree(&in, &configJson, true);
    THashMap<TString, NJson::TJsonValue> threatsMap;
    configJson.GetMap(&threatsMap);
    for (const auto &entry : threatsMap) {
        const TString &threat = entry.first;

        TThreatSettings threatSettings;
        threatSettings.ThreatName = threat;
        THashMap<TString, NJson::TJsonValue> threatSettingsMap;
        entry.second.GetMap(&threatSettingsMap);

        THashMap<TString, NJson::TJsonValue>::const_iterator searchIt;
        // show_in_webmaster
        searchIt = threatSettingsMap.find("show_in_webmaster");
        if (searchIt != threatSettingsMap.end()) {
            threatSettings.Show = searchIt->second.GetBoolean();
            threatSettings.ShowSet = true;
        }
        // template
        searchIt = threatSettingsMap.find("template");
        if (searchIt != threatSettingsMap.end()) {
            const TString displayTemplate = searchIt->second.GetString();
            if (displayTemplate == "title_description") {
                threatSettings.DisplayMode = proto::NONE;
            } else if (displayTemplate == "title_description_urls") {
                threatSettings.DisplayMode = proto::URLS;
            } else if (displayTemplate == "title_description_urls_chains") {
                threatSettings.DisplayMode = proto::URLS_CHAINS;
            } else if (displayTemplate == "title_description_urls_no_date") {
                threatSettings.DisplayMode = proto::URLS_NO_DATE;
            } else {
                ythrow yexception() << "Unknown display template " << displayTemplate << " for threat " << threat;
            }
            threatSettings.DisplayModeSet = true;
        }
        // consequences
        searchIt = threatSettingsMap.find("consequences");
        if (searchIt != threatSettingsMap.end()) {
            threatSettings.Sanction = searchIt->second.GetString();
            threatSettings.SanctionSet = true;
        }
        ThreatSettings.push_back(threatSettings);
    }
}

void TThreatsConfig::SaveThreats() {
    // collect threat names and build trie
    TThreatsTrie::TBuilder threatsTrieBuilder;
    TVector <TString> threatNames;

    for (size_t i = 0; i < ThreatSettings.size(); ++i) {
        threatsTrieBuilder.Add(ThreatSettings[i].ThreatName + ".", i);
        threatNames.push_back(ThreatSettings[i].ThreatName);
    }

    TBufferOutput buffOut;
    threatsTrieBuilder.Save(buffOut);
    ThreatsTrieBuffer = buffOut.Buffer();
    ThreatsTrie = TThreatsTrie(ThreatsTrieBuffer.Data(), ThreatsTrieBuffer.Size());

    // sort threatNames by length
    std::sort(threatNames.begin(), threatNames.end(),
              [](const TString &a, const TString &b) { return a.size() < b.size(); });

    // inherit missing properties from parent
    for (const TString &threatName : threatNames) {
        // find threat
        TThreatSettings &threatSettings = ThreatSettings[ThreatsTrie.Get(threatName + ".")];
        // save display name
        if (DisplayableThreats.contains(threatName)) {
            threatSettings.DisplayName = threatName;
        }
        // remove last dot
        TString parentName = threatName;
        parentName.resize(threatName.size() - 1);
        while (!threatSettings.DisplayModeSet || !threatSettings.SanctionSet || !threatSettings.ShowSet ||
                threatSettings.DisplayName.Empty()) {
            size_t dotPos = parentName.find_last_of(".");
            if (dotPos == TString::npos) {
                // not all properties are set
                break;
            }
            parentName.resize(dotPos);
            // copy properties from parent
            size_t index;
            if (!ThreatsTrie.Find(parentName + ".", &index)) {
                // not found this parent
                continue;
            }
            // parent was found
            const TThreatSettings &parentSettings = ThreatSettings[index];
            // copy props
            if (!threatSettings.ShowSet && parentSettings.ShowSet) {
                threatSettings.Show = parentSettings.Show;
                threatSettings.ShowSet = true;
            }
            if (!threatSettings.DisplayModeSet && parentSettings.DisplayModeSet) {
                threatSettings.DisplayMode = parentSettings.DisplayMode;
                threatSettings.DisplayModeSet = true;
            }
            if (!threatSettings.SanctionSet && parentSettings.SanctionSet) {
                threatSettings.Sanction = parentSettings.Sanction;
                threatSettings.SanctionSet = true;
            }
            if (threatSettings.DisplayName.Empty() && !parentSettings.DisplayName.Empty()) {
                threatSettings.DisplayName = parentSettings.DisplayName;
            }
        }
    }
    // if some property is not set -> showInWebmaster = false
    for (TThreatSettings &threatSettings : ThreatSettings) {
        if (!threatSettings.SanctionSet || !threatSettings.DisplayModeSet || threatSettings.DisplayName.empty()) {
            if (threatSettings.Show) {
                InvalidThreatSettings.push_back(threatSettings);
                threatSettings.Show = false;
            }
        }
    }
}

bool TThreatsConfig::FindMatchingThreat(const TString &threatName, TThreatSettings *matchingThreat) const {
    size_t prefixLen;
    size_t idx;
    bool found = ThreatsTrie.FindLongestPrefix(threatName + ".", &prefixLen, &idx);
    if (found) {
        *matchingThreat = ThreatSettings[idx];
    }
    return found && matchingThreat->Show;
}

bool TThreatsConfig::FindExactThreat(const TString &threatName, TThreatSettings *threat) const {
    size_t idx;
    bool found = ThreatsTrie.Find(threatName + ".", &idx);
    if (found) {
        *threat = ThreatSettings[idx];
    }
    return found && threat->Show;
}

size_t TThreatsConfig::Size() const {
    return ThreatsTrie.Size();
}

} //namespace NWebmaster
