#include <balancer/modules/redirects/redirects.h>

#include <library/cpp/iterator/functools.h>
#include <library/cpp/testing/unittest/registar.h>

using namespace NModRedirects;
using namespace NModRedirects::NImpl;


Y_UNIT_TEST_SUITE(TRedirectsTest) {

    TTemplate GenTemplate(std::initializer_list<TTemplate::value_type> vals) {
        TTemplate res;
        for (auto&& val : vals) {
            if (std::holds_alternative<TString>(val) && !std::get<TString>(val)) {
                continue;
            }
            res.emplace_back(val);
        }
        return res;
    }

    Y_UNIT_TEST(TestTemplate) {
        for (auto root : {TString(), TString("/"), TString("///"), TString("//host/"), TString("https://host/")}) {
            for (auto frag : {TString(), TString("#?")}) {
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + frag),
                    GenTemplate({root + frag}));
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + "{path}" + frag),
                    GenTemplate({root, EPlace::Path, frag}));
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + "{query}" + frag),
                    GenTemplate({root, EPlace::Q_Query, frag}));
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + "{path}{query}" + frag),
                    GenTemplate({root, EPlace::Path, EPlace::Q_Query, frag}));
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + "{path}?{query}" + frag),
                    GenTemplate({root, EPlace::Path, TString("?"), EPlace::Query, frag}));
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + "{path}?a=b&{query}" + frag),
                    GenTemplate({root, EPlace::Path, TString("?a=b&"), EPlace::Query, frag}));
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + "{path}?a=b{query}" + frag),
                    GenTemplate({root, EPlace::Path, TString("?a=b"), EPlace::A_Query, frag}));
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + "{path}?a=b{query}&c=d" + frag),
                    GenTemplate({root, EPlace::Path, TString("?a=b"), EPlace::A_Query, "&c=d" + frag}));
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + "ab{path}?a=b{query}&c=d" + frag),
                    GenTemplate({root + "ab", EPlace::S_Path, TString("?a=b"), EPlace::A_Query, "&c=d" + frag}));
                UNIT_ASSERT_VALUES_EQUAL(
                    GenLocation(root + "ab{path}/cd?a=b{query}&c=d" + frag),
                    GenTemplate({root + "ab", EPlace::S_Path, TString("/cd?a=b"), EPlace::A_Query, "&c=d" + frag}));
            }
        }

        UNIT_ASSERT_VALUES_EQUAL(
            GenLocation("https://yandex.ru/supp/trouble/list.xml?f13397=12786&f15015=18656&$args"),
            GenTemplate({TString("https://yandex.ru/supp/trouble/list.xml?f13397=12786&f15015=18656&"), EPlace::Query})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            GenLocation("https://yandex.by/bus/$request_uri$args"),
            GenTemplate({TString("https://yandex.by/bus/"), EPlace::Path, EPlace::Q_Query})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            GenLocation("https://yandex.by/bus$request_uri$args"),
            GenTemplate({TString("https://yandex.by/bus"), EPlace::S_Path, EPlace::Q_Query})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            GenLocation("https://yandex.ru/support/parking/feedback.xml$args#q1"),
            GenTemplate({TString("https://yandex.ru/support/parking/feedback.xml"), EPlace::Q_Query, TString("#q1")})
        );
    }

    Y_UNIT_TEST(TestSrc) {
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("//yandex", false), "yandex/");
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("//yandex?xxx", false), "yandex/?xxx");
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("//yandex/", false), "yandex/");
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("//yandex/?xxx", false), "yandex/?xxx");
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("//yandex/*", false), "yandex/*");
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("//yandex", true), "yandex/*");
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("//yandex?xxx", true), "yandex/*?xxx");
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("//yandex/*", true), "yandex/*");
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("http://yandex/*", true), "yandex/*");
        UNIT_ASSERT_VALUES_EQUAL(GenSrc("https://yandex/*", true), "yandex/*");
    }

    struct TUrlPartsParams {
        bool Host = false;
        bool Path = false;
        bool Query = false;
        bool Fragment = false;
    };

    TUrlParts UrlParts(TMaybe<TUrlPartsParams> params) {
        TUrlParts parts;
        if (params) {
            parts.host = params->Host;
            parts.path = params->Path;
            parts.query = params->Query;
            parts.fragment = params->Fragment;
        }
        return parts;
    }

    TRewriteConfig MakeRewrite(TString regexp, TString rewrite, bool global = false, TUrlParts parts = {}) {
        TRewriteConfig cfg;
        cfg.set_regexp(regexp);
        cfg.set_rewrite(rewrite);
        cfg.set_global(global);
        cfg.set_url(parts);
        return cfg;
    }

    Y_UNIT_TEST(TestRewrite) {
        UNIT_ASSERT_VALUES_EQUAL(TRewrite(MakeRewrite("(x)[.]xml$", "$1"))
            .Apply("https://yandex.ru/x/x/x.xml?y=y&y#z.z.z"),
            "https://yandex.ru/x/x/x?y=y&y#z.z.z");

        UNIT_ASSERT_VALUES_EQUAL(TRewrite(MakeRewrite("x", "y", false))
            .Apply("https://yandex.ru/x/x/x?x=x&x#x.x.x"), "https://yandex.ru/y/x/x?x=x&x#x.x.x");
        UNIT_ASSERT_VALUES_EQUAL(TRewrite(MakeRewrite("x", "y", true))
            .Apply("https://yandex.ru/x/x/x?x=x&x#x.x.x"), "https://yandex.ru/y/y/y?x=x&x#x.x.x");
        UNIT_ASSERT_VALUES_EQUAL(TRewrite(MakeRewrite("x", "y", true, UrlParts(TUrlPartsParams{.Query=true})))
            .Apply("https://yandex.ru/x/x/x?x=x&x#x.x.x"), "https://yandex.ru/x/x/x?y=y&y#x.x.x");
        UNIT_ASSERT_VALUES_EQUAL(TRewrite(MakeRewrite("x", "y", true, UrlParts(TUrlPartsParams{.Fragment=true})))
            .Apply("https://yandex.ru/x/x/x?x=x&x#x.x.x"), "https://yandex.ru/x/x/x?x=x&x#y.y.y");
        UNIT_ASSERT_VALUES_EQUAL(TRewrite(MakeRewrite("x", "y", true, UrlParts(TUrlPartsParams{
            .Path=true, .Query=true, .Fragment=true
        }))).Apply("https://yandex.ru/x/x/x?x=x&x#x.x.x"), "https://yandex.ru/y/y/y?y=y&y#y.y.y");

        // MINOTAUR-2606
        auto rwr = TRewrite(MakeRewrite(
            "([0-9]+)[.]games[.]([^/]+)(/.*)",
            "$2/games/play/$1$3", false, UrlParts(TUrlPartsParams{.Host=true, .Path=true})));

        for (auto [inp, out] : std::initializer_list<std::pair<TString, TString>>{
            {"https://98073.games.yandex.ru/", "https://yandex.ru/games/play/98073/"},
            {"https://98073.games.yandex.ru/xxx?yyy", "https://yandex.ru/games/play/98073/xxx?yyy"},
        }) {
            UNIT_ASSERT_VALUES_EQUAL(rwr.Apply(inp), out);
        }
    }

    TRedirectConfig Redir(
        TString dst, ui16 code=302, TVector<TRewriteConfig> rewrites = {}, TMaybe<bool> legacyStripping=Nothing())
    {
        TRedirectConfig rdr;
        rdr.set_dst(std::move(dst));
        rdr.set_dst_rewrites(std::move(rewrites));
        rdr.set_code(code);
        if (legacyStripping) {
            rdr.set_legacy_rstrip(*legacyStripping);
        }
        return rdr;
    }

    TForwardConfig Forward(TString dst, TVector<TRewriteConfig> rewrites = {}, TMaybe<bool> legacyStripping=Nothing()) {
        TForwardConfig fwd;
        fwd.set_dst(std::move(dst));
        fwd.set_dst_rewrites(std::move(rewrites));
        if (legacyStripping) {
            fwd.set_legacy_rstrip(*legacyStripping);
        }
        return fwd;
    }

    class TFakeModule : public IModule, public INodeHandle<IModule> {
        IModuleHandle* DoHandle() const override {
            return (IModuleHandle*)this;
        }

        const TString& DoName() const noexcept override {
            return Name;
        }

        IModule* DoConstruct(const TModuleParams&) const override {
            return nullptr;
        }

        TError DoRun(const TConnDescr&) const override {
            return {};
        }

    private:
        TString Name = "fake";
    };

    Y_UNIT_TEST(TestRealRedirects) {
        TFakeModule mod;
        TRedirects r;
        r.AddRedirect("//bus.yandex.ru/*", Redir("https://yandex.ru/bus/$request_uri$args", 301, {
            MakeRewrite("[.]xml$", ""),
            MakeRewrite("[.]php$", ""),
        }));
        r.AddRedirect("//yandex.ru/things/*", Redir("https://store.turbo.site/$args"));
        r.AddRedirect("//feedback2.yandex.ru/parking/partner/*", Redir("https://yandex.ru/support/parking/feedback.xml$args#q1"));
        r.AddRedirect("//www.yandex.aero", Redir("https://avia.yandex.ru/$request_uri$args"));
        r.AddRedirect("//firefox.yandex.ru", Redir("https://yandex.ru/?clid=2873"));
        // this one is fake
        r.AddRedirect("//yandex.ru/abc/?def", Redir("https://yandex.ru/yes"));
        r.AddRedirect("//visual.yandex.ru", Redir("https://yandex.ru/soft/element/vb?redir=vb$args"));
        r.AddRedirect("//pogoda.yandex.ru/*", Redir("https://yandex.ru/pogoda/{path}{query}", 302, {}, false));
        r.AddForward("//weather.yandex.ru/*", Forward("//yandex.ru/pogoda/{path}{query}", {}, false), mod);
        r.AddForward("//visual.yandex.ru/robots.txt", Forward("//visual.s3.yandex.net/robot.txt"), mod);
        r.Compile();
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("bus.yandex.by", "/", ""),
            TResult());
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("bus.yandex.ru", "/", ""),
            TResult(TRedirect{.Location="https://yandex.ru/bus/", .Code=301}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("bus.yandex.ru", "/xxx", "?yyy"),
            TResult(TRedirect{.Location="https://yandex.ru/bus/xxx?yyy", .Code=301}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("bus.yandex.ru", "/xxx/", "?yyy"),
            TResult(TRedirect{.Location="https://yandex.ru/bus/xxx?yyy", .Code=301}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("bus.yandex.ru", "/xxx/zzz", "?yyy"),
            TResult(TRedirect{.Location="https://yandex.ru/bus/xxx/zzz?yyy", .Code=301}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("bus.yandex.ru", "/xxx/zzz.php", "?yyy"),
            TResult(TRedirect{.Location="https://yandex.ru/bus/xxx/zzz?yyy", .Code=301}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("bus.yandex.ru", "/xxx/zzz.xml", "?yyy"),
            TResult(TRedirect{.Location="https://yandex.ru/bus/xxx/zzz?yyy", .Code=301}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("bus.yandex.ru", "/xxx/zzz.php.xml", "?yyy"),
            TResult(TRedirect{.Location="https://yandex.ru/bus/xxx/zzz?yyy", .Code=301}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("bus.yandex.ru", "/xxx/zzz.xml.php", "?yyy"),
            TResult(TRedirect{.Location="https://yandex.ru/bus/xxx/zzz.xml?yyy", .Code=301}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("yandex.ru", "/thing", ""),
            TResult());
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("yandex.ru", "/thingss", ""),
            TResult());
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("yandex.ru", "/things", ""),
            TResult(TRedirect{.Location="https://store.turbo.site/", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("yandex.ru", "/things", "?yyy"),
            TResult(TRedirect{.Location="https://store.turbo.site/?yyy", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("yandex.ru", "/things/xxx", "?yyy"),
            TResult(TRedirect{.Location="https://store.turbo.site/?yyy", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("yandex.ru", "/things/xxx/zzz/", "?yyy"),
            TResult(TRedirect{.Location="https://store.turbo.site/?yyy", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("www.yandex.aero", "/", ""),
            TResult(TRedirect{.Location="https://avia.yandex.ru/", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("www.yandex.aero", "/", "?yyy"),
            TResult(TRedirect{.Location="https://avia.yandex.ru/?yyy", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("www.yandex.aero", "/xxx", "?yyy"),
            TResult(TRedirect{.Location="https://avia.yandex.ru/xxx?yyy", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("www.yandex.aero", "/xxx/zzz", "?yyy"),
            TResult(TRedirect{.Location="https://avia.yandex.ru/xxx/zzz?yyy", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("www.yandex.aero", "/xxx/zzz/", "?yyy"),
            TResult(TRedirect{.Location="https://avia.yandex.ru/xxx/zzz?yyy", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("firefox.yandex.ru", "/", ""),
            TResult(TRedirect{.Location="https://yandex.ru/?clid=2873", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("firefox.yandex.ru", "/a", ""),
            TResult());
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("firefox.yandex.ru", "/", "?a"),
            TResult(TRedirect{.Location="https://yandex.ru/?clid=2873", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("yandex.ru", "/abc", "?def"),
            TResult(TRedirect{.Location="https://yandex.ru/yes", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("yandex.ru", "/abc/", "?def&"),
            TResult(TRedirect{.Location="https://yandex.ru/yes", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("yandex.ru", "/abc", ""),
            TResult());
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("visual.yandex.ru", "", "?xxx"),
            TResult(TRedirect{.Location="https://yandex.ru/soft/element/vb?redir=vb&xxx", .Code=302}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("visual.yandex.ru", "/xxx", "?xxx"),
            (TResult()));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("visual.yandex.ru", "/robots.txt", "?xxx"),
            TResult(TForward{.Location="//visual.s3.yandex.net/robot.txt", .Dst=&mod}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("visual.yandex.ru", "/robots.txt/", "?xxx"),
            TResult(TForward{.Location="//visual.s3.yandex.net/robot.txt", .Dst=&mod}));
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("pogoda.yandex.ru", "/", "?xxx"),
            TResult(TRedirect{.Location="https://yandex.ru/pogoda/?xxx", .Code=302})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("pogoda.yandex.ru", "/yyy", "?xxx"),
            TResult(TRedirect{.Location="https://yandex.ru/pogoda/yyy?xxx", .Code=302})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("pogoda.yandex.ru", "/yyy/", "?xxx&"),
            TResult(TRedirect{.Location="https://yandex.ru/pogoda/yyy/?xxx&", .Code=302})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("pogoda.yandex.ru", "//yyy/", "?xxx&"),
            TResult(TRedirect{.Location="https://yandex.ru/pogoda/yyy/?xxx&", .Code=302})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("weather.yandex.ru", "/", "?xxx"),
            TResult(TForward{.Location="//yandex.ru/pogoda/?xxx", .Dst=&mod})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("weather.yandex.ru", "/yyy", "?xxx"),
            TResult(TForward{.Location="//yandex.ru/pogoda/yyy?xxx", .Dst=&mod})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("weather.yandex.ru", "/yyy/", "?xxx&"),
            TResult(TForward{.Location="//yandex.ru/pogoda/yyy/?xxx&", .Dst=&mod})
        );
        UNIT_ASSERT_VALUES_EQUAL(
            r.Location("weather.yandex.ru", "//yyy/", "?xxx&"),
            TResult(TForward{.Location="//yandex.ru/pogoda/yyy/?xxx&", .Dst=&mod})
        );
    }
}
