local Rx = require 'Rx'

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

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

local function timeoutWithSelector(source, firstTimeout, timeoutDurationSelector, other)

    if other == nil or type(other) == 'string' then
        other = Observable.throw(other or "Timeout")
    end

    return AnonymousObservable(function(observer)

        local subscription = SerialDisposable()
        local timer = SerialDisposable()
        local original = SingleAssignmentDisposable()

        subscription:setDisposable(original)

        local id = 0
        local switched = false

        local function setTimer(timeout)

            local myId = id
            local d = SingleAssignmentDisposable()

            local function timerWins()

                  switched = myId == id
                  return switched

            end

            timer:setDisposable(d)

            d:setDisposable(timeout:subscribe(
                function()

                    if timerWins() then
                        subscription:setDisposable(other:subscribe(observer))
                    end

                    d:dispose()

                end,
                function(...)

                    if timerWins() then
                        observer:onError(...)
                    end

                end,
                function()

                    if timerWins() then
                        subscription:setDisposable(other:subscribe(observer))
                    end

                end))

        end

        setTimer(firstTimeout)

        local function observableWins()

            local result = not switched
            if result then
                id = id + 1
            end

            return result

        end

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

                if observableWins() then

                    id = id + 1

                    observer:onNext(...)

                    local success, timeoutOrError = pcall(timeoutDurationSelector, ...)

                    if not success then
                        return observer:onError(timeoutOrError)
                    end

                    setTimer(timeoutOrError)

                end

            end,
            function(...)

                if observableWins() then

                    id = id + 1
                    observer:onError(...)

                end

            end,
            function()

                if observableWins() then

                    id = id + 1
                    observer:onCompleted()

                end

            end))

        return CompositeDisposable(subscription, timer)

    end, source)

end

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

local function timeout(source, dueTime, other, scheduler)

    if Scheduler:isClassOf(other) then
        scheduler = other
        other = nil
    end

    if other == nil or type(other) == 'string' then
        other = Observable.throw(other or "Timeout")
    end

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

    return AnonymousObservable(function(observer)

        local id = 0
        local original = SingleAssignmentDisposable()
        local subscription = SerialDisposable()
        local timer = SerialDisposable()

        subscription:setDisposable(original)

        local function createTimer()

            local myId = id

            timer:setDisposable(scheduler:scheduleFuture(dueTime, function()

                if id == myId then
                    subscription:setDisposable(other:subscribe(observer))
                end

            end))

        end

        createTimer()

        original:setDisposable(source:subscribe(
            function(...)
                id = id + 1
                observer:onNext(...)
                createTimer()
            end,
            function(...)
                id = id + 1
                observer:onError(...)
            end,
            function()
                id = id + 1
                observer:onCompleted()
            end)
        )

        return CompositeDisposable(subscription, timer)

    end, source)

end

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

return function(self, arg1, arg2, arg3)

    local arg1Type = type(arg1)

    if arg1Type == 'number' then

        return timeout(self, arg1, arg2, arg3)
            -- dueTime, other, selector

    elseif arg1Type == 'function' then

        return timeoutWithSelector(self, Observable.never(), arg1, arg2)
            -- firstTimeout, timeoutDurationSelector, other

    elseif Observable:isClassOf(arg1) and type(arg2) == 'function' then

        return timeoutWithSelector(self, arg1, arg2, arg3)
            -- firstTimeout, timeoutDurationSelector, other

    else

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

    end

end
