local Rx = require 'Rx'

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

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

return function(source, timeSpan, count, scheduler)

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

    return AnonymousObservable(function(observer)
        local timerD = SerialDisposable()
        local groupDisposable = CompositeDisposable(timerD)
        local refCountDisposable = RefCountDisposable(groupDisposable)
        local n = 0
        local windowId = 0
        local s = Subject()
        groupDisposable:add(s)

        local function createTimer(id)
            local m = SingleAssignmentDisposable()
            timerD:setDisposable(m)
            m:setDisposable(scheduler:scheduleFuture(timeSpan, function()
                if id ~= windowId then return end
                n = 0
                windowId = windowId + 1
                local newId = windowId
                s:onCompleted()
                s = Subject()
                groupDisposable:add(s)
                observer:onNext(addRef(s, refCountDisposable))
                createTimer(newId)
            end))
        end

        observer:onNext(addRef(s, refCountDisposable))
        createTimer(0)

        groupDisposable:add(source:subscribe(
            function(...)
                local newId = 0
                local newWindow = false

                s:onNext(...)
                n = n + 1
                if n == count then
                    newWindow = true
                    n = 0
                    windowId = windowId + 1
                    newId = windowId
                    s:onCompleted()
                    s = Subject()
                    groupDisposable:add(s)
                    observer:onNext(addRef(s, refCountDisposable))
                end
                if newWindow then
                    createTimer(newId)
                end
            end,
            function(...)
                s:onError(...)
                observer:onError(...)
            end, function()
                s:onCompleted()
                observer:onCompleted()
            end))

        return refCountDisposable
    end, source)
end
