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

#include <doctest.h>

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

namespace {

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

auto operator<(Tree const& lhs, Tree const& rhs) -> bool;

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

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, std::vector<Tree> const& y) -> bool {
  switch (y.size()) {
    case 0: return false;
    case 1: return x < y[0];
    default: return !(y[0] < x);
  }
}

auto operator<(std::vector<Tree> 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);
}

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: throw std::runtime_error{"bad tree list"};
    }
  } 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 {
    throw std::runtime_error{"bad tree start"};
  }
}

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 const& 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 const& x) { return x < two; });
  auto ix6 = std::count_if(input.begin(), input.end(), [&](auto const& 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(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;
}