#include <algorithm>
#include <cctype>
#include <cstdint>
#include <deque>
#include <iostream>
#include <sstream>
#include <iterator>
#include <map>
#include <optional>
#include <variant>
#include <vector>

#include <doctest.h>

#include <aocpp/Startup.hpp>
#include <aocpp/Overloaded.hpp>

namespace {


using Name = std::string;
using Value = std::int64_t;

struct Literal { Value value; };
struct Variable { Name name; };
using Expression = std::variant<Literal, Variable>;

struct Send { Value sent; };
struct Receive { Value & target; };
using Effect = std::variant<Send, Receive>;

struct Set { Name x; Expression y; };
struct Add { Name x; Expression y; };
struct Mul { Name x; Expression y; };
struct Mod { Name x; Expression y; };
struct Jgz { Expression x; Expression y; };
struct Rcv { Name x; };
struct Snd { Expression x; };
using Instruction = std::variant<Set, Add, Mul, Mod, Jgz, Rcv, Snd>;

struct State {
  std::map<Name, Value> registers;
  std::size_t pc;
};

auto IsName(std::string const& word) {
  return std::all_of(word.begin(), word.end(), [](auto c) { return std::isalpha(c); });
}

auto ParseName(std::istream & in) -> Name
{
  std::string word;
  in >> word;
  if (IsName(word)) { return word; }
  throw std::runtime_error{"bad name"};
}

auto ParseExpression(std::istream & in) -> Expression
{
  std::string word;
  in >> word;
  if (IsName(word)) {
    return Variable{word};
  } else {
    return Literal{std::stoll(word)};
  }
}

auto Parse(std::istream & in) -> std::vector<Instruction> {
  std::vector<Instruction> result;
  std::string op;
  while (in >> op) {
    if ("set" == op) {
      auto x = ParseName(in);
      auto y = ParseExpression(in);
      result.push_back(Set{x,y});
    } else if ("add" == op) {
      auto x = ParseName(in);
      auto y = ParseExpression(in);
      result.push_back(Add{x,y});
    } else if ("mul" == op) {
      auto x = ParseName(in);
      auto y = ParseExpression(in);
      result.push_back(Mul{x,y});
    } else if ("mod" == op) {
      auto x = ParseName(in);
      auto y = ParseExpression(in);
      result.push_back(Mod{x,y});
    } else if ("jgz" == op) {
      auto x = ParseExpression(in);
      auto y = ParseExpression(in);
      result.push_back(Jgz{x,y});
    } else if ("rcv" == op) {
      auto x = ParseName(in);
      result.push_back(Rcv{x});
    } else if ("snd" == op) {
      auto x = ParseExpression(in);
      result.push_back(Snd{x});
    } else {
      throw std::runtime_error{"unknown op"};
    }
  }
  return result;
}

auto Eval(State const& state, Expression expression) -> Value {
  return std::visit(overloaded{
    [](Literal l) -> Value { return l.value; },
    [&](Variable v) -> Value {
      if (auto it = state.registers.find(v.name); it != state.registers.end()) {
        return it->second;
     } else {
      return 0; }
    },
  }, expression);
}

auto Step(State & state, Instruction const& instruction) -> std::optional<Effect> {
  using R = std::optional<Effect>;
  return std::visit(overloaded{
    [&](Set instruction) -> R { state.registers[instruction.x]  = Eval(state, instruction.y); state.pc++; return {}; },
    [&](Add instruction) -> R { state.registers[instruction.x] += Eval(state, instruction.y); state.pc++; return {}; },
    [&](Mul instruction) -> R { state.registers[instruction.x] *= Eval(state, instruction.y); state.pc++; return {}; },
    [&](Mod instruction) -> R { state.registers[instruction.x] %= Eval(state, instruction.y); state.pc++; return {}; },
    [&](Jgz instruction) -> R { state.pc += Eval(state, instruction.x) > 0 ? Eval(state, instruction.y) : 1; return {}; },
    [&](Snd instruction) -> R { state.pc++; return Send{Eval(state, instruction.x)}; },
    [&](Rcv instruction) -> R { state.pc++; return Receive{state.registers[instruction.x]}; },
    },
    instruction);
}

auto BigStep(State & state, std::vector<Instruction> const& program) -> Effect {
  while(state.pc < program.size()) {
    if (auto effect = Step(state, program[state.pc])) {
      return *effect;
    }
  }
  throw std::runtime_error{"program crash"};
}

auto Part1(std::vector<Instruction> const& program) -> Value {
  Value last_sound = -1;
  State state {};

  for (;;) {
    auto effect = BigStep(state, program);
    switch (effect.index()) {
      case 0: last_sound = std::get<0>(effect).sent; break;
      case 1: if (std::get<1>(effect).target > 0) { return last_sound; } break;
    }
  }

  throw std::runtime_error{"part 1 failed"};
}

auto Spin(
  State & state,
  std::vector<Instruction> const& program,
  std::deque<Value> & input,
  std::deque<Value> & output
) -> std::pair<Value*, std::size_t>
{
  std::size_t n = 0;
  for(;;) {
    auto effect = BigStep(state, program);
    switch (effect.index()) {
      case 0:
        output.push_back(std::get<0>(effect).sent);
        n++;
        break;
      case 1: {
        auto & target = std::get<1>(effect).target;
        if (input.empty()) {
          return {&target, n};
        } else {
          target = input.front();
          input.pop_front();
        }
        break;
      }
    }
  }  
}

auto Part2(std::vector<Instruction> const& program) -> std::size_t {
  State m0{};
  State m1{{{"p",1}}};

  std::deque<Value> inputs0;
  std::deque<Value> inputs1;

  auto [stuck0, _] = Spin(m0, program, inputs0, inputs1);
  auto [stuck1, n] = Spin(m1, program, inputs1, inputs0);

  while(!(inputs0.empty() && inputs1.empty())) {
    if (!inputs0.empty()) {
      *stuck0 = inputs0.front(); inputs0.pop_front();
      stuck0 = Spin(m0, program, inputs0, inputs1).first;
    }
    if (!inputs1.empty()) {
      *stuck1 = inputs1.front(); inputs1.pop_front();
      auto res = Spin(m1, program, inputs1, inputs0);
      stuck1 = res.first;
      n += res.second;
    }
  }
  return n;
}

} // namespace

TEST_SUITE("documented examples") {
  TEST_CASE("part 1") {
    std::istringstream in {
      "set a 1\n"
      "add a 2\n"
      "mul a a\n"
      "mod a 5\n"
      "snd a\n"
      "set a 0\n"
      "rcv a\n"
      "jgz a -1\n"
      "set a 1\n"
      "jgz a -2\n"};
    REQUIRE(Part1(Parse(in)) == 4);
  }
  TEST_CASE("part 2") {
    std::istringstream in {
      "snd 1\n"
      "snd 2\n"
      "snd p\n"
      "rcv a\n"
      "rcv b\n"
      "rcv c\n"
      "rcv d\n"};
    REQUIRE(Part2(Parse(in)) == 3);
  }
}

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