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

#include <doctest.h>

#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;
}

auto Part1(Machine machine) -> ValueType {
  return Optimize(std::move(machine), {0, 1, 2, 3, 4}, Compute1);
}

auto Part2(Machine machine) -> ValueType {
  return Optimize(std::move(machine), {5, 6, 7, 8, 9}, Compute2);
}

} // namespace

TEST_SUITE("documented examples") {
  TEST_CASE("part 1") {
    auto eval = [](std::vector<ValueType> pgm) { return Part1(Machine{pgm}); };
    REQUIRE(eval({3,15,3,16,1002,16,10,16,1,16,15,15,4,15,99,0,0}) == 43210);
    REQUIRE(eval({3,23,3,24,1002,24,10,24,1002,23,-1,23,101,5,23,23,1,24,23,23,4,23,99,0,0}) == 54321);
    REQUIRE(eval({3,31,3,32,1002,32,10,32,1001,31,-2,31,1007,31,0,33,1002,33,7,33,1,33,31,31,1,32,31,31,4,31,99,0,0,0}) == 65210);
  }
  TEST_CASE("part 2") {
    auto eval = [](std::vector<ValueType> pgm) { return Part2(Machine{pgm}); };
    REQUIRE(eval({3,26,1001,26,-4,26,3,27,1002,27,2,27,1,27,26,27,4,27,1001,28,-1,28,1005,28,6,99,0,0,5}) == 139629729);
    REQUIRE(eval({3,52,1001,52,-5,52,3,53,1,52,56,54,1007,54,5,55,1005,55,26,1001,54,-5,54,1105,1,12,1,53,54,53,1008,54,0,55,1001,55,1,55,2,53,55,53,4,53,1001,56,-1,56,1005,56,6,99,0,0,0,0,10})== 18216);
  }
}

auto main(int argc, char** argv) -> int {
  auto machine = Machine{ParseStream(aocpp::Startup(argc, argv))};
  std::cout << "Part 1: " << Part1(machine) << std::endl;
  std::cout << "Part 2: " << Part2(std::move(machine)) << std::endl;
}