#include "server.h"
#include "config.h"

#include <balancer/server/modules/common.h>
#include <balancer/server/modules/callback.h>

#include <balancer/kernel/process/master_main.h>
#include <balancer/kernel/process/master_process.h>
#include <balancer/kernel/helpers/cast.h>
#include <balancer/kernel/custom_io/limitio.h>
#include <balancer/kernel/coro/universal_sleep.h>

#include <library/cpp/http/io/compression.h>

#include <util/system/event.h>
#include <util/thread/pool.h>
#include <util/stream/length.h>

namespace {
    class TInputAdapter
            : public IInputStream
    {
    public:
        explicit TInputAdapter(NSrvKernel::IIoInput& input)
            : Input_(input)
        {
        }

        size_t DoRead(void* buf, size_t len) override {
            if (!Error_ && Tmp_.Empty()) {
                Error_ = Input_.Recv(Tmp_, TInstant::Max());
            }

            if (Error_) {
                return 0;
            }

            size_t spaceLeft = len;
            char* const bufBegin = (char*)buf;
            char* writePos = bufBegin;

            while (spaceLeft && !Tmp_.Empty()) {
                NSrvKernel::TChunkPtr chunk = Tmp_.PopFront();

                const size_t toWrite = Min(spaceLeft, chunk->Length());
                memcpy(writePos, chunk->Data(), toWrite);
                chunk->Skip(toWrite);
                writePos += toWrite;
                spaceLeft -= toWrite;

                if (0 == chunk->Length()) {
                    continue;
                } else {
                    Tmp_.PushFront(std::move(chunk));
                }
            }

            return writePos - bufBegin;
        }

        NSrvKernel::TError TakeError() {
            return std::move(Error_);
        }
    private:
        NSrvKernel::IIoInput& Input_;
        NSrvKernel::TChunkList Tmp_;
        NSrvKernel::TError Error_;
    };
}

namespace NBalancerServer {
    THttpRequestEnv::THttpRequestEnv(const NSrvKernel::TConnDescr& connDescr)
        : Request_(*connDescr.Request)
        , StartTime_(connDescr.Properties->Start)
        , RemoteAddress_(connDescr.RemoteAddrStr())
        , Cont_(connDescr.Process().Executor().Running())
    {}

    class THttpRequestEnvImpl
        : public THttpRequestEnv
    {
    public:
        THttpRequestEnvImpl(const TServerOptions& serverOptions,
                        const NSrvKernel::TConnDescr& connDescr)
            : THttpRequestEnv(connDescr)
            , ServerOptions_(serverOptions)
            , Output_(*connDescr.Output)
        {
        }

        IHttpReplyTransport::TSendingRef GetReplyTransport() override {
            if (!ReplyTransport_) {
                ReplyTransport_ = IHttpReplyTransport::TSendingRef(
                    MakeAtomicShared<THttpReplyTransport>(
                        GetCont(),
                        Output_,
                        ServerOptions_.ReplyTransportSize,
                        ServerOptions_.ReplyTransportSendTimeout
                    )
                );

                if (ServerOptions_.EnableResponseCompression) {
                    ReplyTransport_->ChooseEncoding(Headers());
                }
            }

            return ReplyTransport_;
        }

        TAtomicSharedPtr<IHttpReplyTransport> GetReplyTransportHolder() {
            return ReplyTransport_.GetHolder();
        }

        NSrvKernel::TError ReceiveBody(const NSrvKernel::TConnDescr& connDescr, bool* tooLargeBody) {
            Body_.clear();

            RawLimitedInput_ = MakeHolder<NSrvKernel::TLengthLimitInput>(connDescr.Input, ServerOptions_.MaxRequestBodySize);

            InputAdapter_ = MakeHolder<TInputAdapter>(*RawLimitedInput_);
            IInputStream* input = InputAdapter_.Get();

            if (ServerOptions_.DecompressRequestBody) {
                if (auto contentEncoding = Headers().GetFirstValue("content-encoding")) {
                    auto decoderFactory = TCompressionCodecFactory::Instance().FindDecoder(contentEncoding);
                    if (decoderFactory) {
                        BodyDecompressed_ = true;
                        InputDecoder_ = (*decoderFactory)(input);
                        input = InputDecoder_.Get();
                    }
                }
            }

            if (!ServerOptions_.EnableInputStreaming) {
                class TBodySizeLimiter
                        : public IInputStream {
                  public:
                    TBodySizeLimiter(IInputStream& input, size_t size)
                            : Input_(input), Limited_(&input, size) {
                    }

                    size_t DoRead(void* buf, size_t len) override {
                        size_t res = Limited_.Read(buf, len);
                        if (!res && Input_.Skip(1)) {
                            IsTooLarge_ = true;
                        }
                        return res;
                    }

                    bool IsTooLarge() const {
                        return IsTooLarge_;
                    }

                  private:
                    IInputStream& Input_;
                    TLengthLimitedInput Limited_;
                    bool IsTooLarge_ = false;
                };


                TBodySizeLimiter decodedBodyLimiter(*input, ServerOptions_.MaxRequestBodySize);

                if (InputDecoder_) {
                    input = &decodedBodyLimiter;
                }


                TStringOutput so(Body_);

                auto timeBegin = Now();
                ::TransferData(input, &so);
                TransferAndDecodingDuration = Now() - timeBegin;

                if (!decodedBodyLimiter.IsTooLarge()) {
                    Y_PROPAGATE_ERROR(SkipAll(RawLimitedInput_.Get(), TInstant::Max()));
                }

                *tooLargeBody = decodedBodyLimiter.IsTooLarge() ||
                                (RawLimitedInput_->Depleted() && connDescr.Input->TryPeek(TInstant::Max()));

                auto err = InputAdapter_->TakeError();
                InputDecoder_.Reset();
                InputAdapter_.Reset();
                return err;
            }

            return {};
        }

        IInputStream* Input() const override {
            return InputAdapter_.Get();
        }

      private:
        THolder<NSrvKernel::TLengthLimitInput> RawLimitedInput_;
        THolder<TInputAdapter> InputAdapter_;
        THolder<IInputStream> InputDecoder_;
        const TServerOptions& ServerOptions_;
        NSrvKernel::IHttpOutput& Output_;

        IHttpReplyTransport::TSendingRef ReplyTransport_;
    };

    using namespace NSrvKernel;

    struct TCallbackError {
        TError Error;
        TAtomic HasError = 0;
    };

    using TCallbackErrorPtr = TAtomicSharedPtr<TCallbackError>;

    TError TransferReply(IHttpReplyTransport& reply, const TConnDescr& descr, const TCallbackEnv& cbEnv, TCallbackErrorPtr cbError) {
        Y_PROPAGATE_ERROR(reply.Transfer(descr.Process().Executor().Running(), cbEnv.EventWaker, *descr.Output));

        if (reply.IsReplyFinished()) {
            return {};
        } else if (cbError && AtomicGet(cbError->HasError)) {
            return std::move(cbError->Error);
        } else {
            return Y_MAKE_ERROR(yexception() << "not finished reply for unknown reason");
        }
    }

    NSrvKernel::TError SendSimpleResponse(const TConnDescr& descr, TStringBuf responseLine) {
        NSrvKernel::TResponse response;
        Y_PROPAGATE_ERROR(response.Parse(responseLine));
        Y_PROPAGATE_ERROR(descr.Output->SendHead(std::move(response), false, TInstant::Max()));
        return descr.Output->SendEof(TInstant::Max());
    }

    TModuleCallback WrapCallbackImpl(TServerCallback callback, const TServerOptions& serverOptions) {
        return [cb = std::move(callback), &serverOptions](const TConnDescr& descr, const TCallbackEnv& cbEnv) -> NSrvKernel::TError {
            TAtomicSharedPtr<IHttpReplyTransport> reply;
            {
                THttpRequestEnvImpl env(serverOptions, descr);

                bool tooLargeBody = false;
                Y_PROPAGATE_ERROR(env.ReceiveBody(descr, &tooLargeBody));
                if (tooLargeBody) {
                    return SendSimpleResponse(descr, "HTTP/1.1 413 Request Entity Too Large\r\n\r\n");
                }

                Y_PROPAGATE_ERROR(cb(env));
                reply = env.GetReplyTransportHolder();
            }

            if (reply) {
                Y_PROPAGATE_ERROR(TransferReply(*reply, descr, cbEnv, nullptr));
            }

            return {};
        };
    }

    TModuleCallback WrapCallbackImpl(TServerCallback callback, TThreadPool& threadPool, const TServerOptions& serverOptions) {
        return [cb = std::move(callback), &threadPool, &serverOptions](const TConnDescr& descr, const TCallbackEnv& cbEnv) -> NSrvKernel::TError {
            TString body;

            THttpRequestEnvImpl env(serverOptions, descr);

            bool tooLargeBody = false;
            Y_PROPAGATE_ERROR(env.ReceiveBody(descr, &tooLargeBody));
            if (tooLargeBody) {
                return SendSimpleResponse(descr, "HTTP/1.1 413 Request Entity Too Large\r\n\r\n");
            }

            env.GetReplyTransport();
            auto reply = env.GetReplyTransportHolder();
            Y_VERIFY(reply);

            TCallbackErrorPtr cbError = MakeAtomicShared<TCallbackError>();

            auto f = [e = std::move(env), &cb, cbError] () mutable {
                if (TError error = cb(e)) {
                    cbError->Error = std::move(error);
                    AtomicSet(cbError->HasError, 1);
                }
            };

            if (!threadPool.AddFunc(std::move(f))) {
                NSrvKernel::TResponse response;
                Y_PROPAGATE_ERROR(response.Parse("HTTP/1.1 503 Service Unavailable\r\n\r\n"));
                Y_PROPAGATE_ERROR(descr.Output->SendHead(std::move(response), false, TInstant::Max()));
                return descr.Output->SendEof(TInstant::Max());
            }

            Y_PROPAGATE_ERROR(TransferReply(*reply, descr, cbEnv, cbError));
            return {};
        };
    }

    class TStandaloneServer::TImpl {
    public:
        TImpl(const TServerCallback& callback, const TOptions& options, const TInitializationExtension& initExtension)
            : BalancerInstanceConfig_(options)
            , MainTask_(MakeHolder<NProcessCore::TMainTask>(BalancerInstanceConfig_.LuaInstanceConfig,
                BalancerInstanceConfig_.Globals,nullptr, BalancerInstanceConfig_.MainOptions, nullptr))
            , Server_(callback, MainTask_, options)
        {
            if (initExtension) {
                initExtension(MainTask_);
            }
        }

        void Run(NThreading::TPromise<void>* startPromise = nullptr) {
            with_lock (Lock_) {
                Running_ = true;
            }
            Server_.Run();

            Y_DEFER {
                Server_.Wait();
                with_lock (Lock_) {
                    Running_ = false;
                }
                CondVar_.Signal();
            };

            MainTask_->Execute(startPromise);
        }

        void Shutdown() {
            with_lock (Lock_) {
                if (!MainTask_ ) {
                    return;
                }
                MainTask_->StopMaster({
                    .CoolDown = TDuration{},
                    .Timeout = TDuration{},
                    .CloseTimeout = TDuration::Seconds(10),
                });
            }
        }

        void Stop() {
            Shutdown();
            with_lock (Lock_) {
                CondVar_.Wait(Lock_, [this] { return !Running_; });
            }
        }

    private:
        TMutex Lock_;
        TCondVar CondVar_;
        TBalancerInstanceConfig BalancerInstanceConfig_;
        TAtomicSharedPtr<NProcessCore::TMainTask> MainTask_;
        TServer Server_;
        bool Running_ = false;
    };

    TStandaloneServer::TStandaloneServer(const TServerCallback& callback, const TOptions& options,
                                         const TInitializationExtension& initExtension)
        : Impl_(MakeHolder<TImpl>(callback, options, initExtension))
    {
    }

    TStandaloneServer::~TStandaloneServer() {
    }

    void TStandaloneServer::Run(NThreading::TPromise<void>* startPromise) {
        Impl_->Run(startPromise);
    }

    void TStandaloneServer::Shutdown() {
        Impl_->Shutdown();
    }

    void TStandaloneServer::Stop() {
        Impl_->Stop();
    }

    class TServer::TImpl {
    public:
        TImpl(const TServerCallback& callback, TAtomicSharedPtr<NProcessCore::TMainTask> mainTask,
            const TServerOptions& options)
            : Options_(options)
            , BalancerRootConfig_(options)
            , CallbackHandle_(WrapCallback(callback))
            , MainTask_(std::move(mainTask))
        {
            AddCommonServerModules(ModulesFactory_);
            ModulesFactory_.AddHandle(&CallbackHandle_);

            MainTask_->AddModules(BalancerRootConfig_.LuaRootConfig, BalancerRootConfig_.Globals, ModulesFactory_, &ActiveSessions_);
        }

        void Run() {
            if (Options_.Threads) {
                Pool_.Start(Options_.Threads, Options_.QueueSize);
            }
        }

        void Wait() {
            Pool_.Stop();
        }

    private:
        TModuleCallback WrapCallback(const TServerCallback& callback) {
            auto noexceptCb = [cb = callback](THttpRequestEnv& e) -> NSrvKernel::TError {
                try {
                    Y_PROPAGATE_ERROR(cb(e));
                } Y_TRY_STORE(yexception);

                return {};
            };

            if (Options_.Threads) {
                return WrapCallbackImpl(std::move(noexceptCb), Pool_, Options_);
            } else {
                return WrapCallbackImpl(std::move(noexceptCb), Options_);
            }
        }

    private:
        TThreadPool Pool_;

        TServerOptions Options_;
        TBalancerRootConfig BalancerRootConfig_;

        TCallbackModuleHandle CallbackHandle_;
        NSrvKernel::TNodeFactory<NSrvKernel::IModule> ModulesFactory_;

        TAtomicSharedPtr<NProcessCore::TMainTask> MainTask_;

        TAtomic ActiveSessions_ = 0;
    };

    TServer::TServer(const TServerCallback& callback, TAtomicSharedPtr<NSrvKernel::NProcessCore::TMainTask> mainTask,
        const TServerOptions& options)
        : Impl_(MakeHolder<TImpl>(callback, std::move(mainTask), options))
    {
    }

    void TServer::Run() {
        Impl_->Run();
    }

    void TServer::Wait() {
        Impl_->Wait();
    }

    TServer::~TServer() {
    }
}
