#include "mosh-me/pipeline.hpp"
#include <catch2/catch_all.hpp>
#include <spdlog/fmt/fmt.h>
#include <spdlog/fmt/ranges.h>

using namespace mosh_me;

class CountProducer : public AsPipeSource<CountProducer, int> {
private:
  int counter_ = 0;

public:
  PipeData<int> pull() noexcept override {
    if (auto i = counter_++; i < 10) {
      return i;
    } else {
      return std::unexpected(PipeStatus::PIPE_EOF);
    }
  }
};

class Doubler : public AsPipeLink<Doubler, int, int> {
public:
  PipeData<int> pull() noexcept override {
    return fetchInput().and_then([](int x) -> PipeData<int> { return x * 2; });
  }
};

class Multiplier : public AsPipeLink<Multiplier, int, int> {
  int f_;

public:
  Multiplier(int f) : f_(f) {}
  PipeData<int> pull() noexcept override {
    return fetchInput().and_then(
        [this](int x) -> PipeData<int> { return f_ * x; });
  }
};

class Sum2 : public AsPipeLink<Sum2, int, int> {
private:
  std::optional<int> last_;

public:
  PipeData<int> pull() noexcept override {
    if (auto x = fetchInput(); x) {
      if (last_) {
        auto sum = *last_ + *x;
        last_ = std::nullopt;
        return sum;
      } else {
        last_ = *x;
        return pull();
      }
    }
    return std::unexpected(PipeStatus::PIPE_EOF);
  }
};

TEST_CASE("Pipeline") {
  auto pipe = Sum2::link() << Doubler::link() << Multiplier::link(10)
                           << CountProducer::link();

  auto count = 0;
  while (auto s = pipe.pull()) {
    REQUIRE(*s == ((count * 2) * 20) + ((count * 2 + 1) * 20));
    count++;
  }

  REQUIRE(count == 5);
}