local Rx = require 'Rx'

local AnonymousObservable = Rx.AnonymousObservable
local CompositeDisposable = Rx.CompositeDisposable
local Observable = Rx.Observable
local SingleAssignmentDisposable = Rx.SingleAssignmentDisposable

local kNil = {}

local unpack = table.unpack or unpack

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

local function wrapNil(x)
    if x == nil then
        return kNil
    end
    return x
end

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

return function(...)

    local sources = { ... }
    local numSources = select('#', ...)

    assert(type(select(numSources, ...)) ~= 'function', "Rx.Observable.zip: Support for resultSelector not implemented. Try :select() on the result of .zip() without the result selector")

    return AnonymousObservable(function(observer)

        local values = {}
        local sourceResultCounts = {}
            -- The count of how many values have been received for each input.

        local nResultsZipped = 0
            -- The count of how many results have been delivered.

        local numSourcesNotDone = numSources
            -- The count of how many sources that haven't completed.

        for i = 1, numSources do
            sourceResultCounts[i] = 0
        end

        local subscriptions = {}

        for i = 1, numSources do

            local isLastSource = i == numSources

            local function nextValue(...)

                local resultIndex = 1 + sourceResultCounts[i] - nResultsZipped
                sourceResultCounts[i] = sourceResultCounts[i] + 1

                local resultPackedValue = values[resultIndex]
                if not resultPackedValue then
                    resultPackedValue = {}
                    values[resultIndex] = resultPackedValue
                end

                local numValues = select('#', ...)

                if isLastSource then

                    if numValues == 0 then
                        -- If the last observable delivers a zero-length tuple result
                        -- (e.g. none) then treat it as if it were an explicit nil.
                        -- This is similar to how we treat values at the other positions
                        -- so that results don't slide around.  (nones are also nils
                        -- and furthermore extra tuple values are ignored except at
                        -- the last position.)
                        numValues = 1
                    end

                    resultPackedValue.n = numValues + numSources - 1

                    for resultN = 1, numValues do
                        resultPackedValue[i + resultN - 1] = wrapNil(select(resultN, ...))
                    end

                else
                    resultPackedValue[i] = wrapNil(...)
                end

                local nextPackedValue = values[1]
                local allSourcesHaveValue = false

                if nextPackedValue ~= nil then
                    allSourcesHaveValue = true
                    for i = 1, numSources do
                        if nextPackedValue[i] == nil then
                            allSourcesHaveValue = false
                            break
                        end
                    end
                end

                if allSourcesHaveValue then
                    for i = 1, nextPackedValue.n do
                        local value = nextPackedValue[i]
                        if value == kNil then
                            nextPackedValue[i] = nil
                        end
                    end
                    table.remove(values, 1)
                    nResultsZipped = nResultsZipped + 1
                    observer:onNext(unpack(nextPackedValue, 1, nextPackedValue.n))
                elseif numSourcesNotDone == 1 then
                    observer:onCompleted()
                    observer = nil
                    values = nil
                    return
                end

            end

            local function done()
                numSourcesNotDone = numSourcesNotDone - 1
                if numSourcesNotDone == 0 then
                    observer:onCompleted()
                    observer = nil
                    values = nil
                end
            end

            local source = sources[i]
            if Observable:isClassOf(source) then
                local sad = SingleAssignmentDisposable()
                subscriptions[#subscriptions + 1] = sad
                sad:setDisposable(sources[i]:subscribe(
                    nextValue,
                    function(...) return observer:onError(...) end,
                    done))
            else
                for _, value in ipairs(source) do
                    nextValue(value)
                end
                done()
            end
        end

        return CompositeDisposable(unpack(subscriptions))

    end)

end
