local Rx = require 'Rx'

local addRef = Rx.internals.addRef
local AnonymousObservable = Rx.AnonymousObservable
local RefCountDisposable = Rx.RefCountDisposable
local CompositeDisposable = Rx.CompositeDisposable
local Scheduler = Rx.Scheduler
local SerialDisposable = Rx.SerialDisposable
local SingleAssignmentDisposable = Rx.SingleAssignmentDisposable
local Subject = Rx.Subject

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

return function(source, timeSpan, timeShiftOrScheduler, scheduler)
    local timeShift
    if timeShiftOrScheduler == nil then
        timeShift = timeSpan
    end
    if not Scheduler:isClassOf(scheduler) then
        scheduler = Rx.config.defaultScheduler
    end
    if type(timeShiftOrScheduler) == 'number' then
        timeShift = timeShiftOrScheduler
    elseif Scheduler:isClassOf(timeShiftOrScheduler) then
        timeShift = timeSpan
        scheduler = timeShiftOrScheduler
    end

    return AnonymousObservable(function(observer)
        local nextShift = timeShift
        local nextSpan = timeSpan
        local q = {}
        local timerD = SerialDisposable()
        local totalTime = 0
        local groupDisposable = CompositeDisposable(timerD)
        local refCountDisposable = RefCountDisposable(groupDisposable)

        local function createTimer()
            local m = SingleAssignmentDisposable()
            local isSpan = false
            local isShift = false

            timerD:setDisposable(m)
            if nextSpan == nextShift then
                isSpan = true
                isShift = true
            elseif nextSpan < nextShift then
                isSpan = true
            else
                isShift = true
            end

            local newTotalTime
            if isSpan then
                newTotalTime = nextSpan
            else
                newTotalTime = nextShift
            end

            local ts = newTotalTime - totalTime
            totalTime = newTotalTime

            if isSpan then
                nextSpan = nextSpan + timeShift
            end
            if isShift then
                nextShift = nextShift + timeShift
            end

            m:setDisposable(scheduler:scheduleFuture(ts, function()
                if isShift then
                    local s = Subject()
                    table.insert(q, s)
                    groupDisposable:add(s)
                    observer:onNext(addRef(s, refCountDisposable))
                end
                if isSpan then
                    local o = table.remove(q, 1)
                    o:onCompleted()
                    groupDisposable:remove(o)
                end
                createTimer()
            end))
        end

        local subject = Subject()
        table.insert(q, subject)
        groupDisposable:add(subject)
        observer:onNext(addRef(subject, refCountDisposable))
        createTimer()
        groupDisposable:add(source:subscribe(
            function(...)
                local len = #q
                for i = 1, len, 1 do
                    local subject = q[i]
                    subject:onNext(...)
                end
            end,
            function(...)
                local len = #q
                for i = 1, len, 1 do
                    q[i]:onError(...)
                end
                observer:onError(...)
            end,
            function()
                local len = #q
                for i = 1, len, 1 do
                    q[i]:onCompleted()
                end
                observer:onCompleted()
            end))

        return refCountDisposable

    end, source)
end
