local Class = require "luanode.class"

local PriorityQueue = Class()

function PriorityQueue:__init (size)
	size = size or 0
	assert(type(size) == "number")
	local new = Class.base.rawnew(PriorityQueue)

	-- initialize arrays to hold queue elements
	self.m_size = math.max(size, 1)
	new.slots = {}
	for i = 1, self.m_size do
		table.insert(new.slots, {})
	end

	return new
end

function PriorityQueue:size ()

	if self.total == nil then
		local total = 0
		for i = 1, self.m_size do
			total = total + #self.slots[i]
		end
		self.total = total
	end
	return self.total
end

function PriorityQueue:enqueue (obj, priority)
	priority = priority or 0
	assert(type(priority) == "number")

	-- clear cached total value
	self.total = nil
	if priority ~= 0 then
		local original_priority = priority
		if priority < 0 or priority >= self.m_size then
			priority = self.m_size - 1
			error("invalid priority: " .. original_priority .. " must be between 0 and " .. priority)
		end
	end

	table.insert(self.slots[priority + 1], obj)
end

function PriorityQueue:dequeue (callback)

	-- clear cached total value
	self.total = nil
	for _, slot in ipairs(self.slots) do
		if #slot > 0 then
			return table.remove(slot, 1)
		end
	end
	return nil
end




local Pool = Class()

function Pool:__init (factory)

	local new = Class.base.rawnew(Pool)
	
	new.idleTimeoutMillis = factory.idleTimeoutMillis or 30000
	new.reapInterval = factory.reapIntervalMillis or 1000
	new.refreshIdle = not not factory.refreshIdle
	new.availableObjects = {}
	new.__availableObjectsCount = 0 -- number of available resources in the pool
	new.waitingClients = PriorityQueue(factory.priorityRange or 1)
	new.count = 0 -- number of resources in the pool
	new.removeIdleScheduled = false
	new.removeIdleTimer = nil
	new.draining = false

	new.factory = factory


	-- prepare a logger function
	if factory.log then
		new.log = factory.log
	else
		new.log = function() end	-- NOOP
	end

	factory.validate = factory.validate or function() return true end

	assert(type(factory.max) == "number" or factory.max == nil, "max must be a number")
	assert(type(factory.min) == "number" or factory.min == nil, "min must be a number")

	factory.max = tonumber(factory.max) or 1
	factory.min = tonumber(factory.min) or 0

	factory.max = math.max(factory.max, 1)
	factory.min = math.min(factory.min, factory.max)

	new:ensureMinimum()

	return new
end

---
-- Request the client to be destroyed. The factory's destroy handler will also be called.
-- This should be called within an acquire() block as an alternative to release().
-- @param obj The acquired item to be destroyed.
--
function Pool:destroy (obj)

	if not self.availableObjects[obj] then
		self:release(obj)
	end

	if self.availableObjects[obj] then
		self.count = self.count - 1
		self.availableObjects[obj] = nil
		self.__availableObjectsCount = self.__availableObjectsCount - 1
	end
	self.factory.destroy(obj)
	self:ensureMinimum()
end

---
-- Checks and removes the available (idle) clients that have timed out.
--
function Pool:removeIdle ()
	local to_remove = {}
	local now = process.hrtime()

	self.removeIdleScheduled = false

	-- Go through the available (idle) items, check if they have timed out

	for obj, timestamp in pairs(self.availableObjects) do
		if self.count - self.factory.min == #to_remove then
			break
		end
		local diff = (now - timestamp):ms()
		if diff > self.idleTimeoutMillis then
			-- Client timed out, so destroy it
			self.log("removeIdle() destroying obj - now: %s ns timeout: %s ns, diff: %s (ms)", now, timestamp, diff)
			table.insert(to_remove, obj)
		end
	end

	for _, obj in ipairs(to_remove) do
		self:destroy(obj)
	end

	-- replace the available items with the ones to keep
	if self.__availableObjectsCount > 0 then
		self.log("availableObjects.length=%s", self.__availableObjectsCount)
		self:scheduleRemoveIdle()
	else
		self.log("removeIdle() all objects removed")
	end
end

---
-- Schedule removal of idle items in the pool.
--
function Pool:scheduleRemoveIdle ()
	if not self.removeIdleScheduled then
		self.removeIdleScheduled = true
		self.removeIdleTimer = setTimeout(self.removeIdle, self.reapInterval, self)
	end
end

---
-- Try to get a new client to work, and clean up pool unused (idle) items.
--
-- - If there are available clients waiting, shift the first one out (LIFO),
--   and call its callback.
-- - If there are no waiting clients, try to create one if it won't exceed
--   the maximum number of clients.
-- - If creating a new client would exceed the maximum, add the client to
--   the wait list.
--
function Pool:dispense ()
	local waitingCount = self.waitingClients:size()

	self.log("dispense() clients=%s available=%s", waitingCount, self.__availableObjectsCount)
	if waitingCount > 0 then
		while self.__availableObjectsCount > 0 do
			self.log("dispense() reusing obj")
			local obj = next(self.availableObjects)
			if not self.factory:validate(obj) then
				self:destroy(obj)
			else
				self.availableObjects[obj] = nil
				self.__availableObjectsCount = self.__availableObjectsCount - 1
				local clientCb = self.waitingClients:dequeue()
				process.nextTick(function()
					clientCb(nil, obj)
				end)
				return
			end
		end

		if self.count < self.factory.max then
			self:createResource()
		end
	end
end

---
--
function Pool:createResource ()
	self.count = self.count + 1
	self.log("createResource() - creating obj - count=%s min=%s max=%s", self.count, self.factory.min, self.factory.max)

	self.factory.create(function(err, obj)
		local clientCb = self.waitingClients:dequeue()

		if err then
			self.count = self.count - 1
			if clientCb then
				process.nextTick(function()
					clientCb(err, obj)
				end)
			end
			process.nextTick(function()
				self:dispense()
			end)
		else
			if clientCb then
				process.nextTick(function()
					clientCb(err, obj)
				end)
			else
				self:release(obj)
			end
		end
	end)
end

---
--
function Pool:ensureMinimum ()
	if not self.draining and self.count < self.factory.min then
		local diff = self.factory.min - self.count
		for i = 1, diff do
			self:createResource()
		end
	end
end

---
-- Request a new client. The callback will be called,
-- when a new client will be availabe, passing the client to it.
--
-- @param {Function} callback
--   Callback function to be called after the acquire is successful.
--   The function will receive the acquired item as the first parameter.
--
-- @param {Number} priority
--   Optional.  Integer between 0 and (priorityRange - 1).  Specifies the priority
--   of the caller if there are no available resources.  Lower numbers mean higher
--   priority.
--
-- @returns {Object} `true` if the pool is not fully utilized, `false` otherwise.
--
function Pool:acquire (callback, priority)
	if self.draining then
		error("pool is draining and cannot accept work")
	end
	self.waitingClients:enqueue(callback, priority)
	self:dispense()
	return self.count < self.factory.max
end

---
-- Return the client to the pool, in case it is no longer required.
-- @param {Object} obj
--   The acquired object to be put back to the pool.
--
function Pool:release (obj)
	-- check to see if this object has already been released (i.e., is back in the pool of availableObjects)
	if self.availableObjects[obj] then
		error("release called twice for the same resource: " .. tostring(obj))
	end

	--self.log("return to pool")
	local timestamp = process.hrtime()
	self.availableObjects[obj] = timestamp
	self.__availableObjectsCount = self.__availableObjectsCount + 1
	self.log("timeout: %s", timestamp)
	self:dispense()
	self:scheduleRemoveIdle()
end

---
-- Disallow any new requests and let the request backlog dissipate.
--
-- @param {Function} callback
--   Optional. Callback invoked when all work is done and all clients have been
--   released.
--
function Pool:drain (callback)
	self.log("draining")

	-- disable the ability to put more work on the queue.
	self.draining = true

	local function check ()
		if self.waitingClients:size() > 0 then
			-- wait until all client requests have been satisfied.
			setTimeout(check, 100)
		elseif self.__availableObjectsCount ~= self.count then
			-- wait until all objects have been released.
			setTimeout(check, 100)
		else
			if callback then
				callback()
			end
		end
	end
	check()
end

---
-- Forcibly destroys all clients regardless of timeout.  Intended to be
-- invoked as part of a drain.  Does not prevent the creation of new
-- clients as a result of subsequent calls to acquire.
--
-- Note that if factory.min > 0, the pool will destroy all idle resources
-- in the pool, but replace them with newly created resources up to the
-- specified factory.min value.  If this is not desired, set factory.min
-- to zero before calling destroyAllNow()
--
-- @param {Function} callback
--   Optional. Callback invoked after all existing clients are destroyed.
--
function Pool:destroyAllNow (callback)
	self.log("force destroying all objects")
	local willDie = self.availableObjects

	self.availableObjects = {}
	self.__availableObjectsCount = 0
	
	for obj in pairs(willDie) do
		self:destroy(obj)
	end

	self.removeIdleScheduled = false
	if self.removeIdleTimer then
		clearTimeout(self.removeIdleTimer)
		self.removeIdleTimer = nil
	end
	if callback then
		callback()
	end
end

---
-- Decorates a function to use a acquired client from the object pool when called.
--
-- @param {Function} decorated
--   The decorated function, accepting a client as the first argument and
--   (optionally) a callback as the final argument.
--
-- @param {Number} priority
--   Optional.  Integer between 0 and (priorityRange - 1).  Specifies the priority
--   of the caller if there are no available resources.  Lower numbers mean higher
--   priority.
--
function Pool:pooled (decorated, priority)
	
	return function (...)
		local callerArgs, original_args = {...}, {...}
		local args_len = select("#", ...)
		local original_args_len = args_len

		local callerCallback
		-- if the caller supplied a callback, put it aside
		if type(callerArgs[args_len]) == "function" then
			callerCallback = callerArgs[args_len]
			callerArgs[args_len] = nil
			args_len = args_len - 1
		end
		
		self:acquire(function(err, client)
			if err then
				if callerCallback then callerCallback(err) end
				return
			end
			
			-- add the client af the front and our cleanup function at the back
			table.insert(callerArgs, 1, client)
			table.insert(callerArgs, function()
				self:release(client)
				if callerCallback then callerCallback(unpack(original_args, 1, original_args_len)) end
			end)
			decorated(unpack(callerArgs, 1, args_len + 2))
		end, priority)
	end
end

function Pool:getPoolSize ()
	return self.count
end

function Pool:getName ()
	return self.factory.name
end

function Pool:availableObjectsCount ()
	return self.__availableObjectsCount
end

function Pool:waitingClientsCount ()
	return self.waitingClients:size()
end

---
-- just chains drain and stop
function Pool:forceStop (callback)
	self.min = 0
	self:destroyAllNow(callback)
end

return Pool
