local Rx = require 'Rx'

local SingleAssignmentDisposable = Rx.SingleAssignmentDisposable
local AnonymousObservable = Rx.AnonymousObservable
local CompositeDisposable = Rx.CompositeDisposable
local validateNonNegativeInteger = Rx.internals.NumericArgument.validateNonNegativeInteger

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

return function(self, ...)

    local maxConcurrent = ...

    if maxConcurrent == nil then
        return self:flatMergeAll()
    end

    maxConcurrent = validateNonNegativeInteger('maxConcurrent', 1, 1, ...)

    return AnonymousObservable(function(observer)

        local activeCount = 0
        local group = CompositeDisposable()
        local isStopped = false
        local q = {}

        local function subscribe(xs)

            local subscription = SingleAssignmentDisposable()
            group:add(subscription)

            subscription:setDisposable(xs:subscribe(
                function(...)
                    return observer:onNext(...)
                end, function(...)
                    return observer:onError(...)
                end, function()
                    group:remove(subscription)
                    if #q > 0 then
                        local val = table.remove(q, 1)
                        subscribe(val)
                    else
                        activeCount = activeCount - 1
                        if isStopped and activeCount == 0 then
                            observer:onCompleted()
                        end
                    end
                end)
            )

        end

        group:add(self:subscribe(
            function(innerSource)
                if activeCount < maxConcurrent then
                    activeCount = activeCount + 1
                    subscribe(innerSource)
                else
                    q[#q + 1] = innerSource
                end
            end,
            function(...)
                return observer:onError(...)
            end,
            function()
                isStopped = true
                if activeCount == 0 then
                    observer:onCompleted()
                end
            end)
        )

        return group

    end)

end
