#include "master_main.h"
#include "connection.h"
#include "master_process.h"
#include "multithreading.h"

#include <balancer/kernel/balancing/updater.h>
#include <balancer/kernel/helpers/syscalls.h>
#include <balancer/kernel/process/sd/sd.h>
#include <balancer/kernel/process_common/fixed_buffer_io.h>
#include <balancer/kernel/coro/coroutine.h>
#include <kernel/p0f/p0f.h>

#include <sys/resource.h>

#include <util/system/getpid.h>
#include <util/system/tls.h>

using namespace NConfig;
using namespace NSrvKernel;

namespace {
    TErrorOr<rlim_t> GetNoFileLimit() {
        struct rlimit rlim = {};
        Y_PROPAGATE_ERROR(Y_SYSCALL(getrlimit(RLIMIT_NOFILE, &rlim)));
        return rlim.rlim_cur;
    }
}

namespace NSrvKernel::NProcessCore {

    using TAcceptFull = TContListener::ICallBack::TAcceptFull;

    TMasterProcess::TMasterProcess(TMainTask* task)
        : MainTask_(task)
        , InitializerPipes_(std::exchange(task->InitializerPipes_, nullptr))
        , MainConfig_(task->Config())
        , Log_(task->Log_)
        , MasterChannel_(MakeHolder<TW2WChannel<TC2MMessage>>(task->GetCountOfChildren()))
        , MasterCpuStat_(MainTask_->ProcessStatOwner_->GetMasterStat()->GetWorkerCpuStat(0))
        , MasterExecutor_(CreateExecutor(task->MainOptions_, task->Config(), MasterCpuStat_.Get()))
        , AdminCallBack_(task->MainStats_->ConnectionsCounter(), MainConfig_, MasterExecutor_->Executor(), *this, *task, Log_)
        , StatsCallback_(task->MainStats_->ConnectionsCounter(), MainConfig_, MasterExecutor_->Executor(), task->SharedStatsManager(), *this, Log_)
        , AdminListener_(MakeHolder<TOwnListener>(MainConfig_.AdminAddrs,&AdminCallBack_,
            TOwnListener::TOptions()
                .SetListenQueue(MainConfig_.TcpListenQueue)
                .SetBindMode(MainTask_->GetBindMode())))
        , StatsListener_(MakeHolder<TOwnListener>(MainTask_->SharedStatsManager_.Addrs(),&StatsCallback_,
            TOwnListener::TOptions()
                .SetListenQueue(MainConfig_.TcpListenQueue)
                .SetBindMode(MainTask_->GetBindMode())))
        , ChildrenManager_(MakeHolder<TChildrenManager>(MasterExecutor_->Executor(), *task, *MasterChannel_, TChildrenManagerOpts{
            MainConfig_.Watchdog, MainConfig_.NChildren, MainConfig_.WorkerStartDelay, MainConfig_.WorkerStartDuration}, MainConfig_.ChildrenMask))
    {
        ExpectWorkersAfterReload_ = MainConfig_.NChildren;
        IsNewMaster_ = MainTask_->MainOptions_.ReloadConfig;
    }

    void TMasterProcess::Execute(NThreading::TPromise<void>* startPromise) {
        // in threads model we need create one writing thread for every log
        MainTask_->InitLogWritingThreads(MasterExecutor_->Executor(), /*workerId =*/ 0);

        CpuMeasureCoroutine_ = StartCpuAndTimeMeasurer(MasterExecutor_->Executor(), MasterCpuStat_.Get(), nullptr);

        MainTask_->GetProcessStatOwner().RegisterStatForProcess(MainConfig_.ChildrenMask);

        RecvLoopTask_ = TCoroutine{ECoroType::Service, "recv loop", &MasterExecutor_->Executor(), &TMasterProcess::RecvLoop, this};
        FDChecker_ = TCoroutine{ECoroType::Service, "fd_checker", &MasterExecutor_->Executor(), &TMasterProcess::FDChecker, this};

        if (InitializerPipes_) {
            HealthCheckTask_ = TCoroutine{"master_initializer_health_check", &MasterExecutor_->Executor(), &TMasterProcess::HealthCheckInitializer, this};
            InitializerCommunicationTask_ = TCoroutine{"master_initializer_comms", &MasterExecutor_->Executor(), &TMasterProcess::InitializerRead, this};
        }

        StopCont_ = TCoroutine{ECoroType::Service, "master_stop_cont", &MasterExecutor_->Executor(), [this]() {
            TGracefulShutdownOpts opts;
            if (MainTask_->StopChannel_.Receive(opts, TInstant::Max(), MasterExecutor_->Executor().Running()) == EChannelStatus::Success) {
                MasterShutDown(opts);
            }
        }};

        ReopenLogCont_ = TCoroutine{ECoroType::Service, "master_reopen_log_cont", &MasterExecutor_->Executor(), [this]() {
            char c;
            while (!MasterExecutor_->Executor().Running()->Cancelled()) {
                if (MainTask_->ReopenLogChannel_.Receive(c, TInstant::Max(), MasterExecutor_->Executor().Running()) == EChannelStatus::Success) {
                    Y_UNUSED(MainTask_->ReopenAllLogs(Cerr));
                }
            }
        }};

        BannedAddressesCleaner_ = TCoroutine{ECoroType::Service, "ban_cleaner", &MasterExecutor_->Executor(), [this]() {
            TSharedFileExistsChecker banDisabledFile;
            if (MainTask_->Config().BanAddressesDisableFile_) {
                banDisabledFile = MainTask_->GetSharedFiles().FileChecker(MainTask_->Config().BanAddressesDisableFile_,TDuration::Seconds(1));
            }
            while (!MasterExecutor_->Executor().Running()->Cancelled()) {
                if (banDisabledFile.Exists()) {
                    MainTask_->BannedAddresses().Clear();
                } else {
                    MainTask_->BannedAddresses().ProcessTtl(Now());
                }
                MasterExecutor_->Executor().Running()->SleepT(TDuration::Seconds(1));
            }
        }};

#ifdef _linux_
        if (MainConfig_.P0fEnabled) {
            NP0f::LoadBpfProgram(MainConfig_.P0fMapMaxSize);
        }
#endif

        if (!IsNewMaster_) {
            StartStatsListener();
        }
        StartAdminListener();

        ChildrenManager_->Start();

        Log_ << "Master process started" << Endl;
        Log_ << "Config MD5 " << MainTask_->GetConfigHash() << Endl;

        if (startPromise) {
            startPromise->SetValue();
        }
    }

    void TMasterProcess::FDChecker() {
        auto* const cont = MasterExecutor_->Executor().Running();

        TString filename;
        filename += "/proc/" + ::ToString(GetPID()) + "/status";

        size_t fdAllocated = 0;
        rlim_t limit = 0;
        TString data;

        while (!cont->Cancelled()) {
            Y_TRY(TError, err) {
                Y_PROPAGATE_ERROR(GetNoFileLimit().AssignTo(limit))
                MainTask_->MainStats_->NoFileLimit.Set(limit);
                return {};
            } Y_CATCH {
                err.Throw(); // let it die
            }

#if defined(_linux_)
            TFileInput fileInput(filename);
            data = fileInput.ReadAll();

            for (const auto& row : StringSplitter(data).Split('\n')) {
                if (row.Token().Contains("FDSize")) {
                    TryFromString<size_t>(row.Token().After('\t'), fdAllocated);
                    break;
                }
            }
#endif

            MainTask_->MainStats_->FDSize.Set(fdAllocated);

            cont->SleepT(TDuration::Seconds(5));
        }
    }

    void TMasterProcess::HealthCheckInitializer() {
        TCont* cont = MasterExecutor_->Executor().Running();
        while (!cont->Cancelled()) {
            if (cont->SleepT(TDuration::Seconds(10)) == ECANCELED) {
                break;
            }

            try {
                PushMessageToInitializer(cont, HealthMessage());
            } catch (...) {
                Cerr << "Initializer-Master pipe is dead, master is about to shutdown" << Endl;
                MasterShutDown();
                break;
            }
        }
    }

    void TMasterProcess::InitializerRead() {
        TCont* cont = MasterExecutor_->Executor().Running();
        while (!cont->Cancelled()) {
            TInitializerMessage msg;

            try {
                msg = PullMessageFromInitializer(cont);
            } catch (...) {
                break;
            }

            if (msg.HasShutdownMaster()) {
                MasterShutDown();
                break;
            } else if (msg.HasReloadState()) {
                LastReloadStateMsg_ = msg;
                ReloadStateMsgReceiveCV_.notify();
            } else if (msg.HasOldMasterClosedUnistat()) {
                StatsListener_->StartIfNotStarted(MasterExecutor_->Executor());
            } else if (msg.HasOldMasterExited()) {
                IsNewMaster_ = false;
            }
        }
    }

    void TMasterProcess::PushMessageToInitializer(TCont* cont, const TInitializerMessage& message) {
        SOCKET fd = InitializerPipes_->Output().GetHandle();

        TChunkPtr chunk = SerializeMessage(message);
        size_t size = chunk->Length();

        TFixedBufferContWriter::WriteBuffer(cont, fd, reinterpret_cast<char*>(&size), sizeof(size));
        TFixedBufferContWriter::WriteBuffer(cont, fd, chunk->Data(), size);
    }

    TInitializerMessage TMasterProcess::PullMessageFromInitializer(TCont* cont) {
        SOCKET fd = InitializerPipes_->Input().GetHandle();

        size_t size;
        TFixedBufferContReader::ReadBuffer(cont, fd, reinterpret_cast<char*>(&size), sizeof(size));

        TChunkPtr chunk = NewChunkReserve(size);
        chunk->Shrink(size);

        TFixedBufferContReader::ReadBuffer(cont, fd, chunk->Data(), size);
        return NSrvKernel::DeserializeMessage(chunk->Data(), size);
    }

    void TMasterProcess::MasterShutDown(const TGracefulShutdownOpts& opts) {
        if (IsTerminating_) {
            return;
        }
        IsTerminating_ = true;

        TUniversalGuard guard(MainTask_->TreesMutex);
        Y_UNUSED(guard.Lock());

        ChildrenManager_->Shutdown(opts);

        Log_ << "Master has finished waiting for workers and about to terminate itself" << Endl;

        ChildrenManager_.Destroy();

        AdminListener_.Destroy();
        StatsListener_.Destroy();

        MasterExecutor_->Executor().Abort();
    }

    void TMasterProcess::SendReloadSuccessful() {
        try {
            PushMessageToInitializer(MasterExecutor_->Executor().Running(), ReloadSuccessful());
        } catch (...) {
            Log_ << "SendReloadSuccessful error: " << CurrentExceptionMessage() << Endl;
        }
    }

    void TMasterProcess::SendCancelReload() {
        try {
            PushMessageToInitializer(MasterExecutor_->Executor().Running(), CancelReload());
        } catch (...) {
            Log_ << "SendCancelReload error: " << CurrentExceptionMessage() << Endl;
        }
    }

    void TMasterProcess::SendOldMasterClosedUnistat() {
        try {
            PushMessageToInitializer(MasterExecutor_->Executor().Running(), OldMasterClosedUnistat());
        } catch (...) {
            Log_ << "SendOldMasterClosedUnistat error: " << CurrentExceptionMessage() << Endl;
        }
    }

    TMasterProcess::TReloadConfigStatus TMasterProcess::ReloadThroughInitializer(const TReloadConfigOpts& opts) {
        TInitializerMessage reloadStartMessage;
        reloadStartMessage.MutableInitializeReload()->SetConfigPath(opts.NewConfigPath.c_str());
        for (NConfig::TGlobals::const_iterator it = opts.NewGlobals.cbegin(); it != opts.NewGlobals.cend(); ++it) {
            auto* globalArg = reloadStartMessage.MutableInitializeReload()->AddGlobalArgs();
            Y_ENSURE(it->second.IsString(), "Only string globals are supported");
            globalArg->SetKey(it->first.c_str());
            globalArg->SetValue(it->second.GetString().c_str());
        }
        reloadStartMessage.MutableInitializeReload()->SetCoroStackSize(opts.NewCoroStackSize);
        reloadStartMessage.MutableInitializeReload()->SetWorkerStartDelay(opts.WorkerStartDelay.GetValue());
        reloadStartMessage.MutableInitializeReload()->SetSaveGlobals(opts.SaveGlobals);

        TCont* cont = MasterExecutor_->Executor().Running();

        PushMessageToInitializer(cont, reloadStartMessage);

        const TInstant deadline = opts.Timeout.ToDeadLine();
        while (Now() < deadline) {
            PushMessageToInitializer(cont, GetReloadState());
            int ret = ReloadStateMsgReceiveCV_.wait(&MasterExecutor_->Executor());
            if (ret != EWAKEDUP) {
                return {
                    .ErrorMsg = TString("reload config wait failed: ") + LastSystemErrorText(ret),
                };
            }

            auto response = LastReloadStateMsg_.GetReloadState();
            if (response.GetStatus() == TReloadState::SUCCESS) {
                return {
                    .Success = true,
                };
            }
            if (response.GetStatus() == TReloadState::FAILED) {
                return {
                    .ErrorMsg = response.GetError()
                };
            }

            cont->SleepT(TDuration::MilliSeconds(100));
        }

        return {
            .Timeouted = true,
            .ErrorMsg = "reload config doesn't meet the deadline",
        };
    }

    void TMasterProcess::DoLaunchStartupAfterReloadConfigLoop() noexcept {
        --ExpectWorkersAfterReload_;

        if (ExpectWorkersAfterReload_ == 0 && IsNewMaster_ && InitializerPipes_) {
            SendReloadSuccessful();
        }
    }

    void TMasterProcess::ReopenLog(TEventData& event) {
        IOutputStream& out = event.RawOut();

        if (!MainTask_->ReopenAllLogs(out)) {
            event.SetError();
        }
    }

    void TMasterProcess::Version(TEventData& event) {
        const char* id = GetSandboxTaskId();
        event.RawOut() << GetProgramSvnVersion() << '\n' << "Sandbox task: " << (id ? id : "unknown");
    }

    void TMasterProcess::GetConfigTag(TEventData& event) noexcept {
        event.RawOut() << MainConfig_.ConfigTag;
    }

    void TMasterProcess::ResetDnsCache(TEventData&) noexcept {
        ChildrenManager_->ResetDnsCache();
        /*
        * should send to pinger as well
        * as it uses modules :/
        * */
    }

    void TMasterProcess::ShowListen(TEventData& event) noexcept {
        MainTask_->Trees.ShowListen(event);
    }

    void TMasterProcess::DbgWorkersReady(TEventData& event) noexcept {
        size_t workersReady = AtomicGet(MainTask_->LiveWorkersCounter_);
        size_t totalWorkers = MainConfig_.NChildren;

        // TODO(carzil): use json format here
        event.RawOut() << workersReady << '/' << totalWorkers << '\n';
    }

    void TMasterProcess::SandboxTaskId(TEventData& event) noexcept {
        const char* id = GetSandboxTaskId();

        // TODO(carzil): use json format here
        event.RawOut() << "Sandbox task: " << (id ? id : "unknown");
    }

    void TMasterProcess::PrintMasterPid(TEventData& event) {
        event.RawOut() << MainTask_->MainPid_ << Endl;
    }

    void TMasterProcess::ShutDown(TEventData&) {
        MasterShutDown();
    }

    void TMasterProcess::GracefulShutdown(TEventData&, const TGracefulShutdownOpts& opts) {
        MasterShutDown(opts);
    }

    void TMasterProcess::ReloadConfig(TEventData& event, const TReloadConfigOpts& opts) /* noexcept */ {
        IOutputStream& out = event.RawOut();

        TReloadConfigStatus s;

        if (IsNewMaster_) {
            event.SetError();
            out << "Unable to reload config while old master is still alive" << Endl;
            return;
        }

        if (IsReloadingConfig_) {
            event.SetError();
            out << "Reload config is already in progress" << Endl;
            return;
        }
        IsReloadingConfig_ = true;

        try {
            s = ReloadThroughInitializer(opts);
        } catch (...) {
            s.ErrorMsg = "ReloadThroughInitializer error: " + CurrentExceptionMessage();
        }

        if (s.Timeouted) {
            SendCancelReload();
        }

        if (s.Success) {
            StatsListener_.Destroy();
            SendOldMasterClosedUnistat();
            MasterShutDown(opts.GracefulShutdownOpts);
        } else {
            IsReloadingConfig_ = false;
            event.SetError();
            out << "Reload config failed: " << s.ErrorMsg << Endl;
            Log_ << "Reload config failed: " << s.ErrorMsg << Endl;
        }
    }

    void TMasterProcess::DumpSharedFiles(TEventData& event) {
        auto& jsonWriter = event.Out();
        jsonWriter.OpenArray();
        jsonWriter.OpenArray();
        MainTask_->GetSharedFiles().DumpFilesInfo(jsonWriter);
        jsonWriter.CloseArray();
        jsonWriter.CloseArray();
    }

    void TMasterProcess::CallEvent(TEventData& event, bool jsonOut) {
        ChildrenManager_->CallEvent(event, jsonOut);
    }

    void TMasterProcess::ListEventHandlers(TEventData& event) {
        ChildrenManager_->ListEventHandlers(event);
    }

    void TMasterProcess::ProcessJsonEvent(TEventData& eventData) {
        if (eventData.Event() == "dump_shared_files") {
            DumpSharedFiles(eventData);
        } else {
            CallEvent(eventData, true);
        }
    }

    void TMasterProcess::StartStatsListener() {
        StatsListener_->Start(MasterExecutor_->Executor());
    }

    void TMasterProcess::StartAdminListener() {
        AdminListener_->Start(MasterExecutor_->Executor());
    }

    TContExecutor& TMasterProcess::Executor() const noexcept {
        return MasterExecutor_->Executor();
    }

    void TMasterProcess::DoDispose() noexcept {
    }

    bool TMasterProcess::IsNeedShutdown() const noexcept {
        return ChildrenManager_->IsShutdown();
    }

    TSystemLog& TMasterProcess::Log() noexcept {
        return Log_;
    }

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

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

            std::visit(TOverloaded{
                [&](const TWorkerReady&) {
                    DoLaunchStartupAfterReloadConfigLoop();
                },
            }, message);
        }
    }
}
