#pragma once

#include <functional>

#include "Supplier.h"

namespace Framework
{
    template<typename T> class Stream
    {
    private:
        Supplier<T>* supplier;

    public:
        Stream(Supplier<T>* supplier)
            : supplier(supplier)
        {}

        ~Stream()
        {
            supplier->release();
        }

        Stream<T> filter(std::function<bool(T)> f)
        {
            return Stream<T>(new FilteredSupplier<T>(
                dynamic_cast<Supplier<T>*>(supplier->getThis()), f));
        }

        template<typename R> Stream<R> map(std::function<R(T)> f)
        {
            return Stream<R>(new MappedSupplier<T, R>(
                dynamic_cast<Supplier<T>*>(supplier->getThis()), f));
        }

        template<typename R> Stream<R> flatMap(std::function<Stream<R>(T)> f)
        {
            return Stream<R>(new FlatMappedSupplier<T, R>(
                dynamic_cast<Supplier<T>*>(supplier->getThis()),
                [this, f](T arg) {
                    return dynamic_cast<Supplier<R>*>(
                        f(arg).zSupplier()->getThis());
                }));
        }

        void forEach(std::function<void(T)> action)
        {
            while (true)
            {
                try
                {
                    action(supplier->next());
                } catch (Exceptions::EndOfSupplier)
                {
                    return;
                }
            }
        }

        template<typename R>
        R collect(std::function<R(R, T)> combinator, R initialValue)
        {
            while (true)
            {
                try
                {
                    initialValue = combinator(initialValue, supplier->next());
                } catch (Exceptions::EndOfSupplier)
                {
                    return initialValue;
                }
            }
        }

        Supplier<T>* zSupplier()
        {
            return supplier;
        }
    };
} // namespace Framework