#include <version>

#include <cstdint>
#include <iostream>
#include <sstream>
#include <stack>
#include <stdexcept>
#include <tuple>
#include <unordered_set>
#include <vector>

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

#include <doctest.h>

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

namespace {

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

using Input = std::vector<std::vector<aocpp::Coord>>;

const aocpp::Coord TOP {500, 0};

// Input file grammar
class Grammar : public qi::grammar<std::string::const_iterator, Input()> {
  qi::rule<iterator_type, aocpp::Coord> coord;
  qi::rule<iterator_type, std::vector<aocpp::Coord>()> line;
  qi::rule<iterator_type, Input()> input;

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

    coord = qi::long_long [ phx::bind(&aocpp::Coord::x, _val) = _1 ] >>
            "," >>
            qi::long_long [ phx::bind(&aocpp::Coord::y, _val) = _1 ];
    line = coord % " -> " >> "\n";
    input = *line;
  }
};

auto ExpandLines(Input const& lines) -> std::unordered_set<aocpp::Coord>
{
  auto result = std::unordered_set<aocpp::Coord>{};
  for (auto const& line : lines) {
    for (auto const i : boost::irange(std::size_t{1}, line.size())) {
      auto const& a = line[i-1];
      auto const& b = line[i];
     
      if (a.x == b.x) {
        for (auto const y : boost::irange(std::min(a.y, b.y), 1+std::max(a.y, b.y))) {
          result.insert({a.x, y});
        }
      } else if (a.y == b.y) {
        for (auto const x : boost::irange(std::min(a.x, b.x), 1+std::max(a.x, b.x))) {
          result.insert({x, a.y});
        }
      } else {
        throw std::runtime_error{"Bad line segment"};
      }
    }
  }
  return result;
}

auto FindBottom(Input const& input) -> std::int64_t
{
  auto result = std::int64_t{TOP.y};
  for (auto const& line : input) {
    for (auto const& coord : line) {
      result = std::max(result, coord.y);
    }
  }
  return result;
}

enum class Action { Visit, Mark };
  
auto Part1(std::int64_t bottom, std::unordered_set<aocpp::Coord> world) -> std::size_t
{
  auto const starting_size = world.size();
  auto work = std::stack<std::pair<Action, aocpp::Coord>>{};
  work.emplace(Action::Visit, TOP);

  while (!work.empty()) {
    auto & [action, here] = work.top();
    switch (action) {
      case Action::Visit:
        if (!world.contains(here)) {
          
          // Particle is falling off the world, stop simulation
          if (here.y == bottom) {
            return world.size() - starting_size;
          }
          
          action = Action::Mark;
          auto const down = aocpp::Down(here);
          // action and here are invalidated after this line
          work.emplace(Action::Visit, aocpp::Right(down));
          work.emplace(Action::Visit, aocpp::Left(down));
          work.emplace(Action::Visit, down);
        } else {
          work.pop();
        }
        break;
      case Action::Mark:
        work.pop();
        world.insert(here);
        break;
    }
  }

  throw std::runtime_error{"No solution for part 1"};
}

auto Part2(std::int64_t bottom, std::unordered_set<aocpp::Coord> world) -> std::size_t
{
  auto const starting_size = world.size();
  auto work = std::stack<std::pair<Action, aocpp::Coord>>{};
  work.emplace(Action::Visit, TOP);

  while (!work.empty()) {
    auto & [action, here] = work.top();
    
    switch (action) {
      case Action::Visit:
        if (here.y - 2 != bottom && !world.contains(here)) {
          action = Action::Mark;
          auto const down = aocpp::Down(here);
          work.emplace(Action::Visit, aocpp::Right(down));
          work.emplace(Action::Visit, aocpp::Left(down));
          work.emplace(Action::Visit, down);
        } else {
          work.pop();
        }
        break;
      case Action::Mark:
        world.insert(here);
        work.pop();
        break;
    }
  }

  return world.size() - starting_size;
}

} // namespace

TEST_SUITE("2022-14 examples") {
  TEST_CASE("documented example") {
    auto in = std::istringstream{
      "498,4 -> 498,6 -> 496,6\n"
      "503,4 -> 502,4 -> 502,9 -> 494,9\n"
    };
    auto const input = aocpp::ParseGrammar(Grammar{}, in);
    auto const bottom = FindBottom(input);
    auto world = ExpandLines(input);
    CHECK(24 == Part1(bottom, world));
    CHECK(93 == Part2(bottom, std::move(world)));
  }
}

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