#include "client.h"
#include "common.h"
#include "fixed_size_encode.h"

#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/http/parser/httpdecoder.h>
#include <balancer/kernel/http/parser/httpencoder.h>
#include <balancer/kernel/http/parser/response_builder.h>
#include <balancer/kernel/http/parser/request_builder.h>
#include <balancer/kernel/io/iobase.h>
#include <balancer/kernel/module/iface.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/log/errorlog.h>
#include <balancer/kernel/regexp/regexp_re2.h>
#include <balancer/kernel/requester/requester.h>

using namespace NSrvKernel;
using namespace NCache;

namespace {

[[nodiscard]] static TRequest GetStoreRequest(const TString& id, TStringBuf lastModifiedHeaderValue,
                                              TStringBuf etagHeaderValue) noexcept
{
    TRequest request = BuildRequest().Version11().Method(EMethod::PUT).Uri(id);
    if (lastModifiedHeaderValue) {
        request.Headers().Add(TLastModified::Key(), TString(lastModifiedHeaderValue));
    }

    if (etagHeaderValue) {
        request.Headers().Add(TETag::Key(), TString(etagHeaderValue));
    }
    return request;
}

}  // namespace

Y_TLS(cache_client) {
    ui64 CacheResponses = 0;
    ui64 CacheHits = 0;
};

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

        if (!IdRegexp_) {
            ythrow TConfigParseError() << "\"id_regexp\" is not set";
        }

        if (!Server_) {
            ythrow TConfigParseError() << "no server module configured";
        }

        if (!Submodule_) {
            ythrow TConfigParseError() << "no submodule configured";
        }
    }

private:
    START_PARSE {
        STATS_ATTR;

        TString id_regexp;
        ON_KEY("id_regexp", id_regexp) {
            if (id_regexp.empty()) {
                ythrow TConfigParseError() << "empty id regexp";
            }

            IdRegexp_.Reset(new TRegexp(id_regexp));
            return;
        }

        if (key == "server") {
            TSubLoader(Copy(value->AsSubConfig())).Swap(Server_);
            return;
        }

        if (key == "module") {
            TSubLoader(Copy(value->AsSubConfig())).Swap(Submodule_);
            return;
        }
    } END_PARSE

private:
    enum ECacheAction {
        CA_STORE,
        CA_VALIDATE,
    };

    TError Validate(const TConnDescr& descr, const TString& id, TStringBuf lastModifiedHeaderValue,
                    TStringBuf etagHeaderValue) const noexcept
    {
        auto request = GetStoreRequest(id, lastModifiedHeaderValue, etagHeaderValue);
        request.Headers().Add(TCacheAction::Key(), TValidate::Value());

        TRequester requester{*Server_, descr};
        return requester.Request(std::move(request));
    }

    TError Store(const TConnDescr& descr, const TString& id, TResponse response, TChunkList responseBody,
                 TStringBuf lastModifiedHeaderValue, TStringBuf etagHeaderValue) const noexcept
    {
        auto request = GetStoreRequest(id, lastModifiedHeaderValue, etagHeaderValue);
        TChunkList body = FixedSizeEncode(std::move(response), std::move(responseBody));

        TRequester requester{*Server_, descr};
        return requester.Request(std::move(request), std::move(body), false);
    }

private:
    TString Identificate(TRequest* request) const noexcept {
        const EMethod method = request->RequestLine().Method;

        if (method != EMethod::GET && method != EMethod::HEAD) {
            // Other methods may not be cached
            return {};
        }

        TVector<TStringBuf> buffers;

        auto url = request->RequestLine().GetURL();
        if (!IdRegexp_->Extract(url, &buffers, true)) {
            return {};
        }

        std::reverse(buffers.begin(), buffers.end());
        buffers.pop_back();

        if (buffers.empty()) {
            return {};
        }

        TString id = TString(buffers.back());
        buffers.pop_back();

        while (!buffers.empty()) {
            id += "//";
            id += buffers.back();
            buffers.pop_back();
        }

        if (!id.Empty()) {
            id += "//";
            id += ::ToString(method);
        }

        return id;
    }

private:
    TError SendResponse(const TConnDescr& descr, TChunksInput& response) const noexcept {
        TFromBackendDecoder decoder{&response, *descr.Request};
        TResponse decodedResponse;
        Y_PROPAGATE_ERROR(decoder.ReadResponse(decodedResponse, TInstant::Max()));
        TChunkList body;
        Y_PROPAGATE_ERROR(RecvAll(&response, body, TInstant::Max())); // or we need another way to make sure that inner response is full while transfering it
        decodedResponse.Props().ContentLength = body.size();
        const bool notEmpty = decodedResponse.Props().ContentLength > 0;
        Y_PROPAGATE_ERROR(descr.Output->SendHead(std::move(decodedResponse), false, TInstant::Max()));
        if (notEmpty) {
            Y_PROPAGATE_ERROR(descr.Output->Send(std::move(body), TInstant::Max()));
        }
        return descr.Output->SendEof(TInstant::Max());
    }

    TError RespondFromCache(const TConnDescr& descr, TChunkList cachedResponse,
                                          const TMaybe<bool>& cacheNotModifiedForUser) const noexcept
    {
        if (cacheNotModifiedForUser.GetOrElse(false)) {
            descr.ExtraAccessLog << " not modified";
            TResponse response = BuildResponse().Version11()
                                                .Code(HTTP_NOT_MODIFIED)
                                                .ContentLength(0);
            Y_PROPAGATE_ERROR(descr.Output->SendHead(std::move(response), false, TInstant::Max()));
            Y_PROPAGATE_ERROR(descr.Output->SendEof(TInstant::Max()));
        } else {
            descr.ExtraAccessLog << " cached";
            TChunksInput in{std::move(cachedResponse)};
            Y_PROPAGATE_ERROR(SendResponse(descr, in));
        }
        return {};
    }

    THolder<TTls> DoInitTls(IWorkerCtl*) override {
        return MakeHolder<TTls>();
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        const TInstant now = Now();
        const TString id = Identificate(descr.Request);

        if (!id) {
            return Submodule_->Run(descr);
        }

        /* Server response looks like:
         *
         * Cache server's response
         * Real response
         * Real body
         *
         * So there are 2 responses and we need to parse second one by hand
         */
        TResponse serverResponse;
        TChunkList serverOut;

        LOG_ERROR(TLOG_INFO, descr, "Cache: loading \"" << id << "\" from server");

        Y_TRY(TError, error) {
            TRequester requester{*Server_, descr};

            TRequest request = BuildRequest().Version11().Method(EMethod::GET).Uri(id);

            return requester.Request(std::move(request), serverResponse, serverOut);
        } Y_CATCH {
            LOG_ERROR(TLOG_INFO, descr, "Cache: failed to load \""
                << id << "\" from server: " << GetErrorMessage(error));
            return Submodule_->Run(descr);
        }

        ++tls.CacheResponses;

        const ui32 serverResponseStatus = serverResponse.ResponseLine().StatusCode;

        TMaybe<TInstant> lastModified;
        TStringBuf etag;
        TInstant validUntil = TInstant::Zero();
        bool mayValidate = false;

        if (serverResponseStatus == HTTP_OK) {
            ++tls.CacheHits;
            validUntil = TInstant::Max();
            mayValidate = true;

            auto lastModifiedHeaderValue = serverResponse.Headers().GetFirstValue(Default<TLastModified>());

            if (lastModifiedHeaderValue) try {
                lastModified = TLastModified::ParseValue(lastModifiedHeaderValue);
            } catch (TLastModified::TParseException&) {}

            auto validUntilValue = serverResponse.Headers().GetFirstValue(Default<TValidUntil>());

            if (validUntilValue) try {
                validUntil = TValidUntil::ParseValue(validUntilValue);
            } catch (TValidUntil::TParseException&) {}

            etag = serverResponse.Headers().GetFirstValue(Default<TETag>());
        }

        // true - may respond 304 based on cached data
        // false - may not respond 304 because cached data is stale
        // undefined - may not respond 304 because request lacks headers that enable it
        TMaybe<bool> cacheNotModifiedForUser;
        if (lastModified.Defined()) {
            TStringBuf ifModifiedSinceValue = descr.Request->Headers().GetFirstValue(Default<TIfModifiedSince>());
            TMaybe<TInstant> ifModifiedSince;
            if (ifModifiedSinceValue) try {
                ifModifiedSince = TIfModifiedSince::ParseValue(ifModifiedSinceValue);
            } catch (const TIfModifiedSince::TParseException&) {}

            if (ifModifiedSince.Defined()) {
                cacheNotModifiedForUser = ifModifiedSince.GetRef() >= lastModified.GetRef();
                if (ifModifiedSince.GetRef() > lastModified.GetRef()) {
                    // If cache is valid, we should return 304.
                    // If cache is out of date, we can not validate it because we may not replace request's If-Modified-Since value by an earlier one.
                    mayValidate = false;
                } else {
                    // If cache is valid, we still may modify the request because it will not be forwarded anywere.
                    // If cache is out of date, we may place request's If-Modified-Since valie by the latter one: when
                    // backend returns 304, we will serve cached reponse to client.
                    descr.Request->Headers().Delete(TIfModifiedSince::Key());
                    descr.Request->Headers().Add(TIfModifiedSince::Key(),
                                                       TIfModifiedSince::Value(lastModified.GetRef()));
                }
            } else {
                // If cache is valid, we still may modify the request because it will not be forwarded anywere.
                descr.Request->Headers().Delete(TIfModifiedSince::Key());
                descr.Request->Headers().Add(TIfModifiedSince::Key(),
                                                   TIfModifiedSince::Value(lastModified.GetRef()));
            }
        }

        if (etag) {
            if (TStringBuf ifNoneMatchValue = descr.Request->Headers().GetFirstValue(Default<TIfNoneMatch>())) {
                if (cacheNotModifiedForUser.GetOrElse(true)) {
                    cacheNotModifiedForUser = ifNoneMatchValue == etag;
                }
            } else {
                descr.Request->Headers().Add(TETag::Key(), TString(etag));
            }
        }

        if (validUntil >= now) {
            LOG_ERROR(TLOG_INFO, descr, "Cache: got valid cached \"" << id << "\" from server");
            descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), cacheNotModifiedForUser.GetOrElse(false) ? "not modified" : "cached");
            Y_PROPAGATE_ERROR(RespondFromCache(descr, std::move(serverOut), cacheNotModifiedForUser));
        } else {
            LOG_ERROR(TLOG_INFO, descr, "Cache: server has no valid cached \"" << id << '\n');

            TResponse moduleResponse;
            TChunkList moduleResponseBody;
            auto output = MakeHttpOutput([&](TResponse&& response, bool forceClose, TInstant) {
                Y_UNUSED(forceClose);
                moduleResponse = std::move(response);
                return TError{};
            }, [&](TChunkList lst, TInstant) {
                moduleResponseBody.Append(std::move(lst));
                return TError{};
            }, [](THeaders&&, TInstant) -> TError {
                Y_FAIL("trailers are not supported");
            });
            Y_PROPAGATE_ERROR(Submodule_->Run(descr.CopyOut(output)));

            const ui32 moduleResponseStatus = moduleResponse.ResponseLine().StatusCode;

            auto lastModifiedHeaderValue = moduleResponse.Headers().GetFirstValue(Default<TLastModified>());
            auto etagHeaderValue = moduleResponse.Headers().GetFirstValue(Default<TETag>());

            if (mayValidate && moduleResponseStatus == HTTP_NOT_MODIFIED) {
                LOG_ERROR(TLOG_INFO, descr, "Cache: serving validated \"" << id << '\n');
                descr.ExtraAccessLog << " validated";
                Y_DEFER {
                    TString reason("validated ");
                    reason.append(cacheNotModifiedForUser.GetOrElse(false) ? "not modified" : "cached");
                    descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), std::move(reason));
                };

                Y_PROPAGATE_ERROR(RespondFromCache(descr, std::move(serverOut), cacheNotModifiedForUser));

                LOG_ERROR(TLOG_INFO, descr, "Cache: bumping \"" << id << "\" at server");

                Y_TRY(TError, error) {
                    return Validate(descr, id, lastModifiedHeaderValue, etagHeaderValue);
                } Y_CATCH {
                    LOG_ERROR(TLOG_INFO, descr, "Cache: failed to bump \"" << id
                        << "\" at server: " << GetErrorMessage(error));
                }
            } else {
                LOG_ERROR(TLOG_INFO, descr, "Cache: serving backend response");

                TResponse response = moduleResponse;
                response.Props().ChunkedTransfer = false;
                response.Props().ContentLength = moduleResponseBody.size();

                Y_TRY(TError, error) {
                    Y_PROPAGATE_ERROR(descr.Output->SendHead(std::move(response), false, TInstant::Max()));
                    if (!moduleResponseBody.Empty()) {
                        Y_PROPAGATE_ERROR(descr.Output->Send(moduleResponseBody.Copy(), TInstant::Max()));
                    }
                    return descr.Output->SendEof(TInstant::Max());
                } Y_CATCH {
                    descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "client write error");
                    return error;
                };
            }

            if (moduleResponseStatus == HTTP_OK) {
                LOG_ERROR(TLOG_INFO, descr, "Cache: storing \"" << id << "\" to server");
                descr.ExtraAccessLog << " storing";
                TAccessLogSummary* summary = descr.ExtraAccessLog.Summary();
                TAccessLogSummary savedSummary;
                if (summary) {
                    savedSummary = *summary;
                }
                Y_DEFER {
                    descr.ExtraAccessLog.SetSummary(std::move(savedSummary.AnsweredModule), std::move(savedSummary.AnswerReason));
                };

                Y_TRY(TError, error) {
                    return Store(descr, id, std::move(moduleResponse),
                                 std::move(moduleResponseBody), lastModifiedHeaderValue,
                                 etagHeaderValue);
                } Y_CATCH {
                    LOG_ERROR(TLOG_INFO, descr, "Cache: failed to store \"" << id
                        << "\" to server: " << GetErrorMessage(error));
                }
            }
        }
        return {};
    }

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

private:
    THolder<TRegexp> IdRegexp_;
    THolder<IModule> Server_;

    TString StatsAttr_;
};

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