#include "worker_process.h"
#include "util/generic/ptr.h"

#include <balancer/kernel/balancing/updater.h>
#include <balancer/kernel/dns/dns_helper.h>
#include <balancer/kernel/process/sd/sd.h>

#include <util/thread/singleton.h>


using namespace NConfig;
using namespace NSrvKernel;

namespace NSrvKernel::NProcessCore {

    using TAcceptFull = TContListener::ICallBack::TAcceptFull;

    void TWorkerProcess::Execute() {
        FastTlsSingleton<TThreadInfo>()->WorkerProcess = this;

        UpdateStateAfterCleanupMaster();

        if (MainTask_.IsExtendedConfigCheck()) {
            ConfigCheckTask_ = MainTask_.ConfigCheck_->Start(*this, BackendGroups_);
        }

        WorkerExecutor_->Executor().Execute();
    }

    void TWorkerProcess::InitThreadedResolver() {
        DnsCounters_ = MainTask_.ProcessStatOwner_->GetProcessStat(WorkerType())->GetDnsCounters(WorkerId());

        const bool useAsync = MainConfig().Dns.AsyncResolve;
        const auto dnsParams = DnsHelper_->CreateDnsResolver(&WorkerExecutor_->Executor(), SharedFiles(), MainConfig().Dns,
                DnsCounters_.Get(), useAsync, ThreadedQueue("resolver_queue"));
        Log_ << dnsParams << Endl;
    }

    void TWorkerProcess::InitThreadedLogs() {
        Log_ = MainTask_.Log_.GenerateThreadedLog(&WorkerExecutor_->Executor(), WorkerId());
        PingerLog_ = MainTask_.PingerLog_.GenerateThreadedLog(&WorkerExecutor_->Executor(), WorkerId());
        DynamicBalancingLog_ = MainTask_.DynamicBalancingLog_.GenerateThreadedLog(&WorkerExecutor_->Executor(), WorkerId());
        for (auto& log : MainTask_.Logs_) {
            Logs_[log.first] = MakeSimpleShared<TContextSafeLog<TLog>>(log.second->GenerateThreadedLog(&WorkerExecutor_->Executor(), WorkerId()));
        }
    }

    void TWorkerProcess::InitThreadedStuff() {
        if (CpuLimiter_) {
            CpuLimiter_->OnInit(&WorkerExecutor_->Executor(), SharedFiles());
        }

        if (MainConfig().BanAddressesDisableFile_) {
            BanAddressesDisableFileChecker_ = SharedFiles()->FileChecker(MainConfig().BanAddressesDisableFile_,TDuration::Seconds(1));
        }

        InitThreadedResolver();
        MemoryPrealloc();
    }

    void TWorkerProcess::MemoryPrealloc() {
        if (!MainConfig().TempBufPrealloc)
            return;

        // Allocate in TTempBuf memory pool predefined number of buffers
        TVector<TTempBuf> tempStorage;
        tempStorage.reserve(MainConfig().TempBufPrealloc);
        for (size_t i = 0; i < MainConfig().TempBufPrealloc; ++i) {
            tempStorage.push_back(TTempBuf{});
            constexpr size_t allocSize = 4096 * 2;
            Y_VERIFY(tempStorage.back().Size() >= allocSize);
            // Actual page size usually at least 4096, but use smaller step, just in case (we have our own container implementations)
            for (size_t j = 0; j < allocSize; j += 512) {
                *(tempStorage.back().Data() + j) = 42; // write something to actually map system memory
            }
        }
    }

    void TWorkerProcess::CloseListeners(const TGracefulShutdownOpts& opts) noexcept {
        if (!Listeners_.ListenersExists()) {
            Log_ << "Worker graceful shutdown. No listeners" << Endl;
            return;
        }

        auto* cont = WorkerExecutor_->Executor().Running();
        const TInstant cooldownDeadline = opts.CoolDown.ToDeadLine();
        const TInstant gracefulShutdownDeadline = cooldownDeadline + opts.Timeout;
        const TInstant gracefulShutdownCloseDeadline = gracefulShutdownDeadline + opts.CloseTimeout;

        if (opts.CoolDown) {
            Log_ << "Worker graceful shutdown. Cooldown has started" << Endl;
            cont->SleepD(cooldownDeadline);
        }

        if (MainConfig().ShutdownCloseUsingBPF && !opts.SkipBlockRequests) {
            if (TError error = Listeners_.CloseUsingBPF()) {
                Log_ << "Worker graceful shutdown error: " << GetErrorMessage(error);
            }
        } else if (!MainConfig().ShutdownAcceptConnections) {
            Listeners_.Destroy();
        }

        Log_ << "Worker graceful shutdown has started. Connections left:" << Connections_.Size() << Endl;

        int counter = 100;
        auto count = [this, &counter] {
            if (!--counter) {
                Log_ << "Worker graceful shutdown. Connections left:" << Connections_.Size() << Endl;
                counter = 100;
            }
        };

        while (Now() < gracefulShutdownDeadline && (Listeners_.ListenersExists() || !Connections_.Empty()) && !cont->Cancelled()) {
            while (!GracefulShutdownHTTP2Handlers_.Empty()) {
                GracefulShutdownHTTP2Handlers_.PopFront()->OnShutdown();
            }

            count();
            if (cont->SleepT(TDuration::MilliSeconds(10)) == ECANCELED) {
                break;
            }
        }

        Log_ << "Worker graceful shutdown is closing connections. Connections left:" << Connections_.Size() << Endl;
        Listeners_.Destroy();

        while (Now() < gracefulShutdownCloseDeadline && !Connections_.Empty() && !cont->Cancelled()) {
            while (!GracefulShutdownHTTP2Handlers_.Empty()) {
                GracefulShutdownHTTP2Handlers_.PopFront()->OnShutdown();
            }

            while (!GracefulShutdownHandlers_.Empty()) {
                GracefulShutdownHandlers_.PopFront()->OnShutdown();
            }

            count();
            if (cont->SleepT(TDuration::MilliSeconds(10)) == ECANCELED) {
                break;
            }
        }

        Log_ << "Worker graceful shutdown is cancelling connections. Connections left:" << Connections_.Size() << Endl;
        if (!Connections_.Empty()) {
            for (auto& connection : Connections_) {
                connection.GetCont()->Cancel();
            }

            while (!Connections_.Empty()) {
                cont->Join(Connections_.Front()->GetCont());
            }
        }

        Log_ << "Worker graceful shutdown has finished" << Endl;
    }

    void TWorkerProcess::GracefulShutdownLoop(const TGracefulShutdownOpts opts) noexcept {
        CloseListeners(opts);
        Dispose();
    }

    void TWorkerProcess::ShutDown(const TGracefulShutdownOpts& opts) noexcept {
        Log_ << "TWorkerProcess::ShutDown" << Endl;

        if (!opts.SkipBlockRequests) {
            SetBlockNewRequests();
        }

        if (Finished_) {
            return;
        }
        Finished_ = true;

        SetShuttingDown();

        GracefulShutdownCoroutine_ = TCoroutine(
            ECoroType::Service, "WorkerProcess::GracefulShutdownLoop", &WorkerExecutor_->Executor(), &TWorkerProcess::GracefulShutdownLoop, this,
            opts
        );
    }

    void TWorkerProcess::UpdateStateAfterCleanupMaster() noexcept {
        ThreadedQueueList_.Construct();

        MainTask_.Trees.Init(this);

        InitThreadedStuff();

        FuncExecutionCont_ = TCoroutine(ECoroType::Service, "func_execution_cont", &WorkerExecutor_->Executor(),
                                        &TWorkerProcess::FuncExecution, this);

        WorkerExecutor_->Executor().Create<TWorkerProcess, &TWorkerProcess::ReportChildReadyAfterConfigReload>(this, "report_child_ready");

        if (MainTask_.ChannelListener) {
            RequestProcessor_ = TCoroutine(ECoroType::Service, "request_processor", &WorkerExecutor_->Executor(),
                &TWorkerProcess::RunRequestProcessor, this);
        }

        Listeners_.Start(WorkerExecutor_->Executor());
        AtomicIncrement(MainTask_.LiveWorkersCounter_);

        Log_ << "Child started" << Endl;
    }

    void TWorkerProcess::ResetDnsCacheEvent() {
        DnsHelper_->ResetDnsCache();
    }

    void TWorkerProcess::ReportChildReadyAfterConfigReload(TCont*) noexcept {
        SendMasterMessage(TWorkerReady(), TInstant::Max());
    }

    void TWorkerProcess::SendMasterMessage(TC2MMessage message, TInstant deadLine) {
        TCont* runningCont = WorkerExecutor_->Executor().Running();
        auto waker = ThreadLocalEventWaker(runningCont);
        auto* wakerPtr = waker.Get();

        Log_ << "Sending message to master" << Endl;

        auto status = C2MChannel_.Send(message, deadLine, runningCont, wakerPtr);
        if (status == EChannelStatus::Success) {
            Log_ << "Successfully sent message to master" << Endl;
        } else {
            Log_ << "Failed to send master message" << Endl;
        }
    }

    void TWorkerProcess::CallEvent(TEventData& event, bool jsonOut) const noexcept {
        if (jsonOut) {
            EventHandler()->CallJson(event);
        } else {
            EventHandler()->Call(event);
        }
    }

    TWorkerProcess::TWorkerProcess(TMainTask& task, const TMainOptions& options, size_t workerId,
                                   TW2WChannel<TM2CMessage>& m2cChannel, TW2WChannel<TC2MMessage>& c2mChannel)
        : TBaseProcess(workerId, task)
        , MainStats_(MakeHolder<TMainStats>(*MainTask_.MainStats_, WorkerId()))
        , WorkerCpuStat_(MainTask_.ProcessStatOwner_->GetProcessStat(WorkerType())->GetWorkerCpuStat(WorkerId()))
        , WorkerScheduleCallback_(options.ContExecutionTimeLogThreshold
                                  ? MakeHolder<TWorkerScheduleCallback>(options.ContExecutionTimeLogThreshold, Log_)
                                  : nullptr)
        , WorkerExecutor_(CreateExecutor(options, task.Config(), WorkerCpuStat_.Get(), WorkerScheduleCallback_.Get()))
        , CpuLimiter_(task.CpuLimiterConfig_ ? MakeHolder<TCpuLimiter>(*task.CpuLimiterConfig_, *task.CpuLimiterStat_, WorkerId()) : nullptr)
        , CpuMeasureCoroutine_(StartCpuAndTimeMeasurer(WorkerExecutor_->Executor(), WorkerCpuStat_.Get(), CpuLimiter_.Get()))
        , RecvLoopTask_(TCoroutine(ECoroType::Service, "recv loop", &WorkerExecutor_->Executor(), &TWorkerProcess::RecvLoop, this))
        , M2CChannel_(m2cChannel)
        , C2MChannel_(c2mChannel)
        , Listeners_(this)
        , DnsHelper_(MakeHolder<NDns::THelper>())
    {
        InitThreadedLogs();
    }

    void TWorkerProcess::DumpSharedFiles(NJson::TJsonWriter& out) {
        SharedFiles()->DumpFilesInfo(out);
    }

    void TWorkerProcess::AddGracefulShutdownHandler(IGracefulShutdownHandler* handler, bool isHTTP2) {
        if (isHTTP2) {
            GracefulShutdownHTTP2Handlers_.PushBack(handler);
        } else {
            GracefulShutdownHandlers_.PushBack(handler);
        }
    }

    TWorkerProcess::~TWorkerProcess() {
        Dispose();
    }

    void TWorkerProcess::DoDispose() noexcept {
        // Stopping listeners
        Listeners_.Destroy();

        // 1. Dispose all modules
        MainTask_.Trees.Dispose(this);

        // 2. Cancel all the threaded queues
        for (auto i = ThreadedQueueList_.Begin(); i != ThreadedQueueList_.End(); ++i) {
            i->Dispose();
        }

        // 3. Join all the threaded queues
        for (auto i = ThreadedQueueList_.Begin(); i != ThreadedQueueList_.End(); ++i) {
            i->Join();
        }

        // 4. Close all log backends and flush the remaining logs
        Log_.DisposeBackend();
        PingerLog_.DisposeBackend();
        DynamicBalancingLog_.DisposeBackend();
        for (auto& i : Logs_) {
            i.second->DisposeBackend();
        }

        // 5. Cancel all coroutines in WorkerExecutor_
        if (WorkerExecutor_) {
            WorkerExecutor_->Executor().Abort();
        }

        AtomicDecrement(MainTask_.LiveWorkersCounter_);
    }

    TSharedFiles* TWorkerProcess::SharedFiles() noexcept {
        return &MainTask_.GetSharedFiles();
    }

    TThreadedQueue* TWorkerProcess::ThreadedQueue(const TString &name) noexcept {
        return ThreadedQueueList_.FindOrCreate(
            &WorkerExecutor_->Executor(),
            ThreadPool_,
            name
        );
    }

    void TWorkerProcess::AddAccepted(const TAddrDescr& addr) {
        ++Accepted_[addr];
    }

    void TWorkerProcess::AddCanceled(const TAddrDescr& addr) {
        ++Canceled_[addr];
        ++MainStats_->CancelledConns;
    }

    NSrvKernel::IPingerManager& TWorkerProcess::SharedPingerManager() noexcept {
        Y_FAIL();
    }

    TLog* TWorkerProcess::GetLog(const TString& name) {
        return Logs_[name].Get();
    }

    NProcessCore::TChildProcessType TWorkerProcess::WorkerType() const noexcept {
        return NProcessCore::TChildProcessType::Default;
    }

    NDns::IResolver& TWorkerProcess::Resolver() noexcept {
        return DnsHelper_->Resolver();
    }

    TLog* TWorkerProcess::GetDynamicBalancingLog() {
        return &DynamicBalancingLog_;
    }

    bool TWorkerProcess::IsBanned(const NAddr::IRemoteAddr& addr) const noexcept {
        if (BanAddressesDisableFileChecker_.Exists()) {
            return false;
        }
        return MainTask_.BannedAddresses().IsBanned(addr);
    }

    const TMainConfig& TWorkerProcess::MainConfig() const noexcept {
        return MainTask_.Config();
    }

    void TWorkerProcess::FuncExecution() {
        auto& channel = *MainTask_.ChildrenChannels_[WorkerId() - 1];
        auto* runningCont = WorkerExecutor_->Executor().Running();
        auto waker = ThreadLocalEventWaker(runningCont);
        auto* wakerPtr = waker.Get();
        std::function<void(IWorkerCtl*)> func;

        while (channel.Receive(func, TInstant::Max(), runningCont, wakerPtr) != EChannelStatus::Canceled) {
            func(this);
        }
    }

    void TWorkerProcess::RunRequestProcessor() {
        TU2WChannel<IChannelRequest*>& channel = MainTask_.ChannelListener->GetChannel(WorkerId());
        auto* cont = Executor().Running();
        auto waker = ThreadLocalEventWaker(cont);

        while (!cont->Cancelled()) {
            IChannelRequest* contextPtr;
            try {
                if (channel.Receive(contextPtr, TInstant::Max(), cont, waker.Get()) != EChannelStatus::Success) {
                    break;
                }
            } catch (...) {
                Log_ << CurrentExceptionMessage() << Endl;
                break;
            }

            contextPtr->RunRequest(this);
        }
    }

    void TWorkerProcess::RecvLoop() {
        auto runningCont = WorkerExecutor_->Executor().Running();
        auto waker = ThreadLocalEventWaker(runningCont);
        auto* wakerPtr = waker.Get();

        while (!runningCont->Cancelled()) {
            TM2CMessage message;
            auto status = M2CChannel_.Receive(message, TInstant::Max(), runningCont, wakerPtr);
            if (status != EChannelStatus::Success) {
                return;
            }

            std::visit(TOverloaded{
                [&](const TEvent& masterEvent) {
                    TChunksOutputStream out;
                    TStringBuf empty;
                    TEventData eventData{ empty, out };
                    eventData.SetEvent(masterEvent.Event);
                    CallEvent(eventData, masterEvent.JsonOut);
                    Y_UNUSED(masterEvent.OutputChannel->Send(ToString(out.Chunks()), TInstant::Max(), runningCont, wakerPtr));
                },

                [&](const TListEvents& toMaster) {
                    Y_UNUSED(toMaster.OutputChannel->Send(std::move(EventHandler()->List()), TInstant::Max(), runningCont, wakerPtr));
                },

                [&](const TShutDown& opts) {
                    Log_ << "Worker has received graceful shutdown message from master" << Endl;
                    ShutDown({
                        .CoolDown = opts.CoolDown,
                        .Timeout = opts.Timeout,
                        .CloseTimeout = opts.CloseTimeout,
                        .SkipBlockRequests = opts.SkipBlockRequests,
                    });
                },

                [&](const TResetDnsCache&) {
                    ResetDnsCacheEvent();
                }
            }, message);
        }
    }

    TWorkerProcess::TListenersAndCallbacks::TListenersAndCallbacks(TWorkerProcess* process) {
        for (auto& i : process->MainTask_.Trees) {
            if (!i.ListenAddrs.empty()) {
                if (TTree* tree = dynamic_cast<TTree*>(&i.Get())) {
                    ListenerCallbacks_.emplace_back(process, *tree->Entry, tree->TotalConnectionInProgress);
                    Listeners_.emplace_back(i.ListenAddrs, &ListenerCallbacks_.back(),TOwnListener::TOptions()
                        .SetSockBufSize(process->MainConfig().SockBufSize)
                        .SetListenQueue(process->MainConfig().TcpListenQueue)
                        .SetBindMode(process->MainTask_.GetBindMode())
                        .SetIgnoreBindErrors(process->MainTask_.IgnoreBindErrors_)
                        .SetP0fEnabled(process->MainConfig().P0fEnabled)
                    );
                }
            }
        }
    }

    void TWorkerProcess::TListenersAndCallbacks::Start(TContExecutor& executor) {
        for (auto& i : Listeners_) {
            i.Start(executor);
        }
    }

    TError TWorkerProcess::TListenersAndCallbacks::CloseUsingBPF() {
        for (auto& i : Listeners_) {
            if (TError error = i.CloseUsingBPF()) {
                Listeners_.clear();
                return error;
            }
        }
        return {};
    }

    void TWorkerProcess::TListenersAndCallbacks::Destroy() {
        Listeners_.clear();
    }

    bool TWorkerProcess::TListenersAndCallbacks::ListenersExists() const {
        return !Listeners_.empty();
    }
}
