local Rx = require 'Rx'

local AbstractObserver = Rx.AbstractObserver
local AnonymousObservable = Rx.AnonymousObservable
local CompositeDisposable = Rx.CompositeDisposable
local ObservableBase = Rx.ObservableBase
local Scheduler = Rx.Scheduler
local SerialDisposable = Rx.SerialDisposable
local SingleAssignmentDisposable = Rx.SingleAssignmentDisposable
local Subject = Rx.Subject

local validateNonNegativeInteger = Rx.internals.NumericArgument.validateNonNegativeInteger

local unpack = table.unpack or unpack

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

local EnumerablePrototype = {
    class = {}
}

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

local CatchErrorObserver = Rx.Class 'Rx.internals.CatchErrorObserver' {

    AbstractObserver,
        -- superclass

    _observer = false,
        -- (Observer)

    _state = false,
        -- (table)

    _recurseFn = false,
        -- (function)

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

    __init = function(self, state, recurseFn)

        self._observer = state.observer
        self._state = state
        self._recurseFn = recurseFn

    end,

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

    next = function(self, ...)

        return self._observer:onNext(...)

    end,

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

    error = function(self, ...)

        local state = self._state

        state.lastError = { n = select('#', ...), ... }

        self._recurseFn(state)

    end,

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

    completed = function(self, ...)

        return self._observer:onCompleted(...)

    end,

}

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

local CatchErrorObservable = Rx.Class 'Rx.internals.CatchErrorObservable' {

    ObservableBase,
        -- superclass

    _source = false,
        -- (Enumerable)

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

    __init = function(self, source)
        self._source = source
    end,

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

    subscribeCore = function(self, observer)

        local source = self._source

        local success, errOrFn, invariant, pairsState = pcall(source.pairs, source)

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

        local subscription = SerialDisposable()

        local state = {
            observer = observer,
            subscription = subscription,
            lastError = false,
            enumerateFn = errOrFn,
            invariant = invariant,
            pairsState = pairsState,
        }

        local function scheduleMethod(state, recurseFn)

            if not state.subscription:isDisposed() then

                local iterVals = { pcall(state.enumerateFn, state.invariant, state.pairsState) }
                local success = iterVals[1]

                if not success then
                    state.observer:onError(unpack(iterVals, 2))
                    return
                end

                state.pairsState = iterVals[2]

                local currentValue = iterVals[3]

                if currentValue == nil then

                    local lastError = state.lastError

                    if lastError then
                        state.observer:onError(unpack(lastError, 1, lastError.n))
                    else
                        state.observer:onCompleted()
                    end

                    return

                end

                local d = SingleAssignmentDisposable()
                state.subscription:setDisposable(d)

                local catchErrorObserver = CatchErrorObserver(state, recurseFn)

                d:setDisposable(CompositeDisposable(
                    currentValue:subscribe(catchErrorObserver),
                    catchErrorObserver))

            end

        end

        local cancelable = Scheduler.currentThread():scheduleRecursive(state, scheduleMethod)

        return CompositeDisposable(subscription, cancelable)

    end,

}

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

function EnumerablePrototype:catch()

    return CatchErrorObservable(self)

end

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

function EnumerablePrototype:catchWhen(notificationHandler)

    return AnonymousObservable(function(observer)

        local success, errOrFn, invariant, state = pcall(self.pairs, self)

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

        local enumerateFn = errOrFn

        local subscription = SerialDisposable()

        local exceptions = Subject()
        local notifier = Subject()
        local handled = notificationHandler(exceptions)
        local notificationDisposable = handled:subscribe(notifier)

        local function nextStepFn(continue)

            local iterVals = { pcall(enumerateFn, invariant, state) }

            local success = iterVals[1]
            if not success then
                observer:onError(unpack(iterVals, 2))
                return
            end

            state = iterVals[2]

            local currentValue = iterVals[3]

            if currentValue == nil then
                observer:onCompleted()
                return
            end

            local outer = SingleAssignmentDisposable()

            local inner = SingleAssignmentDisposable()

            subscription:setDisposable(CompositeDisposable(inner, outer))

            outer:setDisposable(currentValue:subscribe(
                function(...)
                    observer:onNext(...)
                end,
                function(...)

                    inner:setDisposable(notifier:subscribe(
                        continue,
                        function(...)
                            observer:onError(...)
                        end,
                        function()
                            observer:onCompleted()
                        end))

                    if not exceptions:isDisposed() then
                        exceptions:onNext(...)
                    end

                end,
                function()
                    observer:onCompleted()
                end))

        end

        local cancelable = Scheduler.currentThread():scheduleRecursive(
            function(_, continue)
                if not subscription:isDisposed() then
                    nextStepFn(continue)
                end
            end)

        return CompositeDisposable(notificationDisposable, subscription,
            cancelable, exceptions, notifier)

    end)

end

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

local ConcatObserver = Rx.Class 'Rx.internals.ConcatObserver' {

    AbstractObserver,
        -- superclass

    _state = false,
        -- (table)

    _recurseFn = false,
        -- (function)

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

    __init = function(self, state, recurseFn)
        self._state = state
        self._recurseFn = recurseFn
    end,

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

    next = function(self, ...)
        self._state.observer:onNext(...)
    end,

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

    error = function(self, ...)
        self._state.observer:onError(...)
    end,

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

    completed = function(self)
        self._recurseFn(self._state)
    end,

}

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

local ConcatEnumerableObservable = Rx.Class 'ConcatEnumerableObservable' {

    ObservableBase,
        -- superclass

    _source = false,
        -- (Enumerable)

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

    __init = function(self, source)

        self._source = source

    end,

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

    subscribeCore = function(self, observer)

        local source = self._source

        local success, errOrFn, invariant, pairsState = pcall(source.pairs, source)

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

        local subscription = SerialDisposable()

        local state = {
            observer = observer,
            subscription = subscription,
            enumerateFn = errOrFn,
            invariant = invariant,
            pairsState = pairsState,
        }

        local function scheduleMethod(state, recurseFn)

            if not state.subscription:isDisposed() then

                local iterVals = { pcall(state.enumerateFn, state.invariant, state.pairsState) }
                local success = iterVals[1]

                if not success then
                    state.observer:onError(unpack(iterVals, 2))
                    return
                end

                state.pairsState = iterVals[2]

                local currentValue = iterVals[3]

                if currentValue == nil then
                    state.observer:onCompleted()
                    return
                end

                local d = SingleAssignmentDisposable()

                state.subscription:setDisposable(d)

                local observer = ConcatObserver(state, recurseFn)

                local subscribeSuccess, subscriptionOrErr =
                    pcall(currentValue.subscribe, currentValue, observer)

                if not subscribeSuccess then
                    observer:dispose()
                    error(subscriptionOrErr)
                end

                d:setDisposable(CompositeDisposable(subscriptionOrErr, observer))

            end

        end

        local cancelable = Scheduler.currentThread():scheduleRecursive(state, scheduleMethod)

        return CompositeDisposable(subscription, cancelable)

    end,

}

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

function EnumerablePrototype:concat()

    if not self:hasMethod('pairs') then
        error("concat must be called on Enumerable subclass that provides pairs or __pairs", 3)
    end

    return ConcatEnumerableObservable(self)

end

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

local RepEnumerable, OfEnumerable
    -- forward declaration

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

local function second(_, ...)
    return ...
end

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

function EnumerablePrototype.class:rep(...)

    local repeatCount = ...

    if repeatCount == nil then
        repeatCount = -1
    end

    if repeatCount ~= -1 then
        repeatCount = validateNonNegativeInteger('repeatCount', 2, 1, ...)
    end

    return RepEnumerable(repeatCount, second(...))

end

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

function EnumerablePrototype.class:of(source, selector)
    return OfEnumerable(source, selector)
end

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

local Enumerable = Rx.Class("Rx.internals.Enumerable", EnumerablePrototype)

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

local RepEnumerablePrototype = {

    Enumerable,
        -- superclass

    _value = false,
        -- (table)

    _repeatCount = false,
        -- (number)

}

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

function RepEnumerablePrototype:__init(repeatCount, ...)

    self._value = { n = select('#', ...), ... }
    self._repeatCount = repeatCount

end

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

function RepEnumerablePrototype:__pairs()

    local left = self._repeatCount

    return function(self, step)
        if left == 0 then
            return nil
        end
        if left > 0 then
            left = left - 1
        end
        step = step + 1
        return step, unpack(self._value, 1, self._value.n)
    end, self, 0

end

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

RepEnumerable = Rx.Class("Rx.internals.RepEnumerable", RepEnumerablePrototype)

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

local OfEnumerablePrototype = {

    Enumerable,
        -- superclass

    _source = true,
        -- (Observable)

    _selector = true,
        -- (function)

}

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

function OfEnumerablePrototype:__init(source, selector)

    if type(source) ~= 'table' then
        error('Enumerator.of expects source to be an array', 4)
    end

    selector = selector or function(...) return ... end

    if type(selector) ~= 'function' then
        error('Enumerator.of expects selector to be nil or a function', 4)
    end

    self._source = source
    self._selector = selector

end

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

function OfEnumerablePrototype:__pairs()

    return function(self, index)

        local source = self._source
        local n = source.n or #source

        index = index + 1
        if index > n then
            return nil
        end

        return index, self._selector(source[index], index, source)

    end, self, 0

end

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

OfEnumerable = Rx.Class("Rx.internals.OfEnumerable", OfEnumerablePrototype)

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

return Enumerable
