#include <search/config/virthost.h>
#include <saas/library/metasearch/helpers/attribute.h>
#include <saas/library/sharding/sharding.h>
#include <saas/util/types/interval.h>
#include <saas/searchproxy/common/cgi.h>
#include <search/web/core/rule.h>

namespace {
    const TString DefaultUrlAttribute = "url";
    const TString DefaultUrlParameter = NSearchProxyCgi::text;
    const TString DefaultTemplateParameter = NSearchProxyCgi::tmplate;
    const TString DefaultUserRequestParameter = "user_request";
}

enum class EMode {
    KeyValue,
    Search
};

template <>
void Out<EMode>(IOutputStream& output, TTypeTraits<EMode>::TFuncParam value) {
    switch (value) {
    case EMode::KeyValue:
        output << "KeyValue";
        break;
    case EMode::Search:
        output << "Search";
        break;
    }
}

template <>
EMode FromStringImpl<EMode>(const char* data, size_t length) {
    TStringBuf s(data, length);
    if (s == "KeyValue") {
        return EMode::KeyValue;
    }
    if (s == "Search") {
        return EMode::Search;
    }
    throw yexception() << "unknown XSharding mode: " << s;
}

class TXShardingRuleContext : public IRearrangeRuleContext {
private:
    const NSaas::TShardsDispatcher::TPtr Sharding;
    const EMode Mode;
    TString KeyPrefixParameter;
    const TString UrlParameter;
    const TString TemplateParameter;
    const TString UserRequestParameter;

    TString Attribute;
    TString Service;
    TSet<TString> Urls;
    TSet<NSaas::TKeyPrefix> KeyPrefixes;

    bool Debug = false;

public:
    TXShardingRuleContext(NSaas::TShardsDispatcher::TPtr sharding, EMode mode, const TString& keyPrefixParameter)
        : Sharding(sharding)
        , Mode(mode)
        , KeyPrefixParameter(keyPrefixParameter)
        , UrlParameter(DefaultUrlParameter)
        , TemplateParameter(DefaultTemplateParameter)
        , UserRequestParameter(DefaultUserRequestParameter)
    {
    }

protected:
    void DoInitRuleParams(const TAdjustRuleParams& arp) override {
        Attribute = LocalScheme()["Attribute"].GetString(DefaultUrlAttribute);
        Debug = LocalScheme()["Debug"].IsTrue();
        Service = arp.RequestFields.CgiParam.Get("service");
        if (arp.RequestFields.CgiParam.Has("sgkps")) {
            KeyPrefixParameter = "sgkps";
        } else {
            KeyPrefixParameter = NSearchProxyCgi::kps;
        }

        if (Mode == EMode::KeyValue) {
            ProcessUrlParameter(arp.RequestFields.CgiParam, UrlParameter);
        } else if (Mode == EMode::Search) {
            ProcessUrlParameter(arp.RequestFields.CgiParam, TemplateParameter);
            ProcessUrlParameter(arp.RequestFields.CgiParam, UserRequestParameter);
        }

        if (Urls.empty()) {
            Urls.insert(TString());
        }

        ProcessPrefixParameter(arp.RequestFields.CgiParam, KeyPrefixParameter);
        ProcessPrefixParameter(arp.RequestFields.CgiParam, Service + "_" + KeyPrefixParameter);

        if (KeyPrefixes.empty()) {
            KeyPrefixes.insert(0);
        }
    }

    void DoAdjustClientParams(const TAdjustParams& adjustParams) override {
        if (adjustParams.ClientRequestAdjuster->IsClientEphemeral()) {
            return;
        }

        const TInterval<NSearchMapParser::TShardIndex>& interval = FromString(adjustParams.SourceGroup());
        bool eligible = false;

        TSet<TString> localUrls;
        TSet<NSaas::TKeyPrefix> localKeyPrefixes;
        for (auto&& url : Urls) {
            for (auto&& kp : KeyPrefixes) {
                if (Sharding->CheckSearchInterval(url, kp, interval)) {
                    eligible = true;
                    if (kp) {
                        localKeyPrefixes.insert(kp);
                    }
                    if (url) {
                        localUrls.insert(url);
                    }
                }
            }
        }

        if (!eligible) {
            adjustParams.ClientRequestAdjuster->ClientDontSendRequest();

        } else {
            if (localKeyPrefixes.size() != KeyPrefixes.size()) {
                for (int i = 0; i < adjustParams.ClientRequestAdjuster->ClientFormFieldCount(KeyPrefixParameter.data()); ++i) {
                    adjustParams.ClientRequestAdjuster->ClientFormFieldRemove(KeyPrefixParameter.data(), i);
                }
                if (!localKeyPrefixes.empty()) {
                    const TString& stringKeyPrefixes = JoinStrings(localKeyPrefixes.begin(), localKeyPrefixes.end(), ",");
                    adjustParams.ClientRequestAdjuster->ClientFormFieldInsert(KeyPrefixParameter.data(), stringKeyPrefixes.data());
                }
            }
        }
    }

    void DoRearrangeAfterMerge(TRearrangeParams& rp) override {
        if (Debug) {
            rp.InsertRearrangeResult("Mode", ToString(Mode));
            rp.InsertRearrangeResult("Urls", JoinStrings(Urls.begin(), Urls.end(), ","));
            rp.InsertRearrangeResult("KeyPrefixes", JoinStrings(KeyPrefixes.begin(), KeyPrefixes.end(), ","));
        }
    }

private:

    void ProcessUrlParameter(const TCgiParameters& cgiParam, const TString& paramName) {
        auto range = cgiParam.equal_range(paramName);
        for (auto i = range.first; i != range.second; ++i) {
            AddUrls(i->second);
        }
    }

    void AddUrls(const TString& text) {
        switch (Mode) {
        case EMode::KeyValue:
            return AddKeys(text);
        case EMode::Search:
            return AddSearchUrls(text);
        default:
            ythrow yexception() << "unknown mode " << int(Mode);
        }
    }

    void AddSearchUrls(const TString& text) {
        const TStringBuf value = NUtil::GetAttributeValue(Attribute, text);
        if (!value.empty()) {
            Urls.insert(TString(value));
        }
    }

    void AddKeys(const TString& text) {
        for (const TString& key: SplitString(text, ",")) {
            Urls.insert(key);
        }
    }

    void ProcessPrefixParameter(const TCgiParameters& cgiParam, const TString& paramName) {
        auto range = cgiParam.equal_range(paramName);
        for (auto i = range.first; i != range.second; ++i) {
            AddKeyPrefixes(i->second);
        }
    }

    void AddKeyPrefixes(const TString& value) {
        TStringBuf valueBuf = value;
        TStringBuf s;
        while (valueBuf.NextTok(',', s)) {
            if (s.empty()) {
                continue;
            }
            KeyPrefixes.insert(FromString<NSaas::TKeyPrefix>(s));
        }
    }
};

class TXShardingRule : public IRearrangeRule {
private:
    using TShardingContext = NSaas::TShardsDispatcher::TContext;
public:
    TXShardingRule(const NRearr::TRuleOptions& options, const TString& protocol)
        : Sharding(MakeAtomicShared<NSaas::TShardsDispatcher>(
            options.GetScheme().Has("ShardBy") ? TShardingContext::FromString(options.GetScheme()["ShardBy"].ForceString()) : TShardingContext(NSaas::UrlHash)
        ))
        , Mode(options.GetScheme().Has("Mode") ? FromString<EMode>(options.GetScheme()["Mode"].ForceString()) : EMode::KeyValue)
        , Protocol(protocol)
    {
    }
private:
    IRearrangeRuleContext* DoConstructContext() const override {
        const TString& keyPrefixParameter = (Protocol == "proto") ? "kps" : "prefix";
        return MakeHolder<TXShardingRuleContext>(Sharding, Mode, keyPrefixParameter).Release();
    }
private:
    const NSaas::TShardsDispatcher::TPtr Sharding;
    const EMode Mode;
    const TString Protocol;
};

IRearrangeRule* CreateXShardingRule(const NRearr::TRuleOptions& options, NRearr::TDataRegistry&, const TConstructRearrangeEnv& env) {
    return MakeHolder<TXShardingRule>(options, env.Config.CommonSourceOptions.ProtocolType).Release();
}

REGISTER_REARRANGE_RULE(XSharding, CreateXShardingRule);
REGISTER_REARRANGE_RULE(UrlSharding, CreateXShardingRule);
