#include <balancer/kernel/connection_manager/connection_manager.h>

namespace NConnectionManager {
    TConnectionManager::TOptions::TOptions(const NSrvKernel::ICtl& ctl, NConfig::IConfig& config) {
        config.ForEach(this);
        WorkersCount = ctl.GetCountOfChildren();
    }

    TConnectionManager::TConnectionManager(TOptions options)
        : Options_{std::move(options)}
    {
        Y_VERIFY(Options_.WorkersCount > 0);
    }

    TConnectionManager::TConnectionManager(NSrvKernel::ICtl const& ctl, NConfig::IConfig& config)
        : TConnectionManager{TOptions{ctl, config}}
    {
    }

    size_t TConnectionManager::PreferredChannelCapacity() const noexcept {
        return Options_.WorkersCount * 2;
    }

    TConnectionPool* TConnectionManager::AddConnectionPool(size_t id, size_t limit, NSrvKernel::EPollMode pollMode, TDuration keepaliveTimeout, TDuration connectTimeout) {
        if (Options_.IgnorePerPoolLimit) {
            limit = Max<size_t>();
        }
        with_lock(Mutex_) {
            auto ret = ConnectionPoolMap_.emplace(id, MakeIntrusive<TConnectionPool>(
                    *this, Options_.WorkersCount + 1, limit, pollMode, keepaliveTimeout, connectTimeout));
            Y_VERIFY(ret.second);
            return ret.first->second.Get();
        }
    }

    void TConnectionManager::RemoveConnectionPool(size_t id) {
        with_lock(Mutex_) {
            auto it = ConnectionPoolMap_.find(id);
            if (it != ConnectionPoolMap_.end()) {
                if (Runtime_) {
                    Runtime_->AssignDisposal(it->second);
                    Runtime_->TrySend({});
                }
                ConnectionPoolMap_.erase(it);
            }
        }
    }

    bool TConnectionManager::TryIncConnections() {
        if (++ConnectionsCount_ > Options_.ConnectionsLimit) {
            --ConnectionsCount_;
            return false;
        }
        return true;
    }

    void TConnectionManager::DecConnections() {
        --ConnectionsCount_;
    }

    THolder<TConnectionManager::TRuntime> TConnectionManager::StartRuntime(TContExecutor* executor) {
        THolder<TRuntime> runtime = MakeHolder<TRuntime>(executor, this);
        with_lock (Mutex_) {
            Runtime_ = runtime.Get();
        }
        return runtime;
    }

    void TConnectionManager::ForgetRuntime() {
        with_lock (Mutex_) {
            Runtime_ = nullptr;
        }
    }

    bool TConnectionManager::TrySendToRuntime(std::function<void()> func) {
        with_lock (Mutex_) {
            if (Runtime_) {
                return Runtime_->TrySend(std::move(func));
            }
        }
        return false;
    }

    TConnectionManager::TRuntime::TRuntime(TContExecutor *executor, TIntrusivePtr<TConnectionManager> parent)
        : Executor_(executor)
        , Parent_(std::move(parent))
        , Channel_(Parent_->PreferredChannelCapacity())
    {
        Runtime_ = NSrvKernel::TCoroutine(NSrvKernel::ECoroType::Service, "connection_manager_runtime", Executor_, &TConnectionManager::TRuntime::Run, this);
    }

    TConnectionManager::TRuntime::~TRuntime() {
        Parent_->ForgetRuntime();
        Runtime_.Cancel();
        Runtime_.Join();
    }

    TColdStorage& TConnectionManager::TRuntime::EnsureColdStorage(TIntrusivePtr<TConnectionPool> pool, std::atomic<size_t>& counter) {
        auto it = Storages_.find(pool);
        if (it == Storages_.end()) {
            it = Storages_.emplace(pool, MakeHolder<TColdStorage>(Executor_, pool.Get(), counter)).first;
        }
        return *it->second;
    }

    TColdStorage* TConnectionManager::TRuntime::TryGetColdStorage(TIntrusivePtr<TConnectionPool> pool) {
        auto it = Storages_.find(pool);
        if (it == Storages_.end()) {
            return nullptr;
        }
        return it->second.Get();
    }

    bool TConnectionManager::TRuntime::TrySend(std::function<void()> func) {
        return Channel_.TrySend(std::move(func));
    }

    void TConnectionManager::TRuntime::Run() {
        TCont* cont = Executor_->Running();
        auto waker = NSrvKernel::ThreadLocalEventWaker(cont);
        auto* wakerPtr = waker.Get();

        do {
            std::function<void()> func;
            NSrvKernel::EChannelStatus status = Channel_.Receive(func, TDuration::Seconds(1).ToDeadLine(), cont, wakerPtr);
            if (status == NSrvKernel::EChannelStatus::Success && func) {
                func();
            }
            if (status != NSrvKernel::EChannelStatus::Canceled) {
                ProcessDisposals();
                Parent_->MoveUnusedToColdStorage(Executor_);
            }
        } while (!cont->Cancelled());
        Storages_.clear();
    }

    void TConnectionManager::TRuntime::ProcessDisposals() {
        TList<TIntrusivePtr<TConnectionPool>> toDispose;
        with_lock (Mutex_) {
            ToDispose_.swap(toDispose);
        }
        for (auto& pool : toDispose) {
            Storages_.erase(pool);
        }
    }

    void TConnectionManager::TRuntime::AssignDisposal(TIntrusivePtr<TConnectionPool> pool) {
        with_lock (Mutex_) {
            ToDispose_.push_back(pool);
        }
    }

    void TConnectionManager::MoveUnusedToColdStorage(TContExecutor* executor) {
        // Make a snapshot of pools to not hold manager's mutex for a while
        TVector<TIntrusivePtr<TConnectionPool>> pools;
        with_lock(Mutex_) {
            pools.reserve(ConnectionPoolMap_.size());
            for (const auto& p : ConnectionPoolMap_) {
                pools.push_back(p.second);
            }
        }
        for (auto& pool : pools) {
            pool->MoveUnusedToColdStorage(executor);
        }
    }

    TConnectionManager::TRuntime* TConnectionManager::RuntimeNoLock() const noexcept {
        return Runtime_;
    }
}
