#include <version>

#include <cstdint>
#include <iostream>
#include <set>
#include <sstream>
#include <optional>
#include <vector>

#if __has_include(<coroutine>)
#include <coroutine>
namespace co = std;
#elif __has_include(<experimental/coroutine>)
#include <experimental/coroutine>
namespace co = std::experimental;
#endif

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

#include <doctest.h>

#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
template <typename It>
class Grammar : public qi::grammar<It, Input()> {
  qi::rule<It, aocpp::Coord> coord;
  qi::rule<It, std::vector<aocpp::Coord>()> line;
  qi::rule<It, 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;
  }
};

/// Parse the complete input stream according to the rules defined in 'Grammar'.
/// Throws: std::runtime_error on failed parse
auto Parse(std::istream & in) -> Input
{
  auto result = Input{};
  auto const content = std::string{std::istreambuf_iterator{in}, {}};
  auto b = content.begin(); // updated on successful parse
  auto const e = content.end();

  if (!qi::parse(b, e, Grammar<decltype(b)>{}, result) || b != e) {
    throw std::runtime_error{"Bad input file: " + content};
  }

  return result;
}

auto ExpandLines(Input const& lines) -> std::set<aocpp::Coord>
{
  auto result = std::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 new 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;
}

struct CoordPromise;

struct CoordGenerator : co::coroutine_handle<CoordPromise> {
  using promise_type = struct CoordPromise;
  auto next() -> std::optional<aocpp::Coord>;
};

struct CoordPromise {
  std::optional<aocpp::Coord> output;
  std::exception_ptr e;

  auto get_return_object() -> CoordGenerator { return { CoordGenerator::from_promise(*this)}; }
  auto initial_suspend() noexcept -> co::suspend_always { return {}; }
  auto final_suspend() noexcept -> co::suspend_never { return {}; }
  auto return_void() { output.reset(); }
  auto yield_value(aocpp::Coord coord) -> co::suspend_always { output = coord; return {}; }
  auto unhandled_exception() -> void { e = std::current_exception(); }
};

auto CoordGenerator::next() -> std::optional<aocpp::Coord> {
  resume();
  auto & p = promise();
  if (p.e) {
    std::rethrow_exception(p.e);
  }
  return p.output;
}

auto SearchOrder(std::set<aocpp::Coord> & world, aocpp::Coord here) -> CoordGenerator
{
  co_yield aocpp::Down(here);
  co_yield aocpp::Left(aocpp::Down(here));
  co_yield aocpp::Right(aocpp::Down(here));
  world.insert(here);
}
  
auto Part1(std::int64_t bottom, std::set<aocpp::Coord> world) -> std::size_t
{
  auto const starting_size = world.size();

  auto path = std::vector<CoordGenerator>{};
  path.reserve(bottom - TOP.y);
  path.push_back(SearchOrder(world, TOP));

  while (!path.empty()) {
    if (auto next = path.back().next()) {
      if (!world.contains(*next)) {
        if (next->y == bottom) {
          break;
        }
        path.push_back(SearchOrder(world, *next));
      }
    } else {
      path.pop_back();
    }
  }

  return world.size() - starting_size;
}

auto Part2(std::int64_t bottom, std::set<aocpp::Coord> world) -> std::size_t
{
  auto const starting_size = world.size();
  
  auto path = std::vector<CoordGenerator>{};
  path.reserve(2 + bottom - TOP.y);
  path.push_back(SearchOrder(world, TOP));

  while (!path.empty()) {
    if (auto next = path.back().next()) {
      if (next->y - 2 != bottom && !world.contains(*next)) {
        path.push_back(SearchOrder(world, *next));
      }
    } else {
      path.pop_back();
    }
  }

  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 = Parse(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 = Parse(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;
}