#include "rty_pool.h"
#include "control.h"
#include "stats.h"
#include "messages.h"

#include <saas/library/daemon_base/protos/threading.pb.h>
#include <saas/library/daemon_base/config/thread_opts.h>
#include <saas/library/daemon_base/unistat_signals/messages.h>

#include <saas/util/smartqueue_ng.h> // TSmartMtpQueueNg constuctor (Type=HiLow)
#include <saas/util/smartqueue.h> // TSmartMtpQueue constuctor (Type=Smart)
//#include <util/generic/queue.h> // TThreadPool constuctor (Type=Std,Prio)
#include <saas/util/threadpool.h> // TPriorityThreadPool (Type=Prio,Smart,HiLow)

#include <library/cpp/mediator/messenger.h>

#include <util/system/mutex.h>
#include <util/generic/map.h>

using namespace NUtil;

class TRTYPools::TImpl final : public IMessageProcessor {
public:
    TImpl(const NThreading::TThreadingConfig& conf);
    ~TImpl();
    IThreadFactory* Get(TRTYPools::TLoadtype loadtype) const;
    const NThreading::TPoolOpts& GetOpts(TRTYPools::TLoadtype loadtype) const;
    IThreadPool* CreateThreadPool(TRTYPools::TLoadtype loadtype, IThreadFactory* pool, const TString& name);

public:
    static TMutex Mutex;
    size_t QueueDiagId = 0;

private:
    TMap<TRTYPools::TLoadtype, THolder<IThreadFactory>> Map_;
    TMap<TRTYPools::TLoadtype, THolder<NThreading::TPoolOpts>> Opts_;
    TRTYPoolsStat::TRef Stats_;
    TRTYPoolsControl Controls_; //for hi-low queues only

    TRTYPoolsStat::TRef& GetStats();

    void BindControls(TRTYPools::TLoadtype loadtype, TSmartMtpQueueNg* queue);

public:
    // IMessageProcessor
    TString Name() const override {
        return "TRTYPools";
    }

    bool Process(IMessage* message) override;
};

//
// TRTYPools::TImpl
//
TMutex TRTYPools::TImpl::Mutex;

namespace {
    inline NThreading::TPoolOpts::EPoolType GetPoolType(const NThreading::TPoolOpts& opts) {
        if (opts.GetType() != NThreading::TPoolOpts::Default)
            return opts.GetType(); //specified explicitly

        // to use TSmartMtpQueue, specity Type=Smart explicitly
        if (opts.HasLow() || opts.HasHiMult() || opts.HasHiAdd() || opts.HasStep())
            return NThreading::TPoolOpts::HiLow;
        if (opts.HasCpu() || opts.HasIo() || opts.HasNiceVl() || opts.GetNoAlloc())
            return NThreading::TPoolOpts::Prio;
        return NThreading::TPoolOpts::Std;
    }

    inline TScheduling SetCpuFlags(const NThreading::TPoolOpts& opts) {
        // we have to map our config to TScheduling from util
        TScheduling sch = TScheduling();
        if (opts.GetIo() == NThreading::TPoolOpts::IoIdle)
            sch = sch | TScheduling::IoIdle;

        switch(opts.GetCpu()) {
            case NThreading::TPoolOpts::Rt:
                return sch | TScheduling::CpuRt;
            case NThreading::TPoolOpts::NoRt:
                return sch | TScheduling::CpuNormal;
            case NThreading::TPoolOpts::Batch:
                return sch | TScheduling::CpuBatch;
            default:
                return sch | TScheduling::Default;
        }
    }

    THolder<IThreadFactory> CreateThreadPoolWithOptions(const NThreading::TPoolOpts& opts) {
        NThreading::TPoolOpts::EPoolType type = GetPoolType(opts);

        switch (type) {
            case NThreading::TPoolOpts::Std:
                return THolder<IThreadFactory>();
            case NThreading::TPoolOpts::Smart:
                return MakeHolder<TPriorityThreadPool>(SetCpuFlags(opts), opts.GetNiceVl(), TMaybe<bool>());
            case NThreading::TPoolOpts::Prio:
            case NThreading::TPoolOpts::HiLow:
                return MakeHolder<TPriorityThreadPool>(SetCpuFlags(opts), opts.GetNiceVl(), opts.HasNoAlloc() ? opts.GetNoAlloc() : TMaybe<bool>());
            default:
                ythrow yexception() << "Not implemented" << Endl;
        }
    }

    template<typename TOpt>
    TOpt Construct(const NThreading::TPoolOpts&, const TString&) {
        return TOpt();
    }

    template<>
    TSmartMtpQueue::TOptions Construct<TSmartMtpQueue::TOptions>(const NThreading::TPoolOpts& opts, const TString& threadName) {
        TSmartMtpQueue::TOptions args(opts.GetLow(), TDuration::Seconds(1), opts.GetNoAlloc());
        args.ThreadName = threadName;
        return args;
    }

    template<>
    TSmartMtpQueueNg::TOptions Construct<TSmartMtpQueueNg::TOptions>(const NThreading::TPoolOpts& opts, const TString& threadName) {
        TSmartMtpQueueNg::TOptions args;
        args.BurstJobsPerWorker = opts.GetHiQueueBurst();
        args.JobsPerWorker = opts.GetHiQueueJobs();
        args.MinThreads = opts.GetLow();
        args.MaxThreadsAdd = opts.GetHiAdd();
        args.MaxThreadsMult = opts.GetHiMult();
        args.MaxThreadsLimit = opts.GetHiLimit() != 0 ? opts.GetHiLimit() : TMaybe<size_t>();
        args.StepThreads = opts.GetStep();
        args.ThreadName = threadName;
        return args;
    }


    IThreadPool* CreateQueue(NThreading::TPoolOpts::EPoolType poolType, const NThreading::TPoolOpts& opts, IThreadFactory* pool, const TString& threadName) {
        switch(poolType) {
            case NThreading::TPoolOpts::Default:
            case NThreading::TPoolOpts::Std:
            case NThreading::TPoolOpts::Prio:
                return new TSimpleThreadPool(IThreadPool::TParams().SetFactory(pool).SetThreadName(threadName));
            case NThreading::TPoolOpts::Smart:
                return new TSmartMtpQueue(pool, Construct<NUtil::TSmartMtpQueue::TOptions>(opts, threadName));
            case NThreading::TPoolOpts::HiLow:
                return new TSmartMtpQueueNg(pool, Construct<NUtil::TSmartMtpQueueNg::TOptions>(opts, threadName));
            default:
                ythrow yexception() << "Not implemented" << Endl;
        }
    }
}

TRTYPools::TImpl::TImpl(const NThreading::TThreadingConfig& conf) {
    const auto* descr = NThreading::TLoadtype_descriptor();
    for (auto i = 0; i < descr->value_count(); i++) {
        const auto* loadType = descr->value(i);
        TString name = loadType->name(); //i.e.: Search, Indexer
        TLoadtype key = (TLoadtype)loadType->number();
        const NThreading::TPoolOpts* opts = conf.GetPoolOpts(name);
        if (opts) {
            Map_.insert(std::make_pair(key, CreateThreadPoolWithOptions(*opts)));
            Opts_.insert(std::make_pair(key, MakeHolder<NThreading::TPoolOpts>(*opts)));
        }
    }
    RegisterGlobalMessageProcessor(this);
}

TRTYPools::TImpl::~TImpl() {
    UnregisterGlobalMessageProcessor(this);
}

TRTYPoolsStat::TRef& TRTYPools::TImpl::GetStats() {
    // this->Mutex must be locked
    if (!Stats_)
        Stats_ = MakeAtomicShared<TRTYPoolsStat>();
    return Stats_;
}

inline IThreadFactory* TRTYPools::TImpl::Get(TRTYPools::TLoadtype loadtype) const {
    IThreadFactory* pool;
    with_lock (Mutex) {
        auto it = Map_.find(loadtype);
        pool = it == Map_.end() ? nullptr : it->second.Get();
    }
    return pool ? pool : SystemThreadFactory();
}

const NThreading::TPoolOpts& TRTYPools::TImpl::GetOpts(TRTYPools::TLoadtype loadtype) const {
    // Mutex must be locked
    static const NThreading::TPoolOpts defOpts;

    auto it = Opts_.find(loadtype);
    const NThreading::TPoolOpts* opts = it == Opts_.end() ? nullptr : it->second.Get();
    if (!opts)
        opts = &defOpts;
    return *opts;
}

inline IThreadPool* TRTYPools::TImpl::CreateThreadPool(TRTYPools::TLoadtype loadtype, IThreadFactory* pool, const TString& threadName) {
    Y_ASSERT(pool);
    const NThreading::TPoolOpts defOpts;
    with_lock (Mutex) {
        const NThreading::TPoolOpts& opts = GetOpts(loadtype);
        NThreading::TPoolOpts::EPoolType poolType = GetPoolType(opts);
        IThreadPool* result = CreateQueue(poolType, opts, pool, threadName);

        if (poolType == NThreading::TPoolOpts::HiLow) {
            TSmartMtpQueueNg* queue = dynamic_cast<TSmartMtpQueueNg*>(result);
            if (queue) {
                 BindControls(loadtype, queue);
            };
        }
        return result;
    }
}

void TRTYPools::TImpl::BindControls(TRTYPools::TLoadtype loadtype, TSmartMtpQueueNg* queue) {
    // Mutex must be locked
    TString name = NThreading::TLoadtype_descriptor()->FindValueByNumber(loadtype)->name();
    TRTYPoolsStat::TStatHandlerRef handler = TRTYPoolsStat::CreateHandler(GetStats(), name, ToString(QueueDiagId++));
    queue->SetStatHandler(handler);

    TSmartMtpQueueNg::TUpdatableRef updatable = queue->GetUpdatable();
    Controls_.Add(name, updatable);
}

bool TRTYPools::TImpl::Process(IMessage* message) {
    if (auto* m = dynamic_cast<TMessageUpdateUnistatSignals*>(message)) {
        TRTYPoolsStat::TRef stats = Stats_;
        if (stats) {
            stats->Purge();
        }
        return true;
    }
    if (auto* m = dynamic_cast<TMessageGetQueuePerf*>(message)) {
        TRTYPoolsStat::TRef stats = Stats_;
        if (stats) {
            stats->Process(m);
            return true;
        }
    }
    if (auto* m = dynamic_cast<TMessageSetQueueOpts*>(message)) {
        const auto* descr = NThreading::TLoadtype_descriptor();
        const auto* descrLt = descr->FindValueByName(m->QueueName);
        if (descrLt == nullptr) {
            ythrow yexception() << "bad queue name";
        }

        const NThreading::TLoadtype loadtype = NThreading::TLoadtype(descrLt->number());
        const NThreading::TPoolOpts& opts = GetOpts(loadtype);
        const NThreading::TPoolOpts usrOpts = NThreading::TThreadingConfig::ParsePoolOpts(m->OptionsStr);
        const auto& res = usrOpts;
        Opts_[loadtype] = MakeHolder<NThreading::TPoolOpts>(res);

        const NThreading::TPoolOpts::EPoolType oldPoolType = GetPoolType(opts);
        const NThreading::TPoolOpts::EPoolType newPoolType = GetPoolType(res);
        const bool isLiveUpdate = newPoolType == oldPoolType && newPoolType == NThreading::TPoolOpts::HiLow;
        if (!isLiveUpdate) {
            return true;
        } else {
            const TString& key = descrLt->name();
            auto hiLowOpts = Construct<NUtil::TSmartMtpQueueNg::TOptions>(res, "");
            return Controls_.Update(key, hiLowOpts);
        }
    }
    return false;
}

//
// TRTYPools
//
TRTYPools::TRTYPools() {
}

TRTYPools::~TRTYPools() {
}

void TRTYPools::Initialize(const NThreading::TThreadingConfig& conf) {
    if (conf.GetPresetName() != "RTY")
        return;

    with_lock (TImpl::Mutex) {
        Impl_.Reset(new TRTYPools::TImpl(conf));
    }
}

void TRTYPools::Destroy() {
    Impl_.Destroy();
}

IThreadFactory* TRTYPools::Get(TLoadtype loadtype) {
    return Impl_ ? Impl_->Get(loadtype) : SystemThreadFactory();
}

IThreadPool* TRTYPools::CreateThreadPool(TLoadtype loadtype, IThreadFactory* pool, size_t threadCount, size_t maxQueueSize, const TString& threadName) {
    THolder<IThreadPool> queue;
    if (Y_LIKELY(Impl_)) {
        queue.Reset(Impl_->CreateThreadPool(loadtype, pool ? pool : Impl_->Get(loadtype), threadName));
    } else {
        queue = MakeHolder<TSimpleThreadPool>(IThreadPool::TParams().SetFactory(pool ? pool : SystemThreadFactory()).SetThreadName(threadName));
    }

    queue->Start(threadCount, maxQueueSize);
    return queue.Release();
}



