#include <util/system/yassert.h>

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

namespace NConnectionManager {
    namespace {
        bool IsTimeout(const NSrvKernel::TError& error) {
            const auto* se = error.GetAs<TSystemError>();
            return se && se->Status() == ETIMEDOUT;
        }
    }
    TSocketContainer::TSocketContainer(THolder<NModProxy::TKeepAliveData> connection)
        : SockAddr_(connection->ConcreteAddress_)
        , SocketHolder_(std::move(connection->BackendIo_->SockHolder()))
        , SockAddrInfo_(std::move(connection->Addresses_))
        , KeepAliveReleaseDeadline_(connection->KeepAliveReleaseDeadline_)
        , SslIo_(std::move(connection->SslIo_))
        , SniHost_(std::move(connection->SniHost_)) {
        Y_VERIFY(connection->CanStore());
        if (SslIo_) {
            SslIo_->SetIo(nullptr, nullptr);
        }
    }

    THolder<NModProxy::TKeepAliveData> TSocketContainer::BuildKeepaliveData(TContExecutor* executor, NModProxy::EPollMode pollMode) {
        if (!SocketHolder_) {
            return {};
        }
        auto backendIo = MakeHolder<NModProxy::TBackendIo>(std::move(SocketHolder_), executor, pollMode);
        auto keepAlive = MakeHolder<NModProxy::TKeepAliveData>(
                std::move(backendIo),
                std::move(SockAddrInfo_),
                SockAddr_
            );
        keepAlive->SetKeepAliveReleaseDeadline(KeepAliveReleaseDeadline_);
        if (SslIo_) {
            keepAlive->SslIo_ = std::move(SslIo_);
            keepAlive->SslIo_->SetIo(&keepAlive->BackendIo_->Input(), &keepAlive->BackendIo_->Output());
        }
        keepAlive->SniHost_ = std::move(SniHost_);
        return keepAlive;
    }

    TConnectionPool::TConnectionPool(TConnectionManager& parent, size_t tlsLimit, size_t connectionLimit, NSrvKernel::EPollMode pollMode, TDuration keepaliveTimeout, TDuration connectTimeout)
        : Parent_(parent)
        , Tls_(tlsLimit)
        , ConnectionLimit_(connectionLimit)
        , PollMode_(pollMode)
        , KeepaliveTimeout_(keepaliveTimeout)
        , ConnectTimeout_(connectTimeout)
    {}

    TConnectionPool::~TConnectionPool() {
        for (auto& packed : Storage_) {
            OnExpire();
            packed = {};
        };
        Storage_.clear();
    }

    void TConnectionPool::SetOnExpire(std::function<void()> onExpire) {
        OnExpire_ = onExpire;
    }

    void TConnectionPool::OnExpire() {
        DecConnections();
        if (OnExpire_) {
            OnExpire_();
        }
    }

    TWorkerConnectionPool* TConnectionPool::GetPoolForWorker(NSrvKernel::IWorkerCtl* process) {
        Y_VERIFY(process->WorkerId() < Tls_.size());
        THolder<TWorkerConnectionPool>& pool = Tls_[process->WorkerId()];
        if (!pool) {
            pool.Reset(new TWorkerConnectionPool(&process->Executor(), *this));
        }
        return pool.Get();
    }

    void TConnectionPool::ReleasePoolForWorker(NSrvKernel::IWorkerCtl* process) {
        Y_VERIFY(process->WorkerId() < Tls_.size());
        Tls_[process->WorkerId()].Destroy();
    }

    bool TConnectionPool::AddConnection(THolder<NModProxy::TKeepAliveData>& connection) {
        if (!TryIncConnections()) {
            return false;
        }
        TSocketContainer packed(std::move(connection));
        with_lock(Mutex_) { //TODO: try_lock with timeout
            packed.SetDeadline(KeepaliveTimeout_.ToDeadLine());
            Storage_.push_front(std::move(packed));
        }
        return true;
    }

    NSrvKernel::EPollMode TConnectionPool::PollMode() const noexcept {
        return PollMode_;
    }

    TSocketContainer TConnectionPool::TryPullHotConnection() {
        if (!Count_) {
            return {};
        }

        with_lock (Mutex_) {
            if (!Storage_.empty()) {
                TSocketContainer packed = std::move(Storage_.front());
                DecConnections();
                Storage_.pop_front();
                return packed;
            }
        }

        return {};
    }

    TSocketContainer TConnectionPool::TryPullConnection() {
        TSocketContainer packed = TryPullHotConnection();
        if (!packed.IsEmpty()) {
            return packed;
        }
        if (++RequestedFromCold_ > StoredInCold_) {
            --RequestedFromCold_;
            return {};
        }
        //TODO: use shared channel for all pool's requests?
        auto channel = MakeAtomicShared<NSrvKernel::TW2UChannel<char>>(1u);
        TIntrusivePtr<TConnectionPool> self = this;
        if (Parent_.TrySendToRuntime([self, channel](){
            TColdStorage* storage = self->Parent_.RuntimeNoLock()->TryGetColdStorage(self);
            if (storage) {
                THolder<NModProxy::TKeepAliveData> connection = storage->TryPull();
                --self->RequestedFromCold_;
                if (connection) {
                    TSocketContainer packed{std::move(connection)};
                    with_lock (self->Mutex_) {
                        packed.SetDeadline(self->KeepaliveTimeout_.ToDeadLine());
                        self->Storage_.push_front(std::move(packed));
                    }
                }
            } else {
                --self->RequestedFromCold_;
            }
            Y_UNUSED(channel->TrySend('a'));
        })) {
            char c;
            NSrvKernel::EChannelStatus status = channel->Receive(c, ConnectTimeout_.ToDeadLine());
            if (status == NSrvKernel::EChannelStatus::Canceled) {
                return {};
            }
            return TryPullHotConnection();
        }
        --RequestedFromCold_;
        return {};
    }

    bool TConnectionPool::TryIncConnections() {
        if (++Count_ > ConnectionLimit_) {
            --Count_;
            return false;
        }
        if (!Parent_.TryIncConnections()) {
            --Count_;
            return false;
        }
        return true;
    }

    void TConnectionPool::DecConnections() {
        --Count_;
        Parent_.DecConnections();
    }

    void TConnectionPool::MoveUnusedToColdStorage(TContExecutor* executor) {
        for (;;) {
            TSocketContainer packed;
            with_lock(Mutex_) {
                if (Storage_.empty()) {
                    break;
                }
                TInstant deadline = Storage_.back().Deadline();
                if (deadline <= Now()) {
                    OnExpire();
                } else {
                    TDuration notUsedPeriod = KeepaliveTimeout_.ToDeadLine() - deadline;
                    constexpr TDuration hotStorageDuration = TDuration::Seconds(1);
                    if (notUsedPeriod < hotStorageDuration) {
                        break;
                    }
                    packed = std::move(Storage_.back());
                }
                Storage_.pop_back();
            }
            if (!packed.IsEmpty()) {
                THolder<NModProxy::TKeepAliveData> connection = packed.BuildKeepaliveData(executor,PollMode());
                Parent_.RuntimeNoLock()->EnsureColdStorage(this, StoredInCold_).Add(std::move(connection));
            }
        }
    }

    TColdStorage::TColdStorage(TContExecutor* executor, TConnectionPool* owner, std::atomic<size_t>& stored)
        : Executor_(executor)
        , Owner_(owner)
        , Stored_(stored)
    {
        Background_ = NSrvKernel::TCoroutine{NSrvKernel::ECoroType::Service, "cold_storage_cont", Executor_, [this]{
            BackgroundTask();
        }};
    }

    TColdStorage::~TColdStorage() {
        Background_.Cancel();
        Background_.Join();
        Signalled_.clear();
        for (auto& conn : Connections_) {
            --Stored_;
            conn.Destroy();
            Owner_->OnExpire();
        }
        Connections_.clear();
    }

    void TColdStorage::Add(THolder<NModProxy::TKeepAliveData> connection) {
        bool updateDeadline = Connections_.empty();
        Connections_.push_back(std::move(connection));
        auto it = --Connections_.end();
        (*it)->SocketIo().In().SetOnReadyCallback([this, it](int) {
            Signalled_.insert({it->Get(), it});
            BackgroundCondVar_.notify();
        });
        if (updateDeadline) {
            BackgroundCondVar_.notify();
        }
        ++Stored_;
    }

    THolder<NModProxy::TKeepAliveData> TColdStorage::TryPull() {
        if (Connections_.empty() || !Connections_.back()) {
            // In case !Connections_.back(), we are unlucky to face with a connection
            // that is currently processed in ProcessSignalled
            // Let's hope this is rare, and return no keepalive connection here
            return {};
        }
        --Stored_;
        THolder<NModProxy::TKeepAliveData> connection = std::move(Connections_.back());
        connection->SocketIo().In().SetOnReadyCallback({});
        Signalled_.erase(connection.Get());
        Connections_.pop_back();
        return connection;
    }

    void TColdStorage::BackgroundTask() {
        const auto* const cont = Executor_->Running();
        while (!cont->Cancelled()) {
            while (!Connections_.empty()) {
                if (Connections_.front()->KeepAliveReleaseDeadline() <= Now()) {
                    --Stored_;
                    Signalled_.erase(Connections_.front().Get());
                    Connections_.front()->SocketIo().In().SetOnReadyCallback({});
                    Owner_->OnExpire();
                    Connections_.pop_front();
                } else {
                    break;
                }
            }
            ProcessSignalled();
            if (Connections_.empty()) {
                Y_UNUSED(BackgroundCondVar_.wait(Executor_));
            } else {
                Y_UNUSED(BackgroundCondVar_.wait_until(Executor_, Connections_.front()->KeepAliveReleaseDeadline()));
            }
        }
    }

    void TColdStorage::ProcessSignalled() {
        while (!Signalled_.empty()) {
            auto it = Signalled_.begin();
            NSrvKernel::TChunkList lst;
            // Steal connection from the list, so that no one could take it, while we're reading
            THolder<NModProxy::TKeepAliveData> connection = std::move(*it->second);
            NSrvKernel::TError error = connection->Input().Recv(lst, TInstant::Now());
            if (!lst.Empty() || (error && !IsTimeout(error)) || connection->IsSslShutdown()) {
                // Another end is dead or have unexpected data for us
                --Stored_;
                connection->SocketIo().In().SetOnReadyCallback({});
                Owner_->OnExpire();
                Connections_.erase(it->second);
            } else {
                // Put the connection back, it's still valid and can be reused
                *it->second = std::move(connection);
            }
            Signalled_.erase(it);
        }
    }

    TWorkerConnectionPool::TWorkerConnectionPool(TContExecutor* executor, TConnectionPool& parent)
        : Executor_(executor)
        , Parent_(parent)
    {}

    bool TWorkerConnectionPool::AddConnection(THolder<NModProxy::TKeepAliveData>& connection) {
        return Parent_.AddConnection(connection);
    }

    THolder<NModProxy::TKeepAliveData> TWorkerConnectionPool::TryPullConnection() {
        TSocketContainer packed = Parent_.TryPullConnection();
        return packed.BuildKeepaliveData(Executor_, Parent_.PollMode());
    }
}
