#include "cookie_policy.h"

namespace NModCookiePolicy {
    namespace {
        void CombineResults(TVector<TResponseCookieAction>& results, const TVector<TSetCookie>& added, TCookieErrors<EResponseCookieReason> reasons) {
            if (!added) {
                return;
            }
            auto it = added.begin();
            results.emplace_back(TResponseCookieAction{
                .Info=it->Name,
                .Reasons=reasons,
                .Action=ECookieAction::Add,
            });

            for (++it; it != added.end(); ++it) {
                if (it->Name != results.back().Info) {
                    results.emplace_back(TResponseCookieAction{
                        .Info=it->Name,
                        .Reasons=reasons,
                        .Action=ECookieAction::Add,
                    });
                } else {
                    results.back().Count += 1;
                }
            }
        }
    }

    TPolicyTls::TPolicyTls(const TPolicyHolder& policy, TSharedStatsManager& manager)
        : Policy(policy)
        , CustomTls(Policy.Policy().NewTls(TString(policy.Uuid()), manager))
        , Stats(TString(policy.Uuid()), manager)
        , CookieStats(TString(policy.Uuid()), manager)
    {}

    TPolicyTls::TPolicyTls(const TPolicyTls& tmpl, IWorkerCtl& ctl)
        : Policy(tmpl.Policy)
        , CustomTls(tmpl.CustomTls ? tmpl.CustomTls->Clone(ctl) : THolder<IPolicyCustomTls>())
        , Stats(tmpl.Stats, ctl.WorkerId())
        , CookieStats(tmpl.CookieStats, ctl.WorkerId())
    {}

    TPolicyCtx::TPolicyCtx(TPolicyTls& tls, bool dryRun) noexcept
        : Tls(tls)
        , CustomCtx(tls.Policy.Policy().NewCtx())
        , DryRun(dryRun)
    {}

    TPolicyHolder::TPolicyHolder(TString moduleUuid, THolder<ICookiePolicy> policy)
        : Uuid_(JoinNonemptySeq({moduleUuid, policy->Name()}))
        , Policy_(std::move(policy))
    {}

    bool TPolicyHolder::MatchRequest(const TPolicyTls& tls, const TRequestArgs& args) const noexcept {
        const TPolicyArgs policyArgs{
            .Descr=args.Descr,
            .DomainInfo=args.DomainInfo,
            .IsGdprCookie=args.IsGdprCookie,
            .IsGdprUpdate=args.IsGdprUpdate,
            .IsGdprBCookie=args.IsGdprBCookie,
            .IsGdprBUpdate=args.IsGdprBUpdate,
            .CookieClassifier=args.CookieClassifier,
            .Tls=tls.CustomTls.Get(),
        };
        return Policy_->MatchRequest(policyArgs);
    }

    TRequestCookieActions TPolicyHolder::ApplyToRequestCookies(
        TPolicyCtx& ctx, const TRequestArgs& args
    ) const noexcept {
        auto actions = Policy_->ApplyToRequestCookies(args.Cookies, {
            .Descr=args.Descr,
            .DomainInfo=args.DomainInfo,
            .IsGdprCookie=args.IsGdprCookie,
            .IsGdprUpdate=args.IsGdprUpdate,
            .IsGdprBCookie=args.IsGdprBCookie,
            .IsGdprBUpdate=args.IsGdprBUpdate,
            .CookieClassifier=args.CookieClassifier,
            .Tls=ctx.Tls.CustomTls.Get(),
            .Ctx=ctx.CustomCtx.Get(),
            .DryRun=ctx.DryRun,
        });

        for (auto&& act : actions) {
            switch (act.Action) {
            case ECookieAction::Add:
                if (ctx.DryRun) {
                    ctx.Tls.CookieStats.ReqCookieAddDryRun.Add(act.Count);
                } else {
                    ctx.Tls.CookieStats.ReqCookieAdd.Add(act.Count);
                }
                break;
            case ECookieAction::Drop:
                if (ctx.DryRun) {
                    ctx.Tls.CookieStats.ReqCookieDropDryRun.Add(act.Count);
                } else {
                    ctx.Tls.CookieStats.ReqCookieDrop.Add(act.Count);
                }
                break;
            case ECookieAction::Fix:
                if (ctx.DryRun) {
                    ctx.Tls.CookieStats.ReqCookieFixDryRun.Add(act.Count);
                } else {
                    ctx.Tls.CookieStats.ReqCookieFix.Add(act.Count);
                }
                break;
            }
        }

        return actions;
    }

    TResponseCookieActions TPolicyHolder::ApplyToResponseCookies(
        TPolicyCtx& ctx, const TResponseArgs& args
    ) const noexcept {
        const TPolicyArgs policyArgs{
            .Descr=args.Descr,
            .DomainInfo=args.DomainInfo,
            .IsGdprCookie=args.IsGdprCookie,
            .IsGdprUpdate=args.IsGdprUpdate,
            .IsGdprBCookie=args.IsGdprBCookie,
            .IsGdprBUpdate=args.IsGdprBUpdate,
            .CookieClassifier=args.CookieClassifier,
            .DeletionSoftMax=args.DeletionSoftMax,
            .DeletionHardMax=args.DeletionHardMax,
            .Tls=ctx.Tls.CustomTls.Get(),
            .Ctx=ctx.CustomCtx.Get(),
            .DryRun=ctx.DryRun,
        };

        TResponseCookieActions actions;

        for (auto it = args.Cookies.begin(); it != args.Cookies.end();) {
            if (std::holds_alternative<TSetCookieSyntaxErrors>(it->Cookie())) {
                ctx.Tls.CookieStats.CookieTotal.Add(1);
                ctx.Tls.CookieStats.CookieSkip.Add(1);
            } else if (!it->IsAddedHeader()) {
                ctx.Tls.CookieStats.CookieTotal.Add(1);
                const auto& c = std::get<TSetCookie>(it->Cookie());

                if (Policy_->MatchResponseCookie(c, policyArgs)) {
                    if (auto act = Policy_->ApplyToResponseCookie(it->View(), policyArgs)) {
                        actions.emplace_back(TResponseCookieAction{
                            .Info=it->View().Const().Name,
                            .Reasons=act->Reasons,
                            .Action=act->Action,
                        });
                        ctx.Tls.CookieStats.CookieFail.Add(1);
                        if (ECookieAction::Drop == act->Action) {
                            if (ctx.DryRun) {
                                ctx.Tls.CookieStats.CookieDropDryRun.Add(1);
                            } else {
                                ctx.Tls.CookieStats.CookieDrop.Add(1);
                                it = args.Cookies.erase(it);
                                continue;
                            }
                        } else {
                            if (ctx.DryRun) {
                                ctx.Tls.CookieStats.CookieFixDryRun.Add(1);
                            } else {
                                ctx.Tls.CookieStats.CookieFix.Add(1);
                            }
                        }
                    } else {
                        ctx.Tls.CookieStats.CookiePass.Add(1);
                    }
                } else {
                    ctx.Tls.CookieStats.CookieSkip.Add(1);
                }
            }

            ++it;
        }

        Policy_->AddSetCookies([&](const TVector<TSetCookie>& added, TCookieErrors<EResponseCookieReason> reasons) {
            if (added) {
                CombineResults(actions, added, reasons);
            } else {
                actions.emplace_back(TResponseCookieAction{
                    .Reasons=reasons,
                    .Action=ECookieAction::Add,
                    .Count=0,
                });
            }
            if (ctx.DryRun) {
                ctx.Tls.CookieStats.CookieAddDryRun.Add(added.size());
            } else {
                ctx.Tls.CookieStats.CookieAdd.Add(added.size());
                for (auto a : added) {
                    args.Cookies.emplace_back(TResponseCookie::TAdded(), a);
                }
            }
        }, policyArgs);

        return actions;
    }


    TNameMatchPolicyBase::TNameMatchPolicyBase(TString name, EPolicyMode mode, TString nameRe, bool matchDeletion)
        : ICookiePolicy(name, mode)
        , NameRe_(nameRe)
        , MatchDeletion_(matchDeletion)
    {
        Y_ENSURE_EX(nameRe, TConfigParseError() << "name_re required");
    }

    bool TNameMatchPolicyBase::MatchResponseCookie(const TSetCookie& c, const TPolicyArgs& args) const noexcept {
        return Match(NameRe_, c.Name) && (MatchDeletion_ || !HasDeletion(c, args.Descr.Properties->Start));
    }
}
