local Rx = require 'Rx'

local AnonymousObservable = Rx.AnonymousObservable
local CompositeDisposable = Rx.CompositeDisposable
local Disposable = Rx.Disposable
local Observable = Rx.Observable
local SingleAssignmentDisposable = Rx.SingleAssignmentDisposable

local checkDisposed = Rx.Disposable.checkDisposed
local disposableCreate = Rx.Disposable.create

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

local connectableObservablePrototype = {

    Observable,
    Disposable,
        -- superclasses

    _hasSubscription = false,

    _subscription = false,
        -- (Subscription)

    _sourceObservable = false,
        -- (Observable)

    _subject = false,
        -- (Subject)

    _subjectFactory = false,
        -- (function)

    _selector = false,
        -- (function)

}

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

function connectableObservablePrototype:__init(source, subjectOrSubjectFactory, selector)

    self._sourceObservable = source:asObservable()

    if type(subjectOrSubjectFactory) == 'function' then
        self._subjectFactory = subjectOrSubjectFactory
    else
        self._subject = subjectOrSubjectFactory
    end

    self._selector = selector

end

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

function connectableObservablePrototype:_subscribe(observer)

    if self._subjectFactory and not self._subject then
        self._subject = self._subjectFactory()
    end

    checkDisposed(self._subject)

    local selector = self._selector

    if selector then

        local subscription = selector(self._subject):subscribe(observer)

        self:connect()

        return subscription

    else

        return self._subject:subscribe(observer)

    end

end

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

function connectableObservablePrototype:connect()

    if not self._hasSubscription then

        self._hasSubscription = true

        if self._subjectFactory and not self._subject then
            self._subject = self._subjectFactory()
        end

        local sad = SingleAssignmentDisposable()
        self._subscription = CompositeDisposable(
            sad,
            disposableCreate(function()
                self._hasSubscription = false
            end)
        )

        if not self._subject:isStopped() then

            -- No need to resubcribe to sourceObservable if subject is stopped
            -- as it will just end up ignoring everything.

            sad:setDisposable(self._sourceObservable:subscribe(self._subject))
                -- subscribe() may complete before assigning _subscription so call
                -- it after _subscription has been assigned.

        end

    end

    return self._subscription

end

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

function connectableObservablePrototype:refCount()

    local connectableSubscription
    local count = 0

    return AnonymousObservable(function(observer)

        local subscription = self:subscribe(observer)

        count = count + 1
        if count == 1 then
            connectableSubscription = self:connect()
        end

        return function()

            subscription:dispose()

            count = count - 1
            if count == 0 then
                connectableSubscription:dispose()
            end

        end

    end)

end

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

function connectableObservablePrototype:dispose()

    if not self:isDisposed() then

        if self._hasSubscription then
            self._subscription:dispose()
        end

        if self._subjectFactory and self._subject then

            self._subject:dispose()
            self._subject = nil

        end

        self:Rx_Disposable_dispose()

    end

end

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

return Rx.Class('Rx.ConnectableObservable', connectableObservablePrototype)
