#include "module.h"

#include <balancer/modules/exp_static/calc/calc.h>

#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/fs/kv_file_consumer.h>
#include <balancer/kernel/fs/shared_file_rereader.h>
#include <balancer/kernel/fs/shared_files.h>
#include <balancer/kernel/log/errorlog.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/process/thread_info.h>

#include <util/thread/singleton.h>

using namespace NSrvKernel;
using namespace NModExpStatic;

Y_TLS(exp_static) {
    bool RateFileExists() const noexcept {
        return !RateFileReader.Empty();
    }

    TExpRate Rate;
    TString CombinedSalt;

    TSharedFileReReader RateFileReader;
    TSharedFileReReader::TData RateFileData;
};

MODULE_WITH_TLS_BASE(exp_static, TModuleWithSubModule) {
public:
    using TRateFileReaction = TBaseFileReaction<ui64, TExpRate>;
    using TRateFileConsumer = TKvFileConsumer<TRateFileReaction>;

    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        Config->ForEach(this);

        Y_ENSURE_EX(Submodule_,
            TConfigParseError() << "no module configured for exp_static");
        Y_ENSURE_EX(Conf_.ExpId,
            TConfigParseError() << "empty experiment group id(exp_id) in exp_static");
        Y_ENSURE_EX(Conf_.ContId,
            TConfigParseError() << "empty control group id(cont_id) in exp_static");
        Y_ENSURE_EX(Conf_.ExpId != Conf_.ContId,
            TConfigParseError() << "exp_id and cont_id must differ in exp_static");
    }

private:
    START_PARSE {
        PARSE_EVENTS;

        ON_KEY("exp_id", Conf_.ExpId) {
            return;
        }

        ON_KEY("cont_id", Conf_.ContId) {
            return;
        }

        ON_KEY("salt", Conf_.Salt) {
            return;
        }

        ON_KEY("slots_count", Conf_.SlotsCount) {
            return;
        }

        ON_KEY("buckets_count", Conf_.BucketsCount) {
            return;
        }

        ON_KEY("rate_file", RateFilename_) {
            return;
        }

        {
            Submodule_.Reset(Loader->MustLoad(key, Copy(value->AsSubConfig())).Release());
            return;
        }
    } END_PARSE

    THolder<TTls> DoInitTls(IWorkerCtl* process) override {
        auto tls = MakeHolder<TTls>();
        if (!!RateFilename_) {
            tls->RateFileReader = process->SharedFiles()->FileReReader(RateFilename_, TDuration::Seconds(1));
        }
        return tls;
    }

    bool DoCanWorkWithoutHTTP() const noexcept override {
        return true;
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        UpdateRate(descr, tls);

        if (tls.RateFileExists()) {
            TSockAddr addr;
            Y_PROPAGATE_ERROR(TSockAddr::FromRemoteAddr(descr.RemoteAddr()).AssignTo(addr))
            auto exp = CalcExpByIp(Conf_, tls.Rate, tls.CombinedSalt, addr.Ip());

            if (ESlotType::Exp == exp.Type) {
                descr.Properties->Parent.Experiments[exp.Id] = exp.Bucket;
            } else if (ESlotType::Cont == exp.Type) {
                descr.Properties->Parent.Experiments[exp.Id] = exp.Bucket;
            }
        }

        return Submodule_->Run(descr);
    }

    void UpdateRate(const TConnDescr& descr, TTls& tls) const noexcept {
        const auto& data = tls.RateFileReader.Data();
        if (data.Id() == tls.RateFileData.Id()) {
            return;
        }
        tls.RateFileData = data;

        try {
            TRateFileConsumer consumer;
            ProcessKvData(consumer, tls.RateFileData.Data());
            const auto item = consumer.Storage().find(Conf_.ExpId);
            if (item != consumer.Storage().end()) {
                if (item->second.Value >= 0 && item->second.Value <= 1.0) {
                    tls.Rate = item->second;
                    tls.CombinedSalt = GenCombinedSalt(Conf_, tls.Rate);
                } else {
                    LOG_ERROR(TLOG_ERR, descr, "Rate for exp_id " << Conf_.ExpId
                            << " must be in range 0.0 - 1.0: " << item->second.Value);
                }
            } else {
                tls.Rate.Value = 0;
            }
        } catch (const yexception& e) {
            LOG_ERROR(TLOG_ERR, descr, "Error parsing rate file: " << e.what());
        }
    }

private:
    TString RateFilename_;
    TExpConf Conf_;
};

NSrvKernel::IModuleHandle* NModExpStatic::Handle() {
    return TModule::Handle();
}
