#include "ut_common.h"

#include <balancer/modules/cookie_policy/common/combined_policy.h>
#include <balancer/modules/cookie_policy/common/controls_reader.h>
#include <balancer/modules/cookie_policy/policies/policy_traits.h>

#include <balancer/kernel/cookie/cookie.h>
#include <balancer/kernel/cookie/gdpr/lifetime.h>

#include <balancer/kernel/testing/conn_descr.h>
#include <balancer/kernel/testing/process_mock.h>
#include <balancer/kernel/testing/stats.h>

namespace NModCookiePolicy {
    using namespace NSrvKernel::NTesting;

    namespace {
        template <class TCtl, class TCtlCfg>
        THolder<TCtl> GenCtl(TString data, TString fname, const TConnDescr& descr) {
            if (data) {
                return ParseControls<TCtl, TCtlCfg>(
                    data, "./controls/" + fname, "cookie_policy", descr
                );
            }
            return {};
        }
    }

    TCPTestResult TestCookiePolicy(const TCPTestInput& inp) {
        TCombinedPolicy policy = InitCombinedPolicy(inp.Cfg);

        TStatsFixture stats;
        TCombinedPolicyTls tlsTemplate(policy, stats.Manager());
        stats.FreezeAndInit();

        TProcessMock proc;
        TCombinedPolicyTls tls(tlsTemplate, proc);

        TTestConnDescr descr(proc);

        if (inp.IsSecure) {
            descr.ConnDescr().Properties->UserConnIsSsl = true;
        }

        TLog log("cerr");
        descr.ConnDescr().Properties->Parent.GdprCache = inp.GeoCache;
        descr.ConnDescr().Properties->Start = inp.Now;
        descr.ConnDescr().Properties->Random = inp.Random;
        descr.ConnDescr().ErrorLog = &log;
        descr.ConnDescr().Request->RequestLine().Path = TStringStorage(inp.Path);

        TStringBuilder reqLog;
        descr.ConnDescr().ExtraAccessLog = TAccessLogOutput(&reqLog.Out);

        for (auto&& [n, hs] : std::initializer_list<std::pair<TString, TStrVec>>{
            // While multiple host headers are guaranteed to not happen, they still should not cause a crash.
            {"host", inp.Host},
            {"user-agent", inp.UserAgent},
            {"x-yandex-randomuid", inp.XYandexRandomUid},
            {"x-yandex-eu-request", inp.XYandexEURequest},
            {"x-ip-properties", inp.XIpProperties},
            {"cookie", inp.Cookie},
        }) {
            for (auto&& h : hs) {
                descr.ConnDescr().Request->Headers().Add(n, h);
            }
        }

        THolder<TModeControls> modeCtl = GenCtl<TModeControls, TModeControlsCfg>(inp.ModeCtl, "modectl", descr.ConnDescr());
        tls.SetModeControls(modeCtl.Get());

        THolder<TGdprControls> gdprCtl = GenCtl<TGdprControls, TGdprControlsCfg>(inp.GdprCtl, "gdprctl", descr.ConnDescr());
        tls.SetGdprControls(gdprCtl.Get());

        TResponse resp;
        for (auto&& h : inp.SetCookie) {
            resp.Headers().Add("set-cookie", h);
        }

        TCPTestResult res;
        if (auto ctx = tls.ApplyToRequest(descr.ConnDescr())) {
            res.ReqLog = reqLog;
            res.RespLog = tls.ApplyToResponse(*ctx, descr.ConnDescr(), resp);
        }

        res.Gdpr = descr.ConnDescr().Properties->Gdpr;
        res.GeoCache = descr.ConnDescr().Properties->Parent.GdprCache;

        for (const auto& h : descr.ConnDescr().Request->Headers().GetValuesRef("cookie")) {
            res.Cookie.emplace_back(h.AsStringBuf());
        }
        for (const auto& h : resp.Headers().GetValuesRef("set-cookie")) {
            res.SetCookie.emplace_back(h.AsStringBuf());
        }

        if (inp.WithStats) {
            res.Stats = GetUnistat(stats, true);
        }
        return res;
    }

    TString ToOneCookie(const TSetCookie& c) {
        return TStringBuilder() << c;
    }

    TStrVec ToCookie(const TVector<TSetCookie>& c) {
        TStrVec res;
        for (auto&& sc : c) {
            res.emplace_back(ToOneCookie(sc));
        }
        return res;
    }

    void DoTestPolicy(const TPolicyTestCase& c, TSourceLocation loc) {
        UNIT_ASSERT_VALUES_EQUAL_C(TestCookiePolicy(c.In), c.Out, loc);
    }

    TModuleConfig MakeCombinedPolicy(TString raw) {
        TStringInput sin(raw);
        return NProtoConfig::ParseConfig<TModuleConfig>(
            *NConfig::ConfigParser(sin),
            [](const NProtoConfig::TKeyStack& ctx, const TString& key, NConfig::IConfig::IValue*) {
                UNIT_ASSERT_C(false, "Unknown key " << key << " at " << ctx);
            }
        );
    }

    TStatsVec FixStats(TStatsVec v, ui32 cookieCnt) {
        for (auto&& s : TStatsVec{
            {"cpol-xxx-checked_summ", 1},
            {"cpol-xxx-total_summ", 1},
            {"cpol-xxx-parser-total_summ", 1},
        }) {
            v.emplace_back(s);
        }
        if (cookieCnt) {
            for (auto&& s : TStatsVec{
                {"cpol-xxx-parser-cookie-pass_summ", cookieCnt},
                {"cpol-xxx-parser-cookie-total_summ", cookieCnt},
                {"cpol-xxx-parser-pass_summ", 1},
            }) {
                v.emplace_back(s);
            }
        } else {
            v.emplace_back("cpol-xxx-parser-skip_summ", 1);
        }
        return v;
    }

    TStatsVec FixStatsNoGdpr(TStatsVec v, ui32 cookieCnt) {
        v = FixStats(std::move(v), cookieCnt);
        for (auto&& s : TStatsVec{
            {"cpol-xxx-gdpr-geo-0_summ", 1},
            {"cpol-xxx-gdpr-mod-0_summ", 1},
            {"cpol-xxx-gdpr-src-None_summ", 1},
            {"cpol-xxx-gdpr-safe-geo-0_summ", 1},
            {"cpol-xxx-gdpr-safe-mod-0_summ", 1},
            {"cpol-xxx-gdpr-safe-src-None_summ", 1},
        }) {
            v.emplace_back(s);
        }
        return v;
    }

    TSetCookie IsGdprSetCookie(bool value, TInstant now, std::string domain, ESameSite ss) {
        return {
            .Name="is_gdpr",
            .Value=NCookie::FromStringBuf(value ? "1" : "0"),
            .Path="/",
            .Domain=domain,
            .Expires=TMaxAllowedExpires::Get(now),
            .SameSite=ss,
            .Secure=(ss == ESameSite::None),
        };
    }

    TSetCookie IsGdprBSetCookie(TIsGdprB value, TInstant now, std::string domain, ESameSite ss) {
        return {
            .Name="is_gdpr_b",
            .Value=TBlob::FromString(value.Render()),
            .Path="/",
            .Domain=domain,
            .Expires=TMaxAllowedExpires::Get(now),
            .SameSite=ss,
            .Secure=(ss == ESameSite::None),
        };
    }
}

using namespace NModCookiePolicy;

template <>
void Out<TCPTestResult>(IOutputStream& out, const TCPTestResult& r) {
    out << ".GeoCache=" << r.GeoCache << ",\n";
    out << ".Gdpr=" << r.Gdpr << ",\n";

    out << ".Cookie={\n";
    for (auto&& c : r.Cookie) {
        out << "    " << c.Quote() << ",\n";
    }
    out << "},\n";

    out << ".SetCookie={\n";
    for (auto&& c : r.SetCookie) {
        out << "    " << ToString(c).Quote() << ",\n";
    }
    out << "},\n";

    out << ".ReqLog=" << r.ReqLog.Quote() << ",\n";
    out << ".RespLog=" << r.RespLog.Quote() << ",\n";
    out << ".Stats={\n";
    for (auto&& c : Sorted(r.Stats)) {
        out << "    {" << c.first.Quote() << ", " << c.second << "},\n";
    }
    out << "},\n";
}
