#include "global.h"

bool TGlobalScheduler::HasTasks(const TString& owner) {
    auto context = Instance()->GetContext(owner);
    if (!context) {
        return false;
    }
    return (context->GetTasksCount() != 0);
}

TGlobalScheduler::TOwnerContextPtr TGlobalScheduler::GetContext(const TString& owner) {
    TReadGuard g(RegisterMutex);
    auto p = OwnerContexts.find(owner);
    return p != OwnerContexts.end() ? p->second : nullptr;
}

void TGlobalScheduler::Start() {
    INFO_LOG << "Starting MacroMatrixRebuilder" << Endl;
    Queue = ThreadsCount ? CreateRTYQueue(ThreadsCount, "GlobScheduler") : MakeHolder<TFakeThreadPool>().Release();
    Scheduler = ThreadsCount ? MakeHolder<NBus::NPrivate::TScheduler>() : nullptr;
    INFO_LOG << "Started MacroMatrixRebuilder" << Endl;
    RegisterGlobalMessageProcessor(this);
}

void TGlobalScheduler::Stop() {
    UnregisterGlobalMessageProcessor(this);
    if (Scheduler) {
        INFO_LOG << "Stopping Scheduler" << Endl;
        Scheduler->Stop();
        Scheduler.Destroy();
        INFO_LOG << "Stopped Scheduler" << Endl;
    }
    if (Queue) {
        INFO_LOG << "Stopping Queue" << Endl;
        Queue->Stop();
        Queue.Destroy();
        INFO_LOG << "Stopped Queue" << Endl;
    }
}

void TGlobalScheduler::RegisterImpl(const TString& owner) {
    TOwnerContextPtr context = MakeAtomicShared<TOwnerContext>(owner);
    {
        TWriteGuard g(RegisterMutex);
        CHECK_WITH_LOG(OwnerContexts.emplace(owner, context).second);
        INFO_LOG << "Registered owner: " << owner << Endl;
    }
    {
        TReadGuard g(RegisterMutex);
        TGuard<TMutex> guard(StartStopMutex);
        if (OwnerContexts.size() == 1) {
            Start();
        }
    }
}

void TGlobalScheduler::UnregisterImpl(const TString& owner) {
    TOwnerContextPtr context;
    {
        TWriteGuard g(RegisterMutex);
        auto p = OwnerContexts.find(owner);
        CHECK_WITH_LOG(p != OwnerContexts.end());
        context = p->second;
        CHECK_WITH_LOG(OwnerContexts.erase(owner));
        INFO_LOG << "Unregistered owner: " << owner << Endl;
    }
    if (context) {
        INFO_LOG << "Waiting for inflight tasks to complete for owner " << owner << Endl;
        TWriteGuard guard(context->Lock);
        INFO_LOG << "Tasks finished for owner " << owner << Endl;
    }
    {
        TReadGuard g(RegisterMutex);
        TGuard<TMutex> guard(StartStopMutex);
        if (OwnerContexts.empty()) {
            Stop();
        }
    }
}

void TGlobalScheduler::SetThreadsCountImpl(const ui32 threadsCount) {
    TReadGuard g(RegisterMutex);
    TGuard<TMutex> guard(StartStopMutex);
    if (OwnerContexts.size()) {
        Stop();
    }
    ThreadsCount = threadsCount;
    if (OwnerContexts.size()) {
        Start();
    }
}

bool TGlobalScheduler::Process(IMessage* message) {
    TCollectServerInfo* messCollect = dynamic_cast<TCollectServerInfo*>(message);
    if (messCollect) {
        TReadGuard g(RegisterMutex);
        for (auto&& i : OwnerContexts) {
            auto& info = messCollect->MutableAdditionalInfo("gs_info_" + i.second->Owner);
            info.InsertValue("count_schedule", i.second->Scheduled.Val());
            info.InsertValue("count_executing", i.second->Executing.Val());
            info.InsertValue("count_queued", i.second->Queued.Val());
            info.InsertValue("h_schedule", i.second->HistogramScheduled.GetReport());
            info.InsertValue("h_queued", i.second->HistogramQueued.GetReport());
            info.InsertValue("h_executing", i.second->HistogramExecuting.GetReport());
        }
        return true;
    }
    return false;
}

bool TGlobalScheduler::Schedule(THolder<IScheduledItem>&& item) {
    CHECK_WITH_LOG(item);
    auto lock = NNamedLock::TryAcquireLock(item->GetId());
    if (!lock) {
        return false;
    }
    auto owner = item->GetOwner();
    auto ctx = Instance()->GetContext(owner);
    if (!ctx) {
        return false;
    }
    Schedule(ctx, lock, std::move(item));
    return true;
}

void TGlobalScheduler::Schedule(TOwnerContextPtr context, NNamedLock::TNamedLockPtr lock, THolder<IScheduledItem>&& task) {
    CHECK_WITH_LOG(context);
    TInstant ts = task->GetScheduleTime();
    THolder<NBus::NPrivate::IScheduleItem> wrapper = MakeHolder<TOwnedScheduledItem>(context, lock, std::move(task));
    if (Instance()->Scheduler && ts > Now()) {
        Instance()->Scheduler->Schedule(wrapper.Release());
        DEBUG_LOG << "Scheduled a task for owner " << context->Owner << Endl;
    } else {
        wrapper->Do();
    }
}

void TGlobalScheduler::AddTask(TOwnerContextPtr context, NNamedLock::TNamedLockPtr lock, THolder<IObjectInQueue>&& item, THolder<IScheduledItem>&& next) {
    CHECK_WITH_LOG(context);
    THolder<IObjectInQueue> wrapper = MakeHolder<TOwnedObjectInQueue>(context, lock, std::forward<THolder<IObjectInQueue>>(item), std::forward<THolder<IScheduledItem>>(next));
    Instance()->Queue->SafeAdd(wrapper.Get());
    Y_UNUSED(wrapper.Release());
    DEBUG_LOG << "Added a task for owner " << context->Owner << Endl;
}

TGlobalScheduler::TOwnerLock TGlobalScheduler::TryLockOwner(const TString& owner) {
    TOwnerContextPtr context = Instance()->GetContext(owner);
    if (!context) {
        return nullptr;
    }

    auto lock = MakeHolder<TReadGuard>(context->Lock);
    TOwnerContextPtr recheck = Instance()->GetContext(owner);
    if (!recheck) {
        return nullptr;
    }
    return lock;
}
