#include "main.h"

#include "initializer_message.h"
#include <balancer/kernel/process/initializer_io.h>
#include "master_main.h"
#include "select_poller.h"

#include <balancer/kernel/memory/chunks.h>
#include <balancer/kernel/module/iface.h>
#include <balancer/kernel/process_common/main_options.h>

#include <library/cpp/getopt/opt.h>
#include <library/cpp/svnversion/svnversion.h>

#include <util/generic/ptr.h>
#include <util/generic/queue.h>
#include <util/generic/scope.h>
#include <util/generic/strbuf.h>
#include <util/generic/string.h>
#include <util/string/builder.h>
#include <util/stream/file.h>
#include <util/system/getpid.h>
#include <util/system/mlock.h>

#include <tuple>
#include <utility>

#include <sys/types.h>
#include <sys/wait.h>

#ifdef _linux_
#include <sys/prctl.h>
#endif


#define BALANCER_INITIALIZER_DBG_LOG(msg) if (Y_UNLIKELY(DebugEnabled_)) {\
    Cerr << GetPID() << ' ' << msg << Endl; \
}
#define BALANCER_INITIALIZER_ERR(msg) {\
    Cerr << GetPID() << ' ' << msg << Endl; \
}


namespace NSrvKernel {

    struct TMasterInfo: TNonCopyable {
        TMasterInfo(pid_t pid, TMasterStreamPtr io)
            : Pid(pid)
            , IO(std::move(io))
        {}

        const pid_t Pid = 0;
        TMasterStreamPtr IO;
    };

    struct TNewMasterArgs {
        TString ConfigPath;
        NConfig::TGlobals Globals;
        size_t CoroStackSize;
        TDuration WorkerStartDelay;
    };

    static TStringBuf GetMasterName(bool isNewMaster) {
        return isNewMaster ? "New master" : "Master";
    }

    class TInitializerManager {
        enum class EPipeEvent {
            UNEXPECTED, FROM_MASTER, FROM_NEW_MASTER, TO_MASTER, TO_NEW_MASTER, SIGNALED_EXIT_REQUIRED
        };

    public:
        TInitializerManager(const TString& cfgPath, TGlobals& globals, IModuleFactory* loader, const TMainOptions& mainOptions)
            : ConfigPath_(cfgPath)
            , Globals_(&globals)
            , Loader_(loader)
            , MainOptions_(&mainOptions)
            , Poller_(new NSrvKernel::TSelectPoller())
            , DebugEnabled_(MainOptions_->DebugInitializer)
        {
            TPipe::Pipe(ExitRequiredPipes_.Input(), ExitRequiredPipes_.Output());
            ExitRequiredPipes_.Init();
            Poller_->WaitRead(ExitRequiredPipes_.Input().GetHandle(), reinterpret_cast<void*>(&ExitRequiredPipes_.Input()));

            INSTANCE = this;

            struct sigaction act;
            memset(&act, 0, sizeof(act));
            act.sa_handler = TInitializerManager::SignalHandler;

            for (auto signo: {SIGINT, SIGTERM}) {
                if (sigaction(signo, &act, nullptr) != 0) {
                    ythrow TSystemError() << "failed to setup signal handler";
                }
            }
        }

        ~TInitializerManager() {
            struct sigaction act;
            memset(&act, 0, sizeof(act));
            act.sa_handler = SIG_IGN;

            for (auto signo: {SIGINT, SIGTERM}) {
                if (sigaction(signo, &act, nullptr) != 0) {
                    BALANCER_INITIALIZER_ERR("Failed to reset signal handler");
                }
            }
        }

        int RunInitializer() {
            Master_ = ForkMaster();
            if (IsForkedMaster_) {
                return ForkedMasterError_;
            }

            for (;;) {
                Y_VERIFY(!IsForkedMaster_, "Master process should not run initializer loop");
                TrySwapMaster();

                if (!Master_) {
                    BALANCER_INITIALIZER_ERR("Lost master")
                    return 1;
                }

                const EPipeEvent pipeEvent = Poll();

                switch (pipeEvent) {
                case EPipeEvent::FROM_MASTER:
                case EPipeEvent::FROM_NEW_MASTER: {
                    const auto res = HandleEventFromMaster(pipeEvent == EPipeEvent::FROM_NEW_MASTER);
                    if (res) {
                        return res.GetRef();
                    }
                    break;
                }

                case EPipeEvent::TO_MASTER:
                case EPipeEvent::TO_NEW_MASTER: {
                    THolder<TMasterInfo>& master = (pipeEvent == EPipeEvent::TO_NEW_MASTER) ? NewMaster_ : Master_;

                    const EIoResult result = master->IO->Write();
                    switch (result) {
                        case EIoResult::COMPLETED:
                            DisposeOutputPoller(*master);
                            break;
                        case EIoResult::PARTIALLY_COMPLETED:
                            break;
                        case EIoResult::EOF_OCCURED: {
                            const auto res = OnMasterEOF(pipeEvent == EPipeEvent::TO_NEW_MASTER);
                            if (res) {
                                return res.GetRef();
                            }
                            break;
                         }
                    }
                    break;
                }

                case EPipeEvent::SIGNALED_EXIT_REQUIRED: {
                    BALANCER_INITIALIZER_DBG_LOG("Exit required triggered");
                    char inCh;
                    ExitRequiredPipes_.Input().Read(&inCh, 1);
                    PerformExitProcedure();
                    break;
                }

                case EPipeEvent::UNEXPECTED:
                    BALANCER_INITIALIZER_ERR("Poller seems to be configured improperly");
                    return 1;
                }
            }
        }

        void RegisterSignaledExitRequired() {
            if (!IsForkedMaster_) {
                Y_SCOPE_EXIT(Errno_ = errno) {
                    errno = Errno_;
                };
                char ch = 'e';
                // Output().Write() may throw exceptions => calls for allocation functions,
                // which are not signal-safe
                auto pipeHandle = ExitRequiredPipes_.Output().GetHandle();
                write(pipeHandle, &ch, 1);
            }
        }

        void PerformExitProcedure() {
            if (!ShutdownStarted_) {
                ShutdownStarted_ = true;
                SendMessageToMaster(*Master_, ShutdownMaster(TShutdownMaster::INITIALIZER_SIGTERM));
                if (NewMaster_) {
                    SendMessageToMaster(*NewMaster_, ShutdownMaster(TShutdownMaster::INITIALIZER_SIGTERM));
                }
            }
        }

        static void SignalHandler(int signal) {
            Y_UNUSED(signal);
            INSTANCE->RegisterSignaledExitRequired();
        }

    private:
        void SendMessageToMaster(TMasterInfo& master, const TInitializerMessage& message) {
            master.IO->EnqueueWrite(message);
            SetupOutputPoller(master);
        }

        bool WaitAndFinalize(THolder<TMasterInfo>& master, bool isNewMaster) {
            if (master) {
                int exitStatus;
                pid_t pid = master->Pid;
                int ret;
                do {
                    ret = waitpid(pid, &exitStatus, 0);
                } while (ret < 0 && errno == EINTR);
                if (ret >= 0) {
                    DisposeInputPoller(*master);
                    DisposeOutputPoller(*master);
                    master.Reset();
                    if (ret > 0 && WIFEXITED(exitStatus)) {
                        int status = WEXITSTATUS(exitStatus);
                        if (status > 0) {
                            BALANCER_INITIALIZER_ERR(GetMasterName(isNewMaster) << " " << pid << " exited with code " << status);
                            return true;
                        }
                    } else if (ret > 0 && WIFSIGNALED(exitStatus)) {
                        int sig = WTERMSIG(exitStatus);
                        BALANCER_INITIALIZER_ERR(GetMasterName(isNewMaster) << " " << pid << " crashed with signal " << sig);
                        return true;
                    } else {
                        BALANCER_INITIALIZER_ERR(GetMasterName(isNewMaster) << " " << pid << " destroyed");
                        return true;
                    }
                } else {
                    ythrow TSystemError() << "waitpid failed";
                }
            }
            return false;
        }

        TMaybe<bool> OnMasterEOF(bool isNewMaster) {
            BALANCER_INITIALIZER_DBG_LOG(GetMasterName(isNewMaster) << " EOF");

            THolder<TMasterInfo>& master = isNewMaster ? NewMaster_ : Master_;

            const bool error = WaitAndFinalize(master, isNewMaster);
            if (!Master_ && !NewMaster_) {
                return error;
            } else if (Master_ && isNewMaster && ReloadState_.GetReloadState().GetStatus() != TReloadState::FAILED) {
                ReloadState_ = ReloadFailed("New master exited");
            }
            return Nothing();
        }

        TMaybe<bool> HandleEventFromMaster(bool isNewMaster) {
            THolder<TMasterInfo>& master = isNewMaster ? NewMaster_ : Master_;

            TInitializerMessage message;
            const EIoResult result = master->IO->Read(&message);
            switch (result) {
                case EIoResult::PARTIALLY_COMPLETED:
                    BALANCER_INITIALIZER_DBG_LOG("From " << GetMasterName(isNewMaster) << " partially completed");
                    break;
                case EIoResult::EOF_OCCURED: {
                    const auto res = OnMasterEOF(isNewMaster);
                    if (res) {
                        return res.GetRef();
                    }
                    break;
                }
                case EIoResult::COMPLETED: {
                    BALANCER_INITIALIZER_DBG_LOG("From " << GetMasterName(isNewMaster) << " completed: " << message.AsJSON());
                    bool valid = true;

                    if (isNewMaster) { // Messages from new master
                        if (message.HasReloadState() && IsReloadCompleted(message.GetReloadState())) {
                            ReloadState_ = std::move(message);
                        } else if (!message.HasHealthMessage()) {
                            valid = false;
                        }
                    } else { // Messages from master
                        if (message.HasInitializeReload()) {
                            if (NewMaster_) {
                                static const TString error = "Initializer got reload config request, but a config reload is already in progress";
                                BALANCER_INITIALIZER_ERR(error);
                                ReloadState_ = ReloadFailed(error);
                            } else {
                                ReloadState_ = ReloadInProgress();

                                TGlobals globals;
                                if (message.GetInitializeReload().GetSaveGlobals()) {
                                    globals = *Globals_;
                                }
                                for (const auto& arg : message.GetInitializeReload().GetGlobalArgs()) {
                                    globals[arg.GetKey()] = arg.GetValue();
                                }
                                const size_t coroStackSize = message.GetInitializeReload().GetCoroStackSize();

                                TNewMasterArgs args{
                                    message.GetInitializeReload().GetConfigPath(),
                                    globals,
                                    coroStackSize,
                                    TDuration::FromValue(message.GetInitializeReload().GetWorkerStartDelay()),
                                };
                                NewMaster_ = ForkMaster({args});
                                if (IsForkedMaster_) {
                                    return ForkedMasterError_;
                                }
                            }
                        } else if (message.HasGetReloadState()) {
                            SendMessageToMaster(*Master_, ReloadState_);
                        } else if (message.HasCancelReload()) {
                            if (NewMaster_) {
                                SendMessageToMaster(*NewMaster_, ShutdownMaster(TShutdownMaster::RELOAD_CANCELED));
                            }
                        } else if (message.HasOldMasterClosedUnistat()) {
                            if (NewMaster_) {
                                SendMessageToMaster(*NewMaster_, OldMasterClosedUnistat());
                            } else {
                                BALANCER_INITIALIZER_ERR("New master exited after old master closed unistat");
                            }
                        } else if (!message.HasHealthMessage()) {
                            valid = false;
                        }
                    }

                    if (!valid) {
                        BALANCER_INITIALIZER_ERR("Unexpected message from " << GetMasterName(isNewMaster) << " : " << message.AsJSON());
                    }
                    break;
                }
            }
            return Nothing();
        }

        void TrySwapMaster() {
            if (!Master_ && NewMaster_) {
                BALANCER_INITIALIZER_DBG_LOG("Swap master...");
                Master_.Swap(NewMaster_); // Le Roi est mort, vive le Roi!

                BALANCER_INITIALIZER_DBG_LOG("Pushing reload completed message..");
                SendMessageToMaster(*Master_, OldMasterExited());
            }
        }

        void SetupInputPoller(const TMasterInfo& master) {
            Poller_->WaitRead(master.IO->GetInputHandle(), master.IO->GetInputCookie());
        }

        void DisposeInputPoller(const TMasterInfo& master) {
            Poller_->Unwait(master.IO->GetInputHandle());
        }

        void SetupOutputPoller(const TMasterInfo& master) {
            Poller_->WaitWrite(master.IO->GetOutputHandle(), master.IO->GetOutputCookie());
        }

        void DisposeOutputPoller(const TMasterInfo& master) {
            Poller_->Unwait(master.IO->GetOutputHandle());
        }

        EPipeEvent Poll() {
            void *cookie = Poller_->WaitI();

            if (Master_ && cookie == Master_->IO->GetInputCookie()) {
                BALANCER_INITIALIZER_DBG_LOG("Poll event: FROM_MASTER");
                return EPipeEvent::FROM_MASTER;
            } else if (Master_ && cookie == Master_->IO->GetOutputCookie()) {
                BALANCER_INITIALIZER_DBG_LOG("Poll event: TO_MASTER");
                return EPipeEvent::TO_MASTER;
            } else if (NewMaster_ && cookie == NewMaster_->IO->GetInputCookie()) {
                BALANCER_INITIALIZER_DBG_LOG("Poll event: FROM_NEW_MASTER");
                return EPipeEvent::FROM_NEW_MASTER;
            } else if (NewMaster_ && cookie == NewMaster_->IO->GetOutputCookie()) {
                BALANCER_INITIALIZER_DBG_LOG("Poll event: TO_NEW_MASTER");
                return EPipeEvent::TO_NEW_MASTER;
            } else if (cookie == &ExitRequiredPipes_.Input()) {
                BALANCER_INITIALIZER_DBG_LOG("Poll event: SIGNALED_EXIT_REQUIRED");
                return EPipeEvent::SIGNALED_EXIT_REQUIRED;
            } else {
                BALANCER_INITIALIZER_ERR("Poll event: will return UNEXPECTED, got cookie: " << size_t(cookie));
                return EPipeEvent::UNEXPECTED;
            }
        }

        THolder<TMasterInfo> ForkMaster(TMaybe<TNewMasterArgs> args = TMaybe<TNewMasterArgs>()) {
            const bool reload = args.Defined();
            auto [masterStream, initializerStream] = SetupMasterPipes();

            const pid_t pid = fork();

            if (pid > 0) {
                // init process
                initializerStream.Reset();
                auto master = MakeHolder<TMasterInfo>(pid, std::move(masterStream));
                SetupInputPoller(*master);
                return master;
            } else if (pid == 0) {
                IsForkedMaster_ = true;

                // master process
                try {
                    masterStream.Reset();

                    Poller_.Reset();
                    ExitRequiredPipes_.Close();

                    TGlobals globals = *Globals_;
                    TMainOptions mainOptions = *MainOptions_;
                    TString configPath = ConfigPath_;

                    if (reload) {
                        globals = args->Globals;
                        mainOptions = TMainOptions();
                        mainOptions.ReloadConfig = true;
                        mainOptions.WorkerStartDelay = args->WorkerStartDelay;
                        if (args->CoroStackSize != 0) {
                            mainOptions.SetCoroStackSize(args->CoroStackSize);
                        }
                        configPath = args->ConfigPath;
                    }

                    TString config = TFileInput(configPath).ReadAll();

                    NSrvKernel::NProcessCore::RunMasterMain(globals, config, Loader_, mainOptions, initializerStream.Get());
                } catch (...) {
                    ForkedMasterError_ = true;
                    TString ex = CurrentExceptionMessage();
                    Cerr << ex << Endl;
                }

                return {};
            } else {
                ythrow TSystemError() << "fork failed";
            }
        }

    private:
        static TInitializerManager* INSTANCE;

        const TString ConfigPath_;
        TGlobals* Globals_ = nullptr;
        IModuleFactory* Loader_ = nullptr;
        const TMainOptions* MainOptions_ = nullptr;

        THolder<TMasterInfo> Master_;
        THolder<TMasterInfo> NewMaster_;

        TTwoNonblockingPipes ExitRequiredPipes_;

        THolder<NSrvKernel::TSelectPoller> Poller_;

        TInitializerMessage ReloadState_;

        bool IsForkedMaster_ = false;
        bool ForkedMasterError_ = false;
        bool ShutdownStarted_ = false;

        bool DebugEnabled_ = false;
    };

    TInitializerManager* TInitializerManager::INSTANCE = nullptr;

    void PrintVersionAndExit() {
        Cerr << PROGRAM_VERSION << Endl;
        exit(0);
    }

    int RunMain(const TString& cfgPath, TGlobals& globals, INodeFactory<IModule>* loader, const TMainOptions& options) {
        TInitializerManager initializer(cfgPath, globals, loader, options);

        return initializer.RunInitializer();
    }

    int RunMain(int argc, char** argv, INodeFactory<IModule>* loader, TMainOptions& options) {
        using NLastGetopt::TOpt;
        using NLastGetopt::TOpts;
        using NLastGetopt::TOptsParseResult;

        TGlobals globals;

        TMaybe<size_t> coroutineStackSize;

        TOpts opts;
        opts.AddHelpOption('h');
        opts.AddLongOption('v', "version", "balancer version").NoArgument().Handler(&PrintVersionAndExit);

        TVector<TString> globalKeysAndValues;

        opts.SetFreeArgsMin(0);
        opts.SetFreeArgsMax(1);
        opts.SetFreeArgTitle(0, "CONFIG_FILE", "balancer config file");

        opts.AddCharOption('V', "substitute global key in config with provided value, format: key=value")
            .AppendTo(&globalKeysAndValues).RequiredArgument();
        opts.AddCharOption('C', "coroutine stack size")
            .StoreResultT<size_t>(&coroutineStackSize).RequiredArgument();
        opts.AddCharOption('P', "use specified poller")
            .StoreResult(&options.Poller).RequiredArgument();
        opts.AddCharOption('K', "do not start balancer, just validate config")
            .StoreValue(&options.JustCheckConfig, true).NoArgument();
        opts.AddLongOption("skip-backend-checks", "skip backend checks in config validation")
            .StoreValue(&options.SkipBackendChecks, true).NoArgument();
        opts.AddLongOption("skip-ip-checks", "skip ip checks in config validation")
            .StoreValue(&options.SkipIpChecks, true).NoArgument();
        opts.AddLongOption("skip-extended-checks", "perform basic config validation onlyskip all extended checks in config validation")
            .StoreValue(&options.SkipExtendedChecks, true).NoArgument();
        opts.AddCharOption('S', "list of signals")
            .SetFlag(&options.JustDumpSignals).NoArgument();
        opts.AddCharOption('I', "output initializer debug messages")
            .StoreValue(&options.DebugInitializer, true).NoArgument();
        opts.AddCharOption('L', "lock memory of executable")
            .StoreValue(&options.MemLock, true).NoArgument();

        TOptsParseResult res(&opts, argc, argv);

        for (const auto& keyAndValue : globalKeysAndValues) {
            TStringBuf value = keyAndValue;
            TMaybe<TString> key;
            GetNext(value, '=', key);

            if (!key) {
                ythrow TConfigParseError() << "-V argument must be in key=value format";
            }

            globals[*key] = value;
        }
        if (coroutineStackSize) {
            options.SetCoroStackSize(*coroutineStackSize);
        }

        if (options.MemLock) {
            try {
                LockAllMemory(LockCurrentMemory);
            } catch (...) {
                Cerr << "Can't lock self memory: " << CurrentExceptionMessage() << Endl;
            }
        }

        TVector<TString> freeArgs = res.GetFreeArgs();

        if (!freeArgs.empty()) {
            TString cfgPath = freeArgs[0];

            globals[TString("BALANCER_CFG_NAME")] = cfgPath;

            return RunMain(cfgPath, globals, loader, options);
        } else {
            Cerr << PROGRAM_VERSION << Endl;

            ythrow TConfigParseError() << "Usage: " << argv[0] << " [-V key=value]* [-C coro_stack_size] [-K] CONFIG_FILE)";
        }
    }

    int RunMain(int argc, char** argv, IModuleFactory* loader) {
#ifdef _linux_
        // setcap on balancer executable disables cores
        // see BALANCER-1650
        prctl(PR_SET_DUMPABLE, 1);
#endif

        TMainOptions options;
        return RunMain(argc, argv, loader, options);
    }
}

#undef BALANCER_INITIALIZER_DBG_LOG
#undef BALANCER_INITIALIZER_ERR
