using System.Collections.Generic;

namespace Cryville.Common.Buffers {
	public abstract class CategorizedPool<TCategory, TObject> where TObject : class {
		protected readonly Dictionary<TCategory, ObjectPool<TObject>> Buckets
			= new Dictionary<TCategory, ObjectPool<TObject>>();
		public int RentedCount { get; private set; }
		public TObject Rent(TCategory name) {
			var obj = Buckets[name].Rent();
			RentedCount++;
			return obj;
		}
		public void Return(TCategory name, TObject obj) {
			Buckets[name].Return(obj);
			--RentedCount;
		}
	}
	public class CategorizedPoolAccessor<TCategory, TObject> where TObject : class {
		readonly CategorizedPool<TCategory, TObject> _pool;
		static readonly SimpleObjectPool<Dictionary<TObject, TCategory>> _dictPool
			= new SimpleObjectPool<Dictionary<TObject, TCategory>>(1024);
		Dictionary<TObject, TCategory> _rented;
		public CategorizedPoolAccessor(CategorizedPool<TCategory, TObject> pool) {
			_pool = pool;
		}
		public TObject Rent(TCategory name) {
			var obj = _pool.Rent(name);
			if (_rented == null) _rented = _dictPool.Rent();
			_rented.Add(obj, name);
			return obj;
		}
		public void Return(TObject obj) {
			_pool.Return(_rented[obj], obj);
			_rented.Remove(obj);
		}
		public void ReturnAll() {
			if (_rented == null) return;
			foreach (var obj in _rented) {
				_pool.Return(obj.Value, obj.Key);
			}
			_rented.Clear();
			_dictPool.Return(_rented);
			_rented = null;
		}
	}
}