#include <algorithm>
#include <concepts>
#include <iostream>
#include <optional>
#include <utility>
#include <vector>

#include <aocpp/Startup.hpp>
#include <intcode/intcode.hpp>
using namespace intcode;

namespace {

template <class R> auto Controller(Machine machine, R const &params) {
  auto& i = std::get<Input>(Step(machine)).input;
  std::vector<Machine> amps;
  amps.reserve(params.size());
  for (auto p : params) {
    i = p;
    amps.push_back(machine);
  }
  return amps;
}

template <class R>
auto Feed(R &&amps, ValueType start) -> std::optional<ValueType> {
  for (auto &m : amps) {
    auto e = Step(m);
    if (auto i = std::get_if<Input>(&e)) {
      i->input = start;
      start = StepOutput(m);
    } else {
      return {};
    }
  }
  return {start};
}

auto compute1(Machine machine, std::vector<ValueType> const &params)
    -> ValueType {
  return *Feed(Controller(std::move(machine), params), 0);
}
auto compute2(Machine machine, std::vector<ValueType> const &params)
    -> ValueType {
  auto amps = Controller(std::move(machine), params);
  ValueType last = 0;
  while (auto next = Feed(amps, last)) {
    last = *next;
  }
  return last;
}

template <std::invocable<Machine, std::vector<ValueType> const &> F>
auto optimize(Machine machine, std::vector<ValueType> params, F f) {
  ValueType best = 0;
  do {
    best = std::max(best, f(machine, params));
  } while (std::next_permutation(params.begin(), params.end()));
  return best;
}

} // namespace

auto main(int argc, char** argv) -> int {
  auto machine = Machine{ParseStream(aocpp::Startup(argc, argv))};
  std::cout << "Part 1: " << optimize(machine,            {0, 1, 2, 3, 4}, compute1) << std::endl;
  std::cout << "Part 2: " << optimize(std::move(machine), {5, 6, 7, 8, 9}, compute2) << std::endl;
}