State managing.

StateModel.swift

//
//  StateModel.swift
//
//  Created by Volodymyr Andriienko on 11.04.2022.
//

import Foundation
import RxSwift
import RxCocoa

class StateModel<Action, Event, State>: NSObject {
    let bag = DisposeBag()
    var state: State { stateRelay.value }
    var stateObservable: Observable<State> { stateRelay.asObservable() }

    private let stateRelay: BehaviorRelay<State>
    private let eventRelay = PublishRelay<Event>()
    private let actionRelay = PublishRelay<Action>()
    private let actionScheduler: ImmediateSchedulerType
    private let eventScheduler: ImmediateSchedulerType

    init(
        initial: State,
        actionScheduler: ImmediateSchedulerType = MainScheduler.instance,
        eventScheduler: ImmediateSchedulerType = MainScheduler.instance
    ) {
        self.actionScheduler = actionScheduler
        self.eventScheduler = eventScheduler
        self.stateRelay = BehaviorRelay(value: initial)
        super.init()
    }

    func perform(_ event: Event) {
        eventRelay.accept(event)
    }

    func execute(_ action: Action) {
        actionRelay.accept(action)
    }

    func reduce<Event>(_ handler: @escaping (Event) -> State?) {
        eventRelay
            .compactMap { $0 as? Event }
            .observe(on: eventScheduler)
            .compactMap(handler)
            .observe(on: eventScheduler)
            .bind(to: stateRelay)
            .disposed(by: bag)
    }

    func reduceRun<Event>(_ handler: @escaping (Event) -> Void) {
        eventRelay
            .compactMap { $0 as? Event }
            .observe(on: eventScheduler)
            .compactMap(handler)
            .subscribe(onNext: { _ in })
            .disposed(by: bag)
    }

    func reduceS<Event>(_ handler: @escaping (Event, State) -> State?) {
        eventRelay
            .compactMap { $0 as? Event }
            .observe(on: eventScheduler)
            .withLatestFrom(stateRelay, resultSelector: { ($0, $1) })
            .compactMap(handler)
            .observe(on: eventScheduler)
            .bind(to: stateRelay)
            .disposed(by: bag)
    }

    func reduceRunS<Event>(_ handler: @escaping (Event, State) -> Void) {
        eventRelay
            .compactMap { $0 as? Event }
            .observe(on: eventScheduler)
            .withLatestFrom(stateRelay, resultSelector: { ($0, $1) })
            .compactMap(handler)
            .subscribe(onNext: { _ in })
            .disposed(by: bag)
    }

    func on<Action>(_ handler: @escaping (Action) -> Event?) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .compactMap(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    func onRun<Action>(_ handler: @escaping (Action) -> Void) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .compactMap(handler)
            .subscribe(onNext: { _ in })
            .disposed(by: bag)
    }

    /*
     |  *---------------|
     |       *           ------------|
     |                    *           ---------
     */
    func on<Action>(sequential handler: @escaping (Action) -> Observable<Event>) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .concatMap(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    /*
     | *---------------|
     |      *X
     |                  *--------------|
     |                      *X
     */
    func on<Action>(droppable handler: @escaping (Action) -> Observable<Event>) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .flatMapFirst(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    /*
        | *---X
        |     *------------|
     */
    func on<Action>(restartable handler: @escaping (Action) -> Observable<Event>) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .flatMapLatest(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    /*
     | *--------|
     |   *--------|
     |     *-------------|
     */
    func on<Action>(concurrent handler: @escaping (Action) -> Observable<Event>) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .flatMap(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    func onS<Action>(_ handler: @escaping (Action, State) -> Event?) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .withLatestFrom(stateRelay, resultSelector: { ($0, $1) })
            .compactMap(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    func onRunS<Action>(_ handler: @escaping (Action, State) -> Void) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .withLatestFrom(stateRelay, resultSelector: { ($0, $1) })
            .compactMap(handler)
            .subscribe(onNext: { _ in })
            .disposed(by: bag)
    }

    func onS<Action>(sequential handler: @escaping (Action, State) -> Observable<Event>) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .withLatestFrom(stateRelay, resultSelector: { ($0, $1) })
            .concatMap(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    func onS<Action>(droppable handler: @escaping (Action, State) -> Observable<Event>) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .withLatestFrom(stateRelay, resultSelector: { ($0, $1) })
            .flatMapFirst(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    func onS<Action>(restartable handler: @escaping (Action, State) -> Observable<Event>) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .withLatestFrom(stateRelay, resultSelector: { ($0, $1) })
            .flatMapLatest(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    func onS<Action>(concurrent handler: @escaping (Action, State) -> Observable<Event>) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .withLatestFrom(stateRelay, resultSelector: { ($0, $1) })
            .flatMap(handler)
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }

    // TODO: - Check with async/await
    func onS<Action>(concurrentAsync handler: @escaping (Action, State) async throws -> Event) {
        actionRelay
            .compactMap { $0 as? Action }
            .observe(on: actionScheduler)
            .withLatestFrom(stateRelay, resultSelector: { ($0, $1) })
            .flatMap { action, state in
                Observable<Event>
                    .create { observer in
                        let task = Task {
                            do {
                                let result = try await handler(action, state)
                                observer.onNext(result)
                                observer.onCompleted()
                            } catch {
                                observer.onError(error)
                            }
                        }
                        return Disposables.create {
                            task.cancel()
                        }
                    }
            }
            .observe(on: actionScheduler)
            .bind(to: eventRelay)
            .disposed(by: bag)
    }
}