#include <robot/jupiter/protos/acceptance.pb.h>
#include <robot/jupiter/protos/export.pb.h>
#include <robot/library/yt/static/command.h>
#include <robot/library/yt/static/table.h>

#include <util/generic/size_literals.h>
#include <util/string/vector.h>

#include <wmconsole/version3/processors/antiall/antiall.pb.h>
#include <wmconsole/version3/processors/antiall/threats/threats-config.h>
#include <wmconsole/version3/wmcutil/log.h>
#include <wmconsole/version3/wmcutil/owners.h>
#include <wmconsole/version3/wmcutil/thread.h>
#include <wmconsole/version3/wmcutil/yt/misc.h>

#include <yweb/antispam/achievements/https_achievement/protos/https_achievement.pb.h>

#include <limits>

#include "config.h"
#include "threats/threats-config.h"
#include "../../wmcutil/mr/mr_tasks.h"

namespace NWebmaster {

using namespace NJupiter;

TInputTag<NProto::TWebmasterHostRecord> WebmasterHostsInputTag               (0);
TInputTag<NProto::TWebmasterHostsByOwner> WebmasterHostsByOwnerInputTag      (1);
TInputTag<NAntispam::TAntivirusExportItem> AntispamSourceInputTag            (2);
TInputTag<NAntispam::TAntivirusExportItem> AntivirSourceInputTag             (3);
TInputTag<NProto::THostWithThreats> HostWithThreatsInputTag                  (4);
TInputTag<NProto::THostWithThreats> PrevHostWithThreatsInputTag              (5);

TOutputTag<NProto::TWebmasterHostsByOwner> WebmasterHostsByOwnerOutputTag    (0);
TOutputTag<NAntispam::TAntivirusExportItem> AntivirSourceOutputTag           (1);
TOutputTag<NProto::THostWithThreats> HostWithThreatsOutputTag                (2);
TOutputTag<NProto::THostWithThreats> HostWithThreatsDiffOutputTag            (3);

const size_t MAX_SAMPLES_FOR_VERDICT = 10;

struct TThreatComparator {
    bool operator()(const proto::ThreatInfo &t1, const proto::ThreatInfo &t2) {
        return t1.threat() < t2.threat();
    }
};

struct THostOwnerMapper : public TTaggedMapper {
public:
    THostOwnerMapper() = default;

    void DoTagged(TTagedReader reader, TTagedWriter writer) override final {
        TString owner;
        TString error;
        TAnsipamOwnerCanonizer ownerCanonizer;
        NProto::TWebmasterHostsByOwner out;
        for (; reader.IsValid(); reader.Next()) {
            NProto::TWebmasterHostRecord record = reader.GetRow(WebmasterHostsInputTag);
            if (ownerCanonizer.GetOwner(record.GetHost(), owner, error, true)) {
                out.SetOwner(owner);
                out.SetHosts(record.GetHost());
                writer.AddRow(out, WebmasterHostsByOwnerOutputTag);
            }
        }
    }
};
REGISTER_MAPPER(THostOwnerMapper)

struct THostOwnerReducer : public TTaggedReducer {
public:
    THostOwnerReducer() = default;

    void DoTagged(TTagedReader reader, TTagedWriter writer) override final {
        TString owner = reader.GetRow(WebmasterHostsByOwnerInputTag).GetOwner();
        TStringBuilder hosts;
        for (; reader.IsValid(); reader.Next()) {
            if (!hosts.Empty()) {
                hosts << ";";
            }
            hosts << reader.GetRow(WebmasterHostsByOwnerInputTag).GetHosts();
        }
        NProto::TWebmasterHostsByOwner result;
        result.SetOwner(owner);
        result.SetHosts(hosts);
        writer.AddRow(result, WebmasterHostsByOwnerOutputTag);
    }
};
REGISTER_REDUCER(THostOwnerReducer)

struct TUrlSampleInfo {
public:
    TUrlSampleInfo(const TVector<TString> &chain, time_t timestamp)
        : Chain(chain)
        , Timestamp(timestamp) {
    }

public:
    TVector<TString> Chain;
    time_t Timestamp;
};

class TSamplesTreeNode {
    void AddSampleForNode(const TUrlSampleInfo &sampleInfo) {
        SamplesCount++;
        Samples.push_back(sampleInfo);
    }

public:
    TSamplesTreeNode() = default;

    void AddSample(const TString &url, const TUrlSampleInfo &sampleInfo, size_t offset = 0) {
        SamplesCount++;
        size_t nextSlash = url.find('/', offset);
        if (nextSlash == TString::npos) {
            Children[url.substr(offset)].AddSampleForNode(sampleInfo);
        } else {
            Children[url.substr(offset, nextSlash - offset)].AddSample(url, sampleInfo, nextSlash + 1);
        }
    }

    size_t ThinOutSamples(size_t leaveSamples) {
        if (SamplesCount <= leaveSamples) {
            return SamplesCount;
        }
        if (Samples.size() > 0) {
            if (Children.size() > 0) {
                SamplesCount -= Samples.size();
                Samples.clear();
            } else {
                Samples.crop(leaveSamples);
                SamplesCount = Samples.size();
                return SamplesCount;
            }
        }
        if (SamplesCount <= leaveSamples) {
            return SamplesCount;
        }
        size_t newSize = Samples.size();
        if (Children.size() >= leaveSamples) {
            TMap<TString, TSamplesTreeNode>::iterator childrenIt = Children.begin();
            while (childrenIt != Children.end() && newSize < leaveSamples) {
                newSize += childrenIt->second.ThinOutSamples(1);
                childrenIt++;
            }
            if (childrenIt != Children.end()) {
                Children.erase(childrenIt, Children.end());
            }
        } else {
            TVector<TSamplesTreeNode *> childrenVector;
            for (auto &childEntry : Children) {
                childrenVector.push_back(&childEntry.second);
            }
            std::sort(childrenVector.begin(), childrenVector.end(),
                      [](const TSamplesTreeNode *l, const TSamplesTreeNode *r) -> bool {
                          return l->SamplesCount < r->SamplesCount;
                      });
            TVector<TSamplesTreeNode *>::iterator it = childrenVector.begin();
            for (; it != childrenVector.end(); ++it) {
                TSamplesTreeNode *node = *it;
                auto restBranches = std::distance(it, childrenVector.end());
                size_t allowedPerBranch = (leaveSamples - newSize) / restBranches;
                if (node->SamplesCount > allowedPerBranch) {
                    node->ThinOutSamples(allowedPerBranch);
                }
                newSize += node->SamplesCount;
            }
        }
        SamplesCount = newSize;
        return SamplesCount;
    }

    void Clear() {
        SamplesCount = 0;
        Children.clear();
        Samples.clear();
    }

public:
    size_t SamplesCount = 0;
    TMap<TString, TSamplesTreeNode> Children;
    TVector<TUrlSampleInfo> Samples; // should be able to process same url with different properties
};

struct TSamplesAggregator {
private:
    void FillSamples(proto::VerdictInfo &verdict, const TSamplesTreeNode &node, const TString &urlPrefix) {
        for (const TUrlSampleInfo &sampleInfo : node.Samples) {
            proto::UrlSample *protoSample = verdict.add_samples();
            protoSample->set_url(urlPrefix);
            protoSample->set_update_timestamp(sampleInfo.Timestamp);
            for (const TString &chainItem : sampleInfo.Chain) {
                protoSample->add_infection_chain(chainItem);
            }
        }
        for (const auto &childEntry : node.Children) {
            const TString url = TStringBuilder() << urlPrefix << "/" << childEntry.first;
            FillSamples(verdict, childEntry.second, url);
        }
    }

public:
    void AddSample(const TString &url, const TUrlSampleInfo &info) {
        SamplesTree.AddSample(url, info);
    }

    void FillSamples(proto::VerdictInfo &verdict) {
        SamplesTree.ThinOutSamples(MAX_SAMPLES_FOR_VERDICT);
        FillSamples(verdict, SamplesTree, "");
    }

    void Reset() {
        SamplesTree.Clear();
    }

public:
    TSamplesTreeNode SamplesTree;
};

struct TAntispamExportMapper : public TTaggedMapper {
    TAntispamExportMapper() = default;

    void DoTagged(TTagedReader reader, TTagedWriter writer) override final {
        for (; reader.IsValid(); reader.Next()) {
            NAntispam::TAntivirusExportItem record = reader.GetRow(AntispamSourceInputTag);
            TThreatSettings threatSettings;
            if (!ThreatsConfig.FindMatchingThreat(record.GetThreat(), &threatSettings)) {
                // not supported threat
                continue;
            }

            const TString &host = record.GetHost();
            const TString &threatName = threatSettings.ThreatName;

            NProto::THostWithThreats hostWithThreats;
            proto::ThreatInfo *threat = hostWithThreats.MutableThreats()->add_threats();
            threat->set_threat(threatName);
            threat->set_display_name(threatSettings.DisplayName);
            threat->set_display_mode(threatSettings.DisplayMode);
            threat->set_threat_source_host(host);
            threat->set_lock_mode(proto::TIMED);
            threat->set_propagation_mode(proto::OWNER);
            threat->set_sanction(threatSettings.Sanction);
            if (record.HasLastRecheck()) {
                threat->set_last_recheck(record.GetLastRecheck());
            }
            if (record.HasNextRecheck()) {
                threat->set_next_recheck(record.GetNextRecheck());
            }
            if (record.HasNextInterval()) {
                threat->set_recheck_interval_s(record.GetNextInterval());
            }
            threat->set_allowed_to_recheck(false);
            TString owner, error;
            OwnerCanonizer.GetOwner(host, owner, error, true);

            hostWithThreats.SetOwner(owner);
            hostWithThreats.SetHost(host);
            writer.AddRow(hostWithThreats, HostWithThreatsOutputTag);
        }
    }

private:
    TAnsipamOwnerCanonizer OwnerCanonizer;
    TThreatsConfig ThreatsConfig = TThreatsConfig::CInstance();
};

REGISTER_MAPPER(TAntispamExportMapper)

struct TAntivirExportReducer : public TTaggedReducer {
public:
    TAntivirExportReducer() = default;

    void DoTagged(TTagedReader reader, TTagedWriter writer) override final {
        NAntispam::TAntivirusExportItem record = reader.GetRow(AntivirSourceInputTag);
        TThreatSettings threatSettings;
        if (!ThreatsConfig.FindMatchingThreat(record.GetThreat(), &threatSettings)) {
            // not supported threat
            return;
        }

        const bool needUrls = threatSettings.DisplayMode == proto::URLS ||
            threatSettings.DisplayMode == proto::URLS_CHAINS ||
            threatSettings.DisplayMode == proto::URLS_NO_DATE;
        const bool needChains = threatSettings.DisplayMode == proto::URLS_CHAINS;
        const bool needDate = threatSettings.DisplayMode != proto::URLS_NO_DATE;

        const TString host = record.GetHost();
        TString owner, error;
        OwnerCanonizer.GetOwner(host, owner, error, true);

        NProto::THostWithThreats hostWithThreats;
        hostWithThreats.SetOwner(owner);
        hostWithThreats.SetHost(host);
        proto::ThreatInfo *threatInfo = hostWithThreats.MutableThreats()->add_threats();
        threatInfo->set_threat(threatSettings.ThreatName);
        threatInfo->set_display_name(threatSettings.DisplayName);
        threatInfo->set_sanction(threatSettings.Sanction);
        threatInfo->set_display_mode(threatSettings.DisplayMode);
        threatInfo->set_threat_source_host(host);
        threatInfo->set_propagation_mode(proto::SUBDOMAINS);
        time_t lastUpdate = 0;

        if (record.HasNextInterval()) {
            threatInfo->set_lock_mode(proto::TIMED);
            threatInfo->set_recheck_interval_s(record.GetNextInterval());
            if (record.HasLastRecheck()) {
                threatInfo->set_last_recheck(record.GetLastRecheck());
            }
            if (record.HasNextRecheck()) {
                threatInfo->set_next_recheck(record.GetNextRecheck());
            }
        } else {
            threatInfo->set_lock_mode(proto::UNTIL_UPDATE);
        }

        TSamplesAggregator samplesAggregator;
        while (reader.IsValid()) {
            proto::VerdictInfo *verdictInfo = threatInfo->add_verdicts();
            const TString verdict = record.GetVerdict();
            verdictInfo->set_verdict(verdict);

            while (reader.IsValid() && record.GetVerdict() == verdict) {
                time_t downloadTime = TInstant::ParseIso8601Deprecated(record.GetDownloadTime()).TimeT();
                if (downloadTime > lastUpdate) {
                    lastUpdate = downloadTime;
                }

                TVector<TString> chain;
                if (needChains) {
                    TString commaSeparatedChain = record.GetChain();
                    if (!commaSeparatedChain.empty()) {
                        Split(commaSeparatedChain, ",", chain);
                    }
                }

                if (needUrls) {
                    TUrlSampleInfo sampleInfo(chain, needDate ? downloadTime : 0L);
                    samplesAggregator.AddSample(record.GetPath(), sampleInfo);
                }
                reader.Next();
                if (reader.IsValid()) {
                    record = reader.GetRow(AntivirSourceInputTag);
                }
            }
            samplesAggregator.FillSamples(*verdictInfo);
            samplesAggregator.Reset();
        }
        threatInfo->set_confirmation_date(lastUpdate);
        threatInfo->set_allowed_to_recheck(false);

        writer.AddRow(hostWithThreats, HostWithThreatsOutputTag);
    }
private:
    TAnsipamOwnerCanonizer OwnerCanonizer;
    TThreatsConfig ThreatsConfig = TThreatsConfig::CInstance();
};
REGISTER_REDUCER(TAntivirExportReducer)

struct THostWithThreatsPropagator : public TTaggedReducer {
public:
    THostWithThreatsPropagator() = default;

    void DoTagged(TTagedReader reader, TTagedWriter writer) override final {
        TMaybe<NProto::TWebmasterHostsByOwner> maybeHosts = reader.GetSingleRowMaybe(WebmasterHostsByOwnerInputTag);
        if (!maybeHosts.Defined()) {
            return;
        }
        TVector<TString> webmasterHosts;
        Split(maybeHosts->GetHosts(), ";", webmasterHosts);
        // propagate threats to all hosts
        for (; reader.IsValid(); reader.Next()) {
            NProto::THostWithThreats hostWithThreats = reader.GetRow(HostWithThreatsInputTag);
            const TString &threatHost = hostWithThreats.GetHost();
            for (const proto::ThreatInfo &threatInfo : hostWithThreats.GetThreats().threats()) {
                for (const TString &webmasterHost : webmasterHosts) {
                    if (threatInfo.propagation_mode() == proto::SUBDOMAINS && !NUtils::IsSubdomain(webmasterHost, threatHost)) {
                        continue;
                    }
                    bool allowedToRecheck = NUtils::RemoveWWW(NUtils::RemoveScheme(webmasterHost)) == threatHost;
                    NProto::THostWithThreats result;
                    result.SetHost(webmasterHost);
                    result.SetOwner(hostWithThreats.GetOwner());
                    result.SetFixed(hostWithThreats.GetFixed());
                    result.SetActualSince(hostWithThreats.GetActualSince());
                    proto::ThreatInfo *newThreatInfo = result.MutableThreats()->add_threats();
                    *newThreatInfo = threatInfo;
                    newThreatInfo->set_allowed_to_recheck(allowedToRecheck);
                    writer.AddRow(result, HostWithThreatsOutputTag);
                }
            }
        }
    }
};
REGISTER_REDUCER(THostWithThreatsPropagator)

struct THostWithThreatsReducer : public TTaggedReducer {
public:
    THostWithThreatsReducer() = default;

    void DoTagged(TTagedReader reader, TTagedWriter writer) override final {
        TMaybe<NProto::THostWithThreats> previous = reader.GetSingleRowMaybe(PrevHostWithThreatsInputTag);
        bool hasCurrent = false;
        NProto::THostWithThreats current;
        NProto::THostWithThreats hostWithThreats;
        for (; reader.IsValid(); reader.Next()) {
            hasCurrent = true;
            hostWithThreats = reader.GetRow(HostWithThreatsInputTag);
            for (const proto::ThreatInfo &threatInfo : hostWithThreats.GetThreats().threats()) {
                *current.MutableThreats()->add_threats() = threatInfo;
            }
        }
        // snapshot
        if (hasCurrent) {
            current.SetOwner(hostWithThreats.GetOwner());
            current.SetHost(hostWithThreats.GetHost());
            // sort
            auto *threatsCollection = current.MutableThreats()->mutable_threats();
            std::sort(threatsCollection->begin(), threatsCollection->end(), TThreatComparator());
            // actual since
            current.SetActualSince(previous.Defined() ? previous->GetActualSince() : Now().MilliSeconds());
            writer.AddRow(current, HostWithThreatsOutputTag);
        }
        // diff
        if (writer.HasOutputTag(HostWithThreatsDiffOutputTag)) {
            if (hasCurrent && !previous.Defined()) {
                current.SetFixed(false);
                writer.AddRow(current, HostWithThreatsDiffOutputTag);
            } else if (!hasCurrent && previous.Defined()) {
                previous->SetFixed(true);
                writer.AddRow(*previous, HostWithThreatsDiffOutputTag);
            }
        }
    }
};
REGISTER_REDUCER(THostWithThreatsReducer)

int BuildSnapshotAndDiff(int , const char **) {
    const TConfig &config = TConfig::CInstance();
    NYT::IClientPtr client = NYT::CreateClient(config.MR_SERVER_HOST);
    // check for source dates
    time_t antispamSourceDate = NYTUtils::GetModificationTime(client, config.TABLE_ANTISPAM_SOURCE);
    time_t antivirSourceDate = NYTUtils::GetModificationTime(client, config.TABLE_ANTIVIR_SOURCE);
    time_t antispamLastProcessed = 0;
    time_t antivirLastProcessed = 0;
    try {
        antispamLastProcessed = NYTUtils::GetAttr(client, config.TABLE_SNAPSHOT, config.ATTR_ANTISPAM_SOURCE_DATE).AsInt64();
        antivirLastProcessed = NYTUtils::GetAttr(client, config.TABLE_SNAPSHOT, config.ATTR_ANTIVIR_SOURCE_DATE).AsInt64();
    } catch (yexception& ignored) {
    }
    if (antispamLastProcessed >= antispamSourceDate && antivirLastProcessed >= antivirSourceDate) {
        LOG_INFO("No fresh source tables. Exiting");
        return 0;
    }
    LOG_INFO("Found fresh tables. antispam: %li, antivir: %li", antispamSourceDate, antivirSourceDate);
    NYT::ITransactionPtr tx = client->StartTransaction();
    // group webmaster hosts by owner
    TTable<NProto::TWebmasterHostsByOwner> hostsByOwner(tx, NYTUtils::JoinPath(config.TABLE_ANTIALL_TMP_ROOT, "hosts-by-owner"));
    TTable<NAntispam::TAntivirusExportItem> antivirSourceTable(tx, config.TABLE_ANTIVIR_SOURCE);
    TTable<NProto::THostWithThreats> antivirProcessedTable(tx, config.TABLE_ANTIVIR_PROCESSED);
    TTable<NAntispam::TAntivirusExportItem> antispamSourceTable(tx, config.TABLE_ANTISPAM_SOURCE);
    TTable<NProto::THostWithThreats> antispamProcessedTable(tx, config.TABLE_ANTISPAM_PROCESSED);
    TTable<NProto::THostWithThreats> snapshotTable(tx, config.TABLE_SNAPSHOT);

    if (snapshotTable.Exists()) {
        LOG_INFO("Replacing threats snapshot");
        tx->Move(config.TABLE_SNAPSHOT, config.TABLE_PREV_SNAPSHOT, NYT::TMoveOptions().Force(true));
    }

    TMaybe<yexception> error;
    auto prepareWebmasterHosts = [&] {
        try {
            LOG_INFO("Prepare webmaster hosts started");
            TMapReduceCmd<THostOwnerMapper, THostOwnerReducer>(tx)
                .Input(TTable<NProto::TWebmasterHostRecord>(tx, config.TABLE_WEBMASTER_HOSTS), WebmasterHostsInputTag)
                .IntermediateMapTag(WebmasterHostsByOwnerOutputTag)
                .IntermediateReduceTag(WebmasterHostsByOwnerInputTag)
                .Output(hostsByOwner, WebmasterHostsByOwnerOutputTag)
                .ReduceBy("Owner")
                .Do();
            TSortCmd<NProto::TWebmasterHostsByOwner>(tx, hostsByOwner)
                .By("Owner")
                .Do();
            LOG_INFO("Prepare webmaster hosts finished");
        } catch (yexception &e) {
            LOG_ERROR("Prepare webmaster hosts failed: %s", e.what());
            error = e;
        };
    };

    auto prepareAntispamExport = [&] {
        try {
            if (antispamLastProcessed < antispamSourceDate) {
                LOG_INFO("Preparing antispam source table started");
                TMapCmd<TAntispamExportMapper>(tx)
                    .Input(antispamSourceTable, AntispamSourceInputTag)
                    .Output(antispamProcessedTable, HostWithThreatsOutputTag)
                    .Do();
                TSortCmd<NProto::THostWithThreats>(tx, antispamProcessedTable)
                    .By("Host")
                    .Do();
                LOG_INFO("Preparing antispam source table finished");
            }
        } catch (yexception &e) {
            LOG_ERROR("Preparing antispam source table failed: %s", e.what());
            error = e;
        };
    };

    auto prepareAntivirExport = [&] {
        try {
            if (antivirLastProcessed < antivirSourceDate) {
                LOG_INFO("Preparing antivir source table started");
                TMapReduceCmd<void, TAntivirExportReducer>(tx)
                    .Input(antivirSourceTable, AntivirSourceInputTag)
                    .IntermediateMapTag(AntivirSourceOutputTag)
                    .IntermediateReduceTag(AntivirSourceInputTag)
                    .Output(antivirProcessedTable, HostWithThreatsOutputTag)
                    .ReduceBy({"host", "threat"})
                    .SortBy({"host", "threat", "verdict", "path"})
                    .Do();
                TSortCmd<NProto::THostWithThreats>(tx, antivirProcessedTable)
                    .By("Host")
                    .Do();
                LOG_INFO("Preparing antivir source table finished");
            }
        } catch (yexception &e) {
            LOG_ERROR("Preparing antivir source table failed: %s", e.what());
            error = e;
        };
    };

    NUtils::RunAsync(prepareWebmasterHosts, prepareAntispamExport, prepareAntivirExport);
    if (error.Defined()) {
        throw error.Get();
    }

    LOG_INFO("Building full snapshot and calculating diff");
    TTable<NProto::THostWithThreats> prevSnapshotTable(tx, config.TABLE_PREV_SNAPSHOT);
    const TString timestamp = ToString(Now().MilliSeconds());
    TTable<NProto::THostWithThreats> snapshotDiffTable(tx, NYTUtils::JoinPath(config.TABLE_CHANGES_ROOT, timestamp));

    LOG_INFO("Compressing threats and calculating diff");
    TReduceCmd<THostWithThreatsReducer>(tx)
        .Input(prevSnapshotTable.IfExists(), PrevHostWithThreatsInputTag)
        .Input(antivirProcessedTable, HostWithThreatsInputTag)
        .Input(antispamProcessedTable, HostWithThreatsInputTag)
        .Output(snapshotTable, HostWithThreatsOutputTag)
        .Output(prevSnapshotTable.Exists() ? MakeMaybe(snapshotDiffTable) : Nothing(), HostWithThreatsDiffOutputTag)
        .ReduceBy("Host")
        .Do();

    TSortCmd<NProto::THostWithThreats>(tx, snapshotTable)
        .By("Host")
        .Do();

    if (prevSnapshotTable.Exists()) { // has diff
        LOG_INFO("Propagating changes to real host ids");
        TSortCmd<NProto::THostWithThreats>(tx, snapshotDiffTable)
            .By("Owner")
            .Do();

        TTable<NProto::THostWithThreats> nfTable(tx, NYTUtils::JoinPath(config.TABLE_NOTIFICATIONS_ROOT, timestamp));

        TReduceCmd<THostWithThreatsPropagator>(tx)
            .Input(hostsByOwner, WebmasterHostsByOwnerInputTag)
            .Input(snapshotDiffTable, HostWithThreatsInputTag)
            .Output(nfTable, HostWithThreatsOutputTag)
            .ReduceBy("Owner")
            .Do();

        TSortCmd<NProto::THostWithThreats>(tx, nfTable)
            .By("Host")
            .Do();
    }

    NYTUtils::SetAttr(tx, config.TABLE_SNAPSHOT, config.ATTR_ANTISPAM_SOURCE_DATE, antispamSourceDate);
    NYTUtils::SetAttr(tx, config.TABLE_SNAPSHOT, config.ATTR_ANTIVIR_SOURCE_DATE, antivirSourceDate);

    tx->Commit();
    LOG_INFO("Done");
    return 0;
}

}
