#include <algorithm>
#include <cctype>
#include <cstdint>
#include <deque>
#include <iostream>
#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 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>;

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

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

struct Machine {
  std::vector<Instruction> const& program_;
  std::map<Name, Value> registers_;
  std::size_t pc_;

public:
  Machine(std::vector<Instruction> const& program)
    : program_{program}, registers_{}, pc_{}
  {}

  auto Eval(Expression expression) -> Value {
    return std::visit(overloaded{
      [](Literal l) { return l.value; },
      [&](Variable v) { return registers_[v.name]; },
    }, expression);
  }

  auto Step() -> Effect {
    while (pc_ < program_.size()) {
      if (auto effect = std::visit(overloaded{
        [&](Set instruction) -> std::optional<Effect> { registers_[instruction.x]  = Eval(instruction.y); pc_++; return {}; },
        [&](Add instruction) -> std::optional<Effect> { registers_[instruction.x] += Eval(instruction.y); pc_++; return {}; },
        [&](Mul instruction) -> std::optional<Effect> { registers_[instruction.x] *= Eval(instruction.y); pc_++; return {}; },
        [&](Mod instruction) -> std::optional<Effect> { registers_[instruction.x] %= Eval(instruction.y); pc_++; return {}; },
        [&](Jgz instruction) -> std::optional<Effect> { pc_ += Eval(instruction.x) > 0 ? Eval(instruction.y) : 1; return {}; },
        [&](Snd instruction) -> std::optional<Effect> { pc_++; return Send{Eval(instruction.x)}; },
        [&](Rcv instruction) -> std::optional<Effect> { pc_++; return Receive{registers_[instruction.x]}; },
      },
      program_[pc_]))
      {
        return *effect;
      }
    }
    return Halt{};
  }
};

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

  for(;;) {
    auto effect = m.Step();
    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;
      default: throw std::runtime_error{"program halted"};
    }
  }
}

auto Spin(Machine & m, std::deque<Value> & input, std::deque<Value> & output) -> Value* {
  for(;;) {
    auto effect = m.Step();
    switch (effect.index()) {
      case 0: output.push_back(std::get<0>(effect).sent); break;
      case 1: {
        auto & target = std::get<1>(effect).target;
        if (input.empty()) {
          return &target;
        } else {
          target = input.front();
          input.pop_front();
        }
        break;
      }
      default: throw std::runtime_error{"program halted"};
    }
  }  
}

auto Part2(std::vector<Instruction> const& program) -> std::size_t {
  Machine m0 { program };
  Machine m1 { program };
  m1.registers_["p"] = 1;

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

  auto stuck0 = Spin(m0, inputs0, inputs1);
  auto stuck1 = Spin(m1, inputs1, inputs0);
  
  std::size_t result = inputs0.size();

  while(!(inputs0.empty() && inputs1.empty())) {
    if (!inputs0.empty()) {
      *stuck0 = inputs0.front(); inputs0.pop_front();
      stuck0 = Spin(m0, inputs0, inputs1);
    }
    if (!inputs1.empty()) {
      *stuck1 = inputs1.front(); inputs1.pop_front();
      result -= inputs0.size();
      stuck1 = Spin(m1, inputs1, inputs0);
      result += inputs0.size();
    }
  }

  return result;
}

} // namespace

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