#ifndef PKMNLIB_CONTEXTPOOL_HPP
#define PKMNLIB_CONTEXTPOOL_HPP

#include <angelscript.h>
#include <mutex>
#include <thread>

class ContextPool {
    struct CtxThreadedPoolStorage {
        std::vector<asIScriptContext*> Pool;
        std::mutex Lock;

        ~CtxThreadedPoolStorage() {
            for (auto* ctx : Pool) {
                ctx->Release();
            }
        }
    };

    ArbUt::Dictionary<std::thread::id, CtxThreadedPoolStorage*> _pool;
    asIScriptEngine* _engine;

    CtxThreadedPoolStorage* GetThreadPool() {
        auto id = std::this_thread::get_id();
        if (!_pool.Has(id)) {
            auto t = CtxThreadedPoolStorage();
            _pool.Insert(id, new CtxThreadedPoolStorage());
        }
        return _pool[id];
    }

public:
    ContextPool(asIScriptEngine* engine) : _engine(engine) {}

    ~ContextPool() {
        for (const auto& kv : _pool) {
            delete kv.second;
        }
    }

    asIScriptContext* RequestContext() {
        // Get a context from the pool, or create a new
        asIScriptContext* ctx = nullptr;
        auto* pool = GetThreadPool();
        std::lock_guard<std::mutex> guard(pool->Lock);
        if (!pool->Pool.empty()) {
            ctx = *pool->Pool.rbegin();
            pool->Pool.pop_back();
        } else {
            ctx = _engine->CreateContext();
        }
        return ctx;
    }

    void ReturnContextToPool(asIScriptContext* ctx) {
        ctx->Unprepare();
        auto* pool = GetThreadPool();
        std::lock_guard<std::mutex> guard(pool->Lock);
        pool->Pool.push_back(ctx);
    }
};

#endif // PKMNLIB_CONTEXTPOOL_HPP