local Rx = require 'Rx'

local AnonymousObservable = Rx.AnonymousObservable
local CompositeDisposable = Rx.CompositeDisposable
local Observable = Rx.Observable
local SchedulerArgument = Rx.internals.SchedulerArgument
local SerialDisposable = Rx.SerialDisposable
local SingleAssignmentDisposable = Rx.SingleAssignmentDisposable

local unpack = table.unpack or unpack

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

local function delayRelative(source, delayTime, scheduler)

    scheduler = SchedulerArgument.validateWithoutDefault('Rx.Observable:delay', scheduler)

    return AnonymousObservable(function(observer)

        local active = false
        local cancelable = SerialDisposable()
        local exception = nil
        local q = {}
        local running = false

        local materializedSource = source:materialize():timestamp(scheduler)

        local subscription = materializedSource:subscribe(function(timestamp, notification)

            if notification:kind() == 'E' then

                q = {
                    timestamp = timestamp,
                    value = notification
                }

                exception = notification._errors

                if not running then
                    observer:onError(unpack(exception, 1, exception.n))
                end

            else

                q[#q + 1] = {
                    timestamp = timestamp + delayTime,
                    value = notification
                }

                if not active then

                    active = true

                    local d = SingleAssignmentDisposable()

                    cancelable:setDisposable(d)

                    d:setDisposable(scheduler:scheduleRecursiveFuture(delayTime,
                        function(_, recurseFn)

                            if exception == nil then

                                running = true

                                local result
                                repeat

                                    result = nil

                                    if #q > 0 and q[1].timestamp - scheduler:now() <= 0 then

                                        result = table.remove(q, 1).value

                                        if result ~= nil then
                                            result:accept(observer)
                                        end

                                    end

                                until result == nil

                                running = false

                                local shouldRecurse = #q > 0
                                local recurseDelay

                                if shouldRecurse then
                                    recurseDelay = math.max(0, q[1].timestamp - scheduler:now())
                                else
                                    recurseDelay = 0
                                    active = false
                                end

                                if exception ~= nil then

                                    observer:onError(unpack(exception, 1, exception.n))

                                elseif shouldRecurse then

                                    recurseFn(nil, recurseDelay)

                                end

                            end

                    end))

                end

            end

        end)

        return CompositeDisposable(subscription, cancelable)

    end, source)

end

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

local function delayWithSelector(source, delayDurationSelector, subscriptionDelay)

    return AnonymousObservable(function(observer)

        local delays = CompositeDisposable()
        local subscription = SerialDisposable()
        local atEnd = false

        local function done()

            if atEnd and delays:length() == 0 then
                observer:onCompleted()
            end

        end

        local function start()

            subscription:setDisposable(source:subscribe(
                function(...)

                    local value = { n = select('#', ...), ... }

                    local success, delay = pcall(delayDurationSelector, ...)

                    if not success then
                        local errorMessage = delay
                        return observer:onError(errorMessage)
                    end

                    local d = SingleAssignmentDisposable()
                    delays:add(d)

                    d:setDisposable(delay:subscribe(
                        function()
                            observer:onNext(unpack(value, 1, value.n))
                            delays:remove(d)
                            done()
                        end,
                        function(...)
                            observer:onError(...)
                        end,
                        function()
                            observer:onNext(unpack(value, 1, value.n))
                            delays:remove(d)
                            done()
                        end))

                end,
                function(...)

                    observer:onError(...)

                end,
                function()

                    atEnd = true
                    subscription:dispose()
                    done()

                end))

        end

        if subscriptionDelay == nil then

            start()

        else

            subscription:setDisposable(subscriptionDelay:subscribe(
                start,
                function(...)
                    observer:onError(...)
                end,
                start))

        end

        return CompositeDisposable(subscription, delays)

    end, source)

end

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

return function(self, delayOrSelector, schedulerOrSelector)

    local delayOrSelectorType = type(delayOrSelector)

    if delayOrSelectorType == 'number' then

        return delayRelative(self, delayOrSelector, schedulerOrSelector)

    elseif delayOrSelectorType == 'function' then

        return delayWithSelector(self, delayOrSelector)

    elseif Observable:isClassOf(delayOrSelector)
        and type(schedulerOrSelector) == 'function' then

        return delayWithSelector(self, schedulerOrSelector, delayOrSelector)

    else

        error("Rx.Observable:delay: invalid arguments")

    end

end
