--[[----------------------------------------------------------------------------

WFRxRateLimiter.lua

--------------------------------------------------------------------------------

ADOBE CONFIDENTIAL
------------------
Copyright 2016 Adobe Systems Incorporated
All Rights Reserved.

NOTICE: All information contained herein is, and remain the property of
Adobe Systems Incorporated and its suppliers, if any. The intellectual and
technical concepts contained herein are proprietary to Adobe Systems
Incorporated and its suppliers and may be covered by U.S. and Foreign
Patents, patents in process, and are protected by trade secret or copyright
law.  Dissemination of this information or reproduction of this material is
strictly forbidden unless prior written permission is obtained from
Adobe Systems Incorporated.

------------------------------------------------------------------------------]]

	-- RxLua
local Rx = require 'Rx'
local Observable = Rx.Observable

	-- WF.core.diagnostic_log
local WFDiagnosticLog = require 'WFDiagnosticLog'

	-- WF.core.threading
local WFCurrentThread = require 'WFCurrentThread'

	-- WF.core.utils
local WFLogger = require 'WFLogger'
local WFPerformanceReport = require 'WFPerformanceReport'
local KSDate = require 'KSDate'

--------------------------------------------------------------------------------

local pcall = pcall
local next = next

local logger = WFLogger 'WFRxRateLimiter'

--------------------------------------------------------------------------------

local function makeTaskPool()

	return {
		n = 0,
		cost = 0,
		highestPriorityRunning = 0,
		tasks = {} -- Sorted task list for quick rescheduling.
	}

end

--------------------------------------------------------------------------------

local function sortTasks( a, b )

	-- First sort by priority
	-- Then sort by cost
	-- Then sort by time

	if not b then return true end

	if not a then return false end
	
	if a.priority ~= b.priority then
		return a.priority < b.priority
	end
	
	if a.cost ~= b.cost then
		return a.cost > b.cost
	end
	
	return a.timeStamp > b.timeStamp
	
end

--------------------------------------------------------------------------------

local function recordStaminaSnapshot( pool )

	-- STAMINA: Record current snapshot of how many requests are in flight, and their cost.
	
	WFPerformanceReport.recordSingleValue( 
			"WFRxRateLimiter.inFlightRequests",
			{ 
				requestCount = pool.n,
				requestCost = pool.cost
			} )

end

--------------------------------------------------------------------------------

-- Add a task to a pool. Returns true if this was a new task.

local function taskPool_addTask( pool, observable, cost, maybePriority )

	-- Create new task pool entry
	
	local task = { 
		observable = observable, 
		priority = maybePriority or 0,
		cost = cost or 1,
		timeStamp = KSDate.timeNowForAnimation()
	}
	
	table.insert( pool.tasks, task )
	table.sort( pool.tasks, sortTasks )

	-- Update counts
	
	pool.n = pool.n + 1
	pool.cost = pool.cost + task.cost
	pool.highestPriorityRunning = pool.tasks[ pool.n ].priority
	
	recordStaminaSnapshot( pool )
	
	logger:infof( "Current pool cost: %d", pool.cost )
	
	return true

end

--------------------------------------------------------------------------------

-- Remove a task from a pool.

local function taskPool_removeTask( pool, observable )

	logger:debug( "Attempting to clear taskPool entry" )

	local taskToRemove
	
	-- Remove from sorted task list
	
	for k, task in ipairs( pool.tasks ) do
		
		if task.observable == observable then
			taskToRemove = task
			table.remove( pool.tasks, k )
			break
		end
		
	end
	
	-- Adjust pool counts
	
	if taskToRemove ~= nil then
	
		logger:debug( "Clearing taskPool entry" )

		pool.n = pool.n - 1
		pool.cost = pool.cost - taskToRemove.cost
		
		logger:infof( "Current pool cost: %d", pool.cost )
		
	elseif pool.n > 0 then
	
		-- Check to see whether there are any items outstanding or whether
		-- it has all been GC'd
	
		if next( pool.tasks ) == nil then
			pool.n = 0
		end

	end

	pool.highestPriorityRunning = pool.n > 0 and pool.tasks[ pool.n ].priority or 0

	recordStaminaSnapshot( pool )

end

--------------------------------------------------------------------------------

local function createOnErrorOrCompletedOrDisposed( self, observable, cost, maybePriority, maybeDesc )

	local hasRun = false

	return function()

		if not hasRun then

			logger:infof( "Observable %p Completed of Cost %d, Priority %d, Description %s", 
					observable, cost, maybePriority or 0, maybeDesc )
					
			self._scheduledItems[ observable ] = nil
			self:_finishedTask( observable )
			hasRun = true

		end

	end
	
end

--------------------------------------------------------------------------------

local function createTriggerTask( 
		self, observable, disposable, observer, cost, maybePriority, maybeDesc )

	return function( cancelationReason )

		logger:infof( "Starting task of Cost %d, Priority %d, Description %s", cost, maybePriority or 0, maybeDesc )

		if cancelationReason then

			logger:infof( "Canceled task of Cost %d, Priority %d, Description %s error %s", cost, maybePriority or 0, maybeDesc, cancelationReason )

			observer:onError( cancelationReason )

			disposable:dispose()

			return

		end

		if not disposable:isDisposed() then

			self:_startedTask( observable, cost, maybePriority )
			
			local _onErrorOrCompletedOrDisposed = 
					createOnErrorOrCompletedOrDisposed( self, observable, cost, maybePriority, maybeDesc )

			local perfTable = {
				priority = maybePriority,
				cost = cost,
			}

			local workSubscription = WFPerformanceReport.wrapObservable( 'WFRxRateLimiter: getOrPut', 
					observable, perfTable )
					:finally( _onErrorOrCompletedOrDisposed )
					:subscribe( observer )

			local rateLimiterDisposable = Rx.Disposable.create( function()

					--logger:debugf( "Disposing task of Cost %d, Priority %d", cost, maybePriority or 0 )
					_onErrorOrCompletedOrDisposed()
					workSubscription:dispose()

				end )

			disposable:setDisposable( rateLimiterDisposable )

		else

			-- Even though this task was never started, this signal allows
			-- us to immediately process the next queue item.
			
			self:_finishedTask( observable )

		end

	end
	
end

--------------------------------------------------------------------------------

-- public methods for all rate limiters

local publicMethods = {}

--------------------------------------------------------------------------------

--[[!
	\brief Add an observable to the queue of observable to be processed.

	\param observable (Rx.Observable) The observable ( request ) we want to schedule.
	\param maybePriority (optional, number) Priority of this request.
	\param maybeCost (optional, number) How expensive is the operation. In other words, how many
		slots should it take up.
	\param maybeDesc (optional, string) A optional description of the observable. Used for debugging.
		
	\return <code>Observable</code> An observable that resolves with the values of the first
		observable parameter.
]]--
	
function publicMethods:enqueueObservable( observable, maybePriority, maybeCost, maybeDesc )

	return Observable.create( function( observer )

		local cost = maybeCost or 1

		logger:infof( "Enqueueing observable. Cost %d, Priority %d, Description %s", cost, maybePriority, maybeDesc )

		local disposable = Rx.SerialDisposable()
		
		local triggerTask = createTriggerTask( 
				self, observable, disposable, observer, cost, maybePriority, maybeDesc )

		self:_enqueueTask( triggerTask, cost, maybePriority )

		self._scheduledItems[ observable ] = { disposable, observer }

		return disposable

	end )

end

--------------------------------------------------------------------------------

local function createQueueEntry( action, cost )

	return {
		action = action,
		cost = cost
	}

end

--============================================================================--

-- @mark WFRxRateLimiter namespace

--------------------------------------------------------------------------------

local WFRxRateLimiter = {}

--------------------------------------------------------------------------------

--[[!
	\brief Create a rate limiter which can be used to limit the number of simultaneous
	tasks.

	The WFRxRateLimiter will by-pass cost limits of in-flight requests if higher
	priority requests are issued and no execution slots are available.


	\param makeQueue (WFQueuing object) The data structure to use to store the
		tasks. WFQueueing.makePriorityQueue, WFQueueing.makeQueue, WFQueueing.makeStack 
		are all valid options.
	\param targetSimultaneousCost (number) The maximum simultaneous cost of tasks 
		when equal priority tasks are provided. A single higher priority task than
		the priority of tasks currently running will be allowed to run even though
		it's cost would exceed the targetSimultaneousCost.
	\param allowSingleHighestPriorityInQueue (number) If specified it should be the
		highest priority allowed and will cause the rate limiter to check if another
		task of that priority is queued up, if it is the existing task will be dequeued
		and sent a canceled error and the new task will be queued.  This does not stop
		an already running task. The intention of this is to prevent highest priority
		tasks from clogging up the queue, and the use of this highest priority should
		only be for cases where you issue it when you want a specific task to get
		immediately processed.

	\return <code>WFRxRateLimiter</code> A WFRxRateLimiter object. Note: Please see 
		WFRxRateLimiterTest for how to use this returned object.
		
]]--

function WFRxRateLimiter.makeRateLimiter( makeQueue, targetSimultaneousCost, allowSingleHighestPriorityInQueue )

	local enqueue, dequeue, _, peekHead = makeQueue()
	
	local outstandingTasks = makeTaskPool()

	local result = {}
	
	for k, v in pairs( publicMethods ) do
		result[ k ] = v
	end
	
	result._scheduledItems = setmetatable( {}, { __mode = 'k' } )
		-- key: ( Rx.Observable)
		-- value: (table) { Rx.SerialDisposable(), Rx.observer )

	result._enqueue = enqueue

	----------------------------------------------------------------------------

	local startingActions = false
	
	local function startActions( self )
	
		logger:debug( "Checking for more tasks to process.." )
	
		if startingActions then
			logger:debug( "Check already in progress.. aborting." )
			return
		end
		
		startingActions = true


		local maybePriority = peekHead()
		local priority = type( maybePriority ) == 'number' and maybePriority or 0
		local isHigherPriorityThanRunningTasks = priority > outstandingTasks.highestPriorityRunning

		while isHigherPriorityThanRunningTasks or outstandingTasks.cost < targetSimultaneousCost do

			local maybePriorityOrAction, maybeAction = peekHead()
			
			if maybePriorityOrAction == nil then
				break
			end

			local action = maybePriorityOrAction
			local cost = 1
			
			if type( action ) ~= 'function' then
				action = maybeAction.action
				cost = maybeAction.cost
			end

			if isHigherPriorityThanRunningTasks then
				isHigherPriorityThanRunningTasks = false
			elseif cost + outstandingTasks.cost > targetSimultaneousCost then
				logger:debugf( "cost is too high. current: %d, task: %d, aborting.", outstandingTasks.cost, cost )
				break
			end

			-- Remove the peek'd entry, since we're about to process it.

			dequeue()
				
			
			if not action or type( action ) ~= 'function' then
				
				-- If the task has no action, this is a no-op.
				
				logger:debugf( "Task has no action." )
				break
				
			end
			
			pcall( action )
			
		end

		startingActions = false
		
	end
	
	----------------------------------------------------------------------------
	
	result._enqueueTask = function( self, taskFn, taskCost, maybeTaskPriority )

		if not taskFn then
			return
		end

		local item = createQueueEntry( taskFn, taskCost )

		if maybeTaskPriority then


			if maybeTaskPriority == allowSingleHighestPriorityInQueue then

				local maybePriority, maybeAction = peekHead()
				local highestPriorityInQueue = type( maybePriority ) == 'number' and maybePriority or 0

				if highestPriorityInQueue == allowSingleHighestPriorityInQueue then
					dequeue()

					pcall( maybeAction.action, 'canceledByNewRequest' )
				end

			end

			enqueue( maybeTaskPriority, item )

		else
			enqueue( item )
		end

		startActions( self )
	
	end
	
	----------------------------------------------------------------------------
	
	result._startedTask = function( self, observable, cost, maybePriority )
	
		return taskPool_addTask( outstandingTasks, observable, cost, maybePriority )

	end

	----------------------------------------------------------------------------

	result._finishedTask = function( self, observable )
	
		-- Remove the task from the pool
	
		taskPool_removeTask( outstandingTasks, observable )
		
		-- Try to start actions
		
		startActions( self )

	end
	
	----------------------------------------------------------------------------

	WFDiagnosticLog.addLogSectionFunc( function()

		local threadName = WFCurrentThread.DEBUG_getThreadName() or '( unknown )'

		local activeTasks = {}

		local now = KSDate.timeNowForAnimation()

		for _, task in ipairs( outstandingTasks.tasks ) do

			local logEntry = {
				cost = task.cost,
				priority = task.priority,
				timeRunning = now - task.timeStamp,
			}

			table.insert( activeTasks, logEntry )

		end

		local scheduledItemCount = 0
		for _, _ in pairs( result._scheduledItems ) do
			scheduledItemCount = scheduledItemCount + 1
		end

		return {

			name = "WFRxRateLimiter: " .. threadName,
			id = 'WFRxRateLimiter',
			weight = 402,
			outputForAutomation = true,
			parentSection = WFDiagnosticLog.Sections.HTTP,
			
			body = {
				targetSimultaneousCost = targetSimultaneousCost,
				outstandingTasks = activeTasks,
				outstandingTasksCount = outstandingTasks.n,
				outstandingTasksCost = outstandingTasks.cost,
				highestPriorityRunning = outstandingTasks.highestPriorityRunning,
				scheduledItemCount = scheduledItemCount,
			}
		}

	end )

	return result

end

--------------------------------------------------------------------------------

return WFRxRateLimiter
