#include <cstdint>
#include <iostream>
#include <set>
#include <sstream>
#include <tuple>
#include <vector>
#include <variant>

#include <doctest.h>

#include <aocpp/Startup.hpp>

namespace {

struct Tree {
  std::variant<std::int64_t, std::vector<Tree>> var;

  auto operator<(Tree const& rhs) const -> bool {
    if (var.index() == 0) {
      if (rhs.var.index() == 0) {
        return std::get<0>(var) < std::get<0>(rhs.var);
      } else {
        return Tree{std::vector{*this}} < rhs;
      }
    } else {
      if (rhs.var.index() == 0) {
        return *this < Tree{std::vector{rhs}};
      } else {
        return std::get<1>(var) < std::get<1>(rhs.var);
      }
    }
  }
};

using Input = std::vector<Tree>;

auto ParseTree(std::string::iterator & begin, std::string::iterator end) -> Tree
{
  if (*begin == '[') {
    begin++;
    if (*begin == ']') {
      begin++;
      return Tree{std::vector<Tree>{}};
    }
    std::vector<Tree> subtrees;
    top: subtrees.emplace_back(ParseTree(begin, end));
    switch (*begin++) {
      case ',': goto top;
      case ']': return Tree{std::move(subtrees)};
      default: assert(false);
    }
  } else if (isdigit(*begin)) {
    auto number_end = std::find_if_not(begin, end, isdigit);
    std::string number_str {begin, number_end};
    begin = number_end;
    std::int64_t number = std::stoll(number_str);
    return Tree{number};
  } else {
    assert(false);
  }
}

auto Parse(std::istream & in) -> Input
{
  Input result;
  std::string line;
  while (std::getline(in, line)) {
    auto cursor = line.begin();
    result.push_back(ParseTree(cursor, line.end()));
    std::getline(in, line);
    cursor = line.begin();
    result.push_back(ParseTree(cursor, line.end()));
    std::getline(in, line); // skip blank
  }

  return result;
}

auto Part1(Input const& input) -> std::int64_t
{
  std::int64_t result {0};
  for (std::size_t i = 0; i < input.size(); i += 2) {
    if (input[i] < input[i+1]) {
      result += i/2 + 1;
    }
  }
  return result;
}

auto Part2(Input input) -> std::int64_t
{
  Tree two {std::vector{Tree{std::vector{Tree{2}}}}};
  Tree six {std::vector{Tree{std::vector{Tree{6}}}}};
  auto ix2 = std::count_if(input.begin(), input.end(), [&](auto x) { return x < two; });
  auto ix6 = std::count_if(input.begin(), input.end(), [&](auto x) { return x < six; });
  return (ix2+1) * (ix6+2);
}

} // namespace

TEST_SUITE("2022-13 examples") {
  TEST_CASE("simple ordering") {
    CHECK(Tree{1} < Tree{2});
    CHECK(!(Tree{2} < Tree{1}));
    CHECK(Tree{std::vector{Tree{1}}} < Tree{std::vector{Tree{2}}});
    CHECK(Tree{1} < Tree{std::vector{Tree{2}}});
    CHECK(!(Tree{std::vector{Tree{2}}} < Tree{1}));
    CHECK(Tree{std::vector<Tree>{}} < Tree{std::vector{Tree{2}}});
    CHECK(Tree{std::vector<Tree>{}} < Tree{std::vector{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 input = Parse(in);
    CHECK(13 == Part1(input));
    CHECK(140 == Part2(input));
  }
}

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