#include <array>
#include <cstdint>
#include <iostream>
#include <sstream>
#include <variant>
#include <vector>

#include <boost/phoenix.hpp>
#include <boost/range/adaptor/indexed.hpp>
#include <boost/spirit/include/qi.hpp>

#include <doctest.h>

#include <aocpp/Startup.hpp>

namespace {

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

struct Tree;
using Trees = std::vector<Tree>;
using Input = std::vector<std::array<Tree, 2>>;

struct Tree {
  std::variant<std::int64_t, Trees> rep;
};

auto operator<(Tree const& x, std::int64_t y) -> bool;
auto operator<(std::int64_t x, Tree const& y) -> bool;

auto operator<(std::int64_t x, Trees const& y) -> bool {
  switch (y.size()) {
    case 0: return false;
    case 1: return x < y[0];
    default: return !(y[0] < x);
  }
}

auto operator<(Trees const& x, std::int64_t y) -> bool {
  switch (x.size()) {
    case 0: return true;
    default: return x[0] < y;
  }
}

auto operator<(Tree const& x, std::int64_t y) -> bool {
  return std::visit([y](auto const& x_) { return x_ < y; }, x.rep);
}

auto operator<(std::int64_t x, Tree const& y) -> bool {
  return std::visit([x](auto const& y_) { return x < y_; }, y.rep);
}

auto operator<(Tree const& lhs, Tree const& rhs) -> bool {
  return std::visit(std::less(), lhs.rep, rhs.rep);
}

template <typename It>
class Grammar : public qi::grammar<It, Input()> {
  qi::rule<It, std::int64_t> leaf;
  qi::rule<It, Trees()> trees;
  qi::rule<It, Tree()> tree;
  qi::rule<It, std::array<Tree, 2>()> treepair; 
  qi::rule<It, Input()> input;

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

    leaf  = qi::long_long;
    trees = "[" >> -(tree % ",") >> "]";
    tree  = leaf  [ bind(&Tree::rep, _val) = _1 ]
          | trees [ bind(&Tree::rep, _val) = _1 ];
    treepair = (tree >> "\n" >> tree >> "\n") [(_val[0] = _1, _val[1] = _2)];
    input = -(treepair % "\n");
  }
};

auto Parse(std::istream & in) -> Input
{
  Input result;
  std::string content {std::istreambuf_iterator{in}, {}};
  auto b {content.begin()};
  auto const e {content.end()};

  if (!qi::parse(b, e, Grammar<decltype(b)>{}, result) || b != e) {
    throw std::runtime_error{"tree parser failed"};
  }

  return result;
}

auto Part1(Input const& input) -> std::size_t
{
  std::int64_t result {0};
  for (auto const& [i,pair] : boost::adaptors::index(input)) {
    if (pair[0] < pair[1]) {
      result += i + 1;
    }
  }
  return result;
}

auto Part2(Input const& input) -> std::size_t
{
  Tree two{2}, six{6}; // Indistinguishable from [[2]] and [[6]]
  std::size_t ix2{1}, ix6{2};
  for (auto const& pair : input) {
    for (auto const& tree : pair) {
      if (tree < six) {
        ix6++;
        if (tree < two) {
          ix2++;
        }
      }
    }
  }
  return ix2 * ix6;
}

} // namespace

TEST_SUITE("2022-13 examples") {
  TEST_CASE("simple ordering") {
    CHECK(Tree{1} < Tree{2});
    CHECK(!(Tree{2} < Tree{1}));
    CHECK(Tree{Trees{Tree{1}}} < Tree{Trees{Tree{2}}});
    CHECK(Tree{1} < Tree{Trees{Tree{2}}});
    CHECK(!(Tree{Trees{Tree{2}}} < Tree{1}));
    CHECK(Tree{Trees{}} < Tree{Trees{Tree{2}}});
    CHECK(Tree{Trees{}} < Tree{Trees{Tree{2}}});
    CHECK(Tree{1} < Tree{Trees{Tree{2},Tree{3}}});
    CHECK(!(Tree{Trees{Tree{2},Tree{3}}} < Tree{1}));
    CHECK(Tree{2} < Tree{Trees{Tree{2},Tree{3}}});
    CHECK(!(Tree{Trees{Tree{2},Tree{3}}} < Tree{2}));
  }
  TEST_CASE("documented example") {
    std::istringstream in {
      "[1,1,3,1,1]\n"
      "[1,1,5,1,1]\n"
      "\n"
      "[[1],[2,3,4]]\n"
      "[[1],4]\n"
      "\n"
      "[9]\n"
      "[[8,7,6]]\n"
      "\n"
      "[[4,4],4,4]\n"
      "[[4,4],4,4,4]\n"
      "\n"
      "[7,7,7,7]\n"
      "[7,7,7]\n"
      "\n"
      "[]\n"
      "[3]\n"
      "\n"
      "[[[]]]\n"
      "[[]]\n"
      "\n"
      "[1,[2,[3,[4,[5,6,7]]]],8,9]\n"
      "[1,[2,[3,[4,[5,6,0]]]],8,9]\n"
    };
    auto const input = Parse(in);
    CHECK(13 == Part1(input));
    CHECK(140 == Part2(input));
  }
}

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