#pragma once

#include <saas/util/named_lock.h>
#include <saas/util/queue.h>

#include <library/cpp/messagebus/scheduler/scheduler.h>
#include <library/cpp/logger/global/global.h>

#include <util/generic/map.h>
#include <util/system/rwlock.h>
#include <saas/library/daemon_base/daemon/messages.h>
#include <library/cpp/mediator/messenger.h>
#include <saas/library/histogram/time_slide.h>

class TGlobalScheduler: public IMessageProcessor {
public:
    class IScheduledItem {
    public:
        virtual ~IScheduledItem() {}

        virtual const TString& GetId() const = 0;
        virtual const TString& GetOwner() const = 0;
        virtual TInstant GetScheduleTime() const = 0;
        virtual THolder<IObjectInQueue> GetQueuedItem() const = 0;
        virtual THolder<IScheduledItem> GetNextScheduledItem(TInstant now) const = 0;
    };

    template <class T>
    class TScheduledItem
        : public IScheduledItem
        , public IObjectInQueue
    {
    private:
        using TThis = T;

    public:
        TScheduledItem(const TString& owner, const TString& id, TInstant scheduleTime)
            : Owner(owner)
            , Id(id)
            , ScheduleTime(scheduleTime)
        {
        }

        virtual const TString& GetId() const override {
            return Id;
        }
        virtual const TString& GetOwner() const override {
            return Owner;
        }
        virtual TInstant GetScheduleTime() const override {
            return ScheduleTime;
        }

        virtual THolder<IObjectInQueue> GetQueuedItem() const override {
            return MakeHolder<TThis>(*static_cast<const TThis*>(this));
        }
        virtual THolder<IScheduledItem> GetNextScheduledItem(TInstant /*now*/) const override {
            return nullptr;
        }

    private:
        const TString Owner;
        const TString Id;
        const TInstant ScheduleTime;
    };

private:
    struct TOwnerContext {
        const TString Owner;
        mutable TRWMutex Lock;

        TTimeSlidedHistogram HistogramExecuting;
        TTimeSlidedHistogram HistogramScheduled;
        TTimeSlidedHistogram HistogramQueued;

        TAtomicCounter Executing;
        TAtomicCounter Scheduled;
        TAtomicCounter Queued;

        ui64 GetTasksCount() const {
            return Executing.Val() + Scheduled.Val() + Queued.Val();
        }

        TOwnerContext(const TString& owner)
            : Owner(owner)
            , HistogramExecuting(10, 300, 0, 100, 10)
            , HistogramScheduled(10, 300, 0, 100, 10)
            , HistogramQueued(10, 300, 0, 100, 10)
        {
        }
    };
    using TOwnerContextPtr = TAtomicSharedPtr<TOwnerContext>;
    using TOwnerContexts = TMap<TString, TOwnerContextPtr>;
    using TOwnerLock = THolder<TReadGuard>;

    class TOwnedScheduledItem: public NBus::NPrivate::IScheduleItem {
    private:
        const TOwnerContextPtr Context;
        NNamedLock::TNamedLockPtr IdLock;
        THolder<IScheduledItem> Item;
        TGuard<TAtomicCounter> Counter;
        TTimeSlidedHistogram::TGuard GuardHistogramScheduled;
    public:
        TOwnedScheduledItem(TOwnerContextPtr context, NNamedLock::TNamedLockPtr lock, THolder<IScheduledItem>&& item)
            : NBus::NPrivate::IScheduleItem(item->GetScheduleTime())
            , Context(context)
            , IdLock(lock)
            , Item(std::forward<THolder<IScheduledItem>>(item))
            , Counter(context->Scheduled)
            , GuardHistogramScheduled(context->HistogramScheduled)
        {
        }
        virtual void Do() override {
            if (auto lock = TryLockOwner(Context->Owner)) {
                if (Item) {
                    auto queued = Item->GetQueuedItem();
                    auto next = Item->GetNextScheduledItem(Now());
                    AddTask(Context, IdLock, std::move(queued), std::move(next));
                }
            }
        }
    };

    class TOwnedObjectInQueue: public IObjectInQueue {
    private:
        const TOwnerContextPtr Context;
        NNamedLock::TNamedLockPtr IdLock;
        THolder<IObjectInQueue> Object;
        THolder<IScheduledItem> Next;
        TGuard<TAtomicCounter> Counter;
        TTimeSlidedHistogram::TGuard GuardHistogramEnqueued;
    public:
        TOwnedObjectInQueue(TOwnerContextPtr context, NNamedLock::TNamedLockPtr lock, THolder<IObjectInQueue>&& object, THolder<IScheduledItem>&& next)
            : Context(context)
            , IdLock(lock)
            , Object(std::forward<THolder<IObjectInQueue>>(object))
            , Next(std::forward<THolder<IScheduledItem>>(next))
            , Counter(context->Queued)
            , GuardHistogramEnqueued(context->HistogramQueued)
        {
        }
        virtual void Process(void* threadSpecificResource) override {
            THolder<TOwnedObjectInQueue> cleanup(this);
            if (auto lock = TryLockOwner(Context->Owner)) {
                GuardHistogramEnqueued.Release();
                if (Object) {
                    TTimeSlidedHistogram::TGuard g(Context->HistogramExecuting);
                    TGuard<TAtomicCounter> counter(Context->Executing);
                    Object.Release()->Process(threadSpecificResource);
                }
                if (Next) {
                    Schedule(Context, IdLock, std::move(Next));
                }
            }
        }
    };

private:
    THolder<NBus::NPrivate::TScheduler> Scheduler;
    THolder<IThreadPool> Queue;
    TOwnerContexts OwnerContexts;
    TRWMutex RegisterMutex;
    TMutex StartStopMutex;

    ui32 ThreadsCount = 4;

private:
    TOwnerContextPtr GetContext(const TString& owner);
    void Start();
    void Stop();

    void RegisterImpl(const TString& owner);
    void UnregisterImpl(const TString& owner);
    void SetThreadsCountImpl(const ui32 threadsCount);

public:

    TGlobalScheduler() {
    }

    ~TGlobalScheduler() {
        CHECK_WITH_LOG(OwnerContexts.empty());
    }

    virtual bool Process(IMessage* message) override;

    virtual TString Name() const override {
        return "TGlobalScheduler_" + ToString((ui64)this);
    }

    static TGlobalScheduler* Instance() {
        return Singleton<TGlobalScheduler>();
    }

    static void SetThreadsCount(const ui32 threadsCount) {
        Instance()->SetThreadsCountImpl(threadsCount);
    }

    static void Unregister(const TString& owner) {
        return Instance()->UnregisterImpl(owner);
    }

    static void Register(const TString& owner) {
        return Instance()->RegisterImpl(owner);
    }

    static bool Schedule(THolder<IScheduledItem>&& item);
    static bool HasTasks(const TString& owner);

private:
    static void Schedule(TOwnerContextPtr context, NNamedLock::TNamedLockPtr lock, THolder<IScheduledItem>&& task);
    static void AddTask(TOwnerContextPtr context, NNamedLock::TNamedLockPtr lock, THolder<IObjectInQueue>&& item, THolder<IScheduledItem>&& next);

    static TOwnerLock TryLockOwner(const TString& owner);
};
