local Rx = require 'Rx'

local AbstractObserver = Rx.AbstractObserver
local ObservableBase = Rx.ObservableBase
local CompositeDisposable = Rx.CompositeDisposable
local SingleAssignmentDisposable = Rx.SingleAssignmentDisposable
local SerialDisposable = Rx.SerialDisposable
local observableCatch = Rx.Observable.catch

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

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

    AbstractObserver,
        -- superclass

    _observer = false,
        -- (Observer)

    _subscription = false,
        -- (Disposable)

    _handler = false,
        -- (function)

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

    __init = function(self, observer, subscription, handler)
        self._observer = observer
        self._subscription = subscription
        self._handler = handler
    end,

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

    next = function(self, ...)
        self._observer:onNext(...)
    end,

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

    error = function(self, ...)

        local result

        local success, msg = pcall(self._handler, ...)

        if success then

            result = msg

            local d = SingleAssignmentDisposable()
            self._subscription:setDisposable(d)
            d:setDisposable(result:subscribe(self._observer))

        else

            self._observer:onError(msg)

        end

    end,

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

    completed = function(self)
        self._observer:onCompleted()
    end,

}

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

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

    ObservableBase,
        -- superclass

    _source = false,
        -- (Observable)

    _fn = false,
        -- (function)

    __init = function(self, source, fn)

        self._source = source
        self._fn = fn

    end,

    subscribeCore = function(self, observer)

        local d1 = SingleAssignmentDisposable()

        local subscription = SerialDisposable()

        subscription:setDisposable(d1)

        local catchObserver = CatchObserver(observer, subscription, self._fn)
        d1:setDisposable(self._source:subscribe(catchObserver))

        return CompositeDisposable(subscription, catchObserver)

    end,

}

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

return function(self, handlerOrSecond)
    if type(handlerOrSecond) == 'function' then
        return CatchObservable(self, handlerOrSecond)
    else
        return observableCatch(self, handlerOrSecond)
    end
end
