#include "module.h"
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/storage/storage.h>

using namespace NSrvKernel;
using namespace NSrvKernel::NStorage;

#include "helpers.h"
#include "cache_processor.h"

Y_TLS(cache2) {
    TTls(const TSharedCounter& counter, size_t workerId)
            : CacheHit(counter, workerId)
    {}

    TSharedCounter CacheHit;
};

MODULE_WITH_TLS_BASE(cache2, TModuleWithSubModule) {
public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        Config->ForEach(this);

        CacheHitCounter_ = TSharedCounter(
                Control->SharedStatsManager().MakeCounter(
                        TString::Join("cache2-", (Uuid_ ? Uuid_ + "-": ""), "cache_hit")).AllowDuplicate().Build());

        Storage_ = MakeHolder<TStorage> (ShardNumber_);
        Storage_->CacheTTL_ = CacheTTL_;
        CacheProcessor_ = MakeHolder<TCacheProcessor>(IgnoreCacheControl_);

        Control->RegisterStorage(Storage_.Get());
    }

private:
    START_PARSE {
        ON_KEY("cache_ttl", CacheTTL_) {
            return;
        }
        ON_KEY("shard_number", ShardNumber_) {
            return;
        }
        ON_KEY("ignore_cgi", IgnoreCGI_) {
            StorageKeyGenerator_.SetIgnoreCGI(IgnoreCGI_);
            return;
        }
        ON_KEY("ignore_cache_control", IgnoreCacheControl_) {
            return;
        }
        ON_KEY("uuid", Uuid_) {
            return;
        }
        {
            Submodule_.Reset(Loader->MustLoad(key, Copy(value->AsSubConfig())).Release());
            return;
        }
    } END_PARSE

private:
    THolder<TTls> DoInitTls(IWorkerCtl* process) override {
        return MakeHolder<TTls>(*CacheHitCounter_, process->WorkerId());
    }

    TError PerformRequestNoCache(const TConnDescr& descr) const noexcept {
        return Submodule_->Run(descr);
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        TStorageKey key = StorageKeyGenerator_.Generate(*descr.Request);
        auto strategy = CacheProcessor_->ProcessRequest(*descr.Request);

        auto storageItem = Storage_->Get(key);
        descr.ExtraAccessLog << " " << TString(descr.Request->RequestLine().Path.AsStringBuf());
        // cache miss
        if (!storageItem.Contains) {
            TString tmpList;
            TResponse tmpResponse;

            descr.ExtraAccessLog << " cache miss";

            auto output = MakeHttpOutput([&](TResponse&& response, bool forceClose, TInstant deadline) {
                tmpResponse = response;
                CacheProcessor_->UpdateStrategyProcessingResponse(response, &strategy);
                return descr.Output->SendHead(std::move(response), forceClose, deadline);
            }, [&](TChunkList lst, TInstant deadline) {
                tmpList += TString(Union(lst)->AsStringBuf());
                return descr.Output->Send(std::move(lst), deadline);
            }, [&](THeaders&& trailers, TInstant deadline) {
                return descr.Output->SendTrailers(std::move(trailers), deadline);
            });

            Y_TRY(TError, error) {
                Y_PROPAGATE_ERROR(PerformRequestNoCache(descr.CopyOut(output)));
                return {};
            } Y_CATCH {
                CacheProcessor_->UpdateStrategyProcessingException(&strategy);
                return error;
            }

            if (strategy.ShouldFlushToCache) {
                Storage_->StoreData(key, std::move(tmpResponse), tmpList);
            }

            return {};
        } else { // cache hit
            if (strategy.GetFromCache) {
                descr.ExtraAccessLog << " cache hit " << storageItem.InsertTime;
                descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "cache hit");
                Y_PROPAGATE_ERROR(descr.Output->SendHead(std::move(storageItem.Response), false, TInstant::Max()));
                if (storageItem.Body.size() > 0) {
                    Y_PROPAGATE_ERROR(descr.Output->Send(TChunkList(storageItem.Body), TInstant::Max()));
                }
                ++tls.CacheHit;
            } else {
                descr.ExtraAccessLog << " avoid cache " << storageItem.InsertTime;
                Y_PROPAGATE_ERROR(PerformRequestNoCache(descr));
                return {};
            }
        }

        return descr.Output->SendEof(TInstant::Max());
    }

private:
    size_t ShardNumber_ = 10;
    bool IgnoreCGI_ = true;
    bool IgnoreCacheControl_ = true;
    TDuration CacheTTL_ = TDuration::Seconds(5);
    TKeyGenerator StorageKeyGenerator_;
    mutable THolder<TCacheProcessor> CacheProcessor_;
    mutable THolder<TStorage> Storage_;
    TMaybe<TSharedCounter> CacheHitCounter_;
    TString Uuid_;
};

IModuleHandle* NModCache2::Handle() {
    return TModule::Handle();
}

