#include "aoc.h"
#include <functional>
#include <map>
#include <set>

namespace aoc2022 {

std::map<line_view, monkey_math> monkeys;

struct monkey_tree {
  monkey_math* name;
  monkey_tree* left = nullptr;
  monkey_tree* right = nullptr;

  monkey_tree(monkey_math* n) : name(n) {}
};

void find(monkey_tree* t, line_view name, monkey_tree** m, std::set<monkey_math*>& ms) {
  if (*m == nullptr && t->left)
    find(t->left, name, m, ms);
  if (*m == nullptr && t->right)
    find(t->right, name, m, ms);

  if (t->name->op[0] == name) {
    *m = t;
    ms.insert(t->name);
  }
  if (*m != nullptr) {
    if (t->left == *m)
      ms.insert(t->name);
    if (t->right == *m)
      ms.insert(t->name);
    *m = t;
  }
}

int64_t get_value(line_view name, monkey_tree* t) {
  auto& m = monkeys[name];
  if (m.value == INT64_MAX) {
    t->left = new monkey_tree{&monkeys[m.op[1]]};
    t->right = new monkey_tree{&monkeys[m.op[2]]};

    int64_t v1 = get_value(m.op[1], t->left);
    int64_t v2 = get_value(m.op[2], t->right);
    std::function<int64_t(int64_t, int64_t)> fs[] = {
        [](int64_t x1, int64_t x2) { return x1 + x2; },
        [](int64_t x1, int64_t x2) { return x1 - x2; },
        [](int64_t x1, int64_t x2) { return x1 * x2; },
        [](int64_t x1, int64_t x2) { return x1 / x2; },
    };
    m.value = fs[(int)m.todo](v1, v2);
  }
  return m.value;
}

int64_t get_value(int64_t value, const std::set<monkey_math*>& ms, monkey_tree* t) {
  // std::cout << t->name->op[0] << " needs to be " << value << std::endl;
  if (t->name->op[0] == "humn") {
    return value;
  } else {
    bool b = ms.find(t->left->name) == ms.end();
    int64_t v0 = b ? monkeys[t->left->name->op[0]].value : monkeys[t->right->name->op[0]].value;
    monkey_tree* t0 = b ? t->right : t->left;
    switch (t->name->todo) {
    case oper::add:
      return get_value(value - v0, ms, t0);
    case oper::minus:
      return get_value(b ? v0 - value : value + v0, ms, t0);
    case oper::multiply:
      return get_value(value / v0, ms, t0);
    case oper::division:
      return get_value(b ? v0 / value : value * v0, ms, t0);
    case none: break;
    }
  }
  return 0;
}

std::pair<int64_t, int64_t> day21(line_view file) {
  per_line(file, [](line_view lv) {
    auto m = monkey_math{lv};
    // m.print();
    monkeys.insert({m.op[0], m});
    return true;
  });
  auto& m = monkeys["root"];
  monkey_tree tree{&m};
  int64_t n1 = get_value("root", &tree);

  monkey_tree* target = nullptr;
  std::set<monkey_math*> ms;
  find(&tree, "humn", &target, ms);

  auto n2 = get_value(monkeys["qwqj"].value, ms, tree.left);
  // auto n2 = get_value(monkeys["sjmn"].value, ms, tree.left);
  return {n1, n2};
}
} // namespace aoc2022