#include <cstdint>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <variant>
#include <vector>

#include <boost/phoenix.hpp>
#include <boost/spirit/include/qi.hpp>

#include <doctest.h>

#include <z3++.h>

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

namespace {

namespace phx = boost::phoenix;
namespace qi = boost::spirit::qi;

enum class Op { Add, Sub, Mul, Div };

struct Expr {
  std::string lhs;
  std::string rhs;
  Op op;
};

struct Entry {
  std::string lvalue;
  std::variant<std::int64_t, Expr> rvalue;
};

using Input = std::vector<Entry>;

// Input file grammar
class Grammar : public qi::grammar<std::string::const_iterator, Input()> {
  qi::rule<iterator_type, std::string()> variable;
  qi::rule<iterator_type, Op()> op;
  qi::rule<iterator_type, Expr()> expr;
  qi::rule<iterator_type, std::variant<std::int64_t, Expr>()> rhs;
  qi::rule<iterator_type, Entry()> line;
  qi::rule<iterator_type, Input()> input;

public:
  Grammar() : Grammar::base_type{input} {
    using namespace qi::labels;

    variable = qi::as_string[+qi::alpha];
    op =
      qi::string("+") [ _val = Op::Add ] |
      qi::string("-") [ _val = Op::Sub ] |
      qi::string("*") [ _val = Op::Mul ] |
      qi::string("/") [ _val = Op::Div ];
    expr =
      variable [phx::bind(&Expr::lhs, _val) = _1] >> " " >>
      op [phx::bind(&Expr::op, _val) = _1] >> " " >>
      variable [phx::bind(&Expr::rhs, _val) = _1];
    rhs = expr | qi::long_long;
    line =
      variable [ phx::bind(&Entry::lvalue, _val) = _1] >> ": " >>
      rhs [ phx::bind(&Entry::rvalue, _val) = _1] >> "\n";
    input = *line;
  }
};

auto Eval(
  std::unordered_map<std::string, double> & values,
  std::unordered_map<std::string, Expr> const& exprs,
  std::string const& var)
  -> double
{
  if (auto const it = values.find(var); it != values.end()) {
    return it->second;
  }
  auto const& [lhs,rhs,op] = exprs.at(var);
  auto const l = Eval(values, exprs, lhs);
  auto const r = Eval(values, exprs, rhs);
  double o;
  switch (op) {
    case Op::Add: o = l + r; break;
    case Op::Sub: o = l - r; break;
    case Op::Mul: o = l * r; break;
    case Op::Div: o = l / r; break;
  }
  values.emplace(var, o);
  return o;
}

auto Part1(Input const& input) -> std::int64_t
{
  std::unordered_map<std::string, double> values;
  std::unordered_map<std::string, Expr> exprs;
  for (auto const& entry : input) {
    std::visit(overloaded {
      [&](Expr const& expr) { exprs.emplace(entry.lvalue, expr); },
      [&](std::int64_t val) { values.emplace(entry.lvalue, val); }
    }, entry.rvalue);
  }
  return Eval(values, exprs, "root");
}

auto Part2(Input const& input) -> std::int64_t
{
  auto ctx = z3::context{};
  auto const solve_eqs = z3::tactic{ctx, "solve-eqs"};
  auto const smt = z3::tactic{ctx, "smt"};
  auto solver = (solve_eqs & smt).mk_solver();

  std::unordered_map<std::string, z3::expr> constants;
  for (auto const& entry : input) {
    constants.emplace(entry.lvalue, ctx.int_const(entry.lvalue.c_str()));
  }

  for (auto const& entry : input) {
    if (entry.lvalue == "root") {
      if (auto * const expr = std::get_if<Expr>(&entry.rvalue)) {
        solver.add(constants.at(expr->lhs) == constants.at(expr->rhs));
      } else {
        throw std::runtime_error{"malformed root"};
      }
    } else if (entry.lvalue != "humn") {
      auto const rhs =
        std::visit(overloaded {
          [&](std::int64_t literal) {
            return ctx.int_val(literal);
          },
          [&](Expr const& expr) {
            auto const l = constants.at(expr.lhs);
            auto const r = constants.at(expr.rhs);
            switch (expr.op) {
              case Op::Add: return l + r;
              case Op::Sub: return l - r;
              case Op::Mul: return l * r;
              case Op::Div:
                solver.add(l % r == 0);
                return l / r;
            }
          }
        }, entry.rvalue);
      solver.add(constants.at(entry.lvalue) == rhs);
    }
  }

  if (solver.check() != z3::sat) {
    throw std::runtime_error{"no solution to part 2"};
  }

  return solver.get_model().eval(constants.at("humn")).get_numeral_int64();
}

} // namespace

TEST_SUITE("2022-21 examples") {
  TEST_CASE("documented example") {
    auto in = std::istringstream{
      "root: pppw + sjmn\n"
      "dbpl: 5\n"
      "cczh: sllz + lgvd\n"
      "zczc: 2\n"
      "ptdq: humn - dvpt\n"
      "dvpt: 3\n"
      "lfqf: 4\n"
      "humn: 5\n"
      "ljgn: 2\n"
      "sjmn: drzm * dbpl\n"
      "sllz: 4\n"
      "pppw: cczh / lfqf\n"
      "lgvd: ljgn * ptdq\n"
      "drzm: hmdt - zczc\n"
      "hmdt: 32\n"
    };
    auto const input = aocpp::ParseGrammar(Grammar(), in);
    CHECK(152 == Part1(input));
    CHECK(301 == Part2(input));
  }
}

auto Main(std::istream & in, std::ostream & out) -> void
{  
  auto const input = aocpp::ParseGrammar(Grammar(), in);
  out << "Part 1: " << Part1(input) << std::endl;
  out << "Part 2: " << Part2(input) << std::endl;
}