#include <algorithm>
#include <cstdint>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <iterator>
#include <stdexcept>
#include <vector>
#include <set>
#include <string>
#include <tuple>
#include <bitset>
#include <deque>
#include <cctype>
#include <queue>
#include <map>

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

using namespace aocpp;

namespace {

Coord(* const directions[4])(Coord) = {Up, Down, Left, Right};

using Name = std::string;
using Portals = std::map<std::string, Coord>;
using Distances = std::map<Name,std::vector<std::pair<Name, std::int64_t>>>;

auto FindPortals(Grid const& grid) -> Portals {
  Portals portals;

  std::int64_t w = grid.rows[0].size();
  std::int64_t h = grid.rows.size();

  for (std::int64_t x = 1; x < w-1; x++) {
    for (std::int64_t y = 1; y < h-1; y++) {
      Coord const c = {x,y};
      auto const v = grid[c];
      if (std::isupper(v)) {
        char polarity = x==1 || x==w-2 || y==1 || y==h-2
                      ? '-' : '+';
        if (grid[Up(c)] == '.') {
          portals[{polarity, v, grid[Down(c)]}] = Up(c);
        } else if (grid[Down(c)] == '.') {
          portals[{polarity, grid[Up(c)], v}] = Down(c);
        } else if (grid[Left(c)] == '.') {
          portals[{polarity, v, grid[Right(c)]}] = Left(c);
        } else if (grid[Right(c)] == '.') {
          portals[{polarity, grid[Left(c)], v}] = Right(c);
        }
      }
    }
  }
  return portals;
}

auto FindDistancesFrom(
  Grid const& grid,
  std::map<Coord, Name> const& names,
  std::string const start_name,
  Coord const start
) {
  std::vector<std::pair<Name, std::int64_t>> result;

  std::deque<std::pair<std::int64_t, Coord>> todo {{0, start}};

  std::set<Coord> seen;

  for (; !todo.empty(); todo.pop_front()) {
    auto const [steps, here] = todo.front();

    // Don't visit the same coordinate twice
    if (!seen.insert(here).second) continue;

    auto const c = grid[here];
    if (c != '.') continue; // avoid walls

    // success, we've found a key, record the path
    if (auto it = names.find(here); it != names.end() && it->second != start_name) {
        result.emplace_back(it->second, steps);
        continue; // don't walk beyond the portal
    }

    // Visit all neighbors
    for (auto const fn : directions) {
      todo.emplace_back(steps+1, fn(here));
    }
  }

  return result;
}

auto OtherName(Name name) {
  name[0] = name[0] == '+' ? '-' : '+';
  return name;
}

auto FindDistances(Grid const& grid, Portals const& portals) {
  Distances distances;

  std::map<Coord, Name> names;
  for (auto const& [k,v] : portals) {
    names[v] = k;
  }

  for (auto const& [start_name, start_coord] : portals) {
    distances[start_name] = FindDistancesFrom(grid, names, start_name, start_coord);
  }

  return distances;
}

auto SolveMaze(Distances const& distances, bool const recursive) -> std::int64_t
{
  // Track current positions and current set of keys in easy to compare form
  using Visited = std::pair<std::int64_t, Name>;
  std::set<Visited> seen;

  // Priority queue returning lowest path cost states first.
  using PqElt = std::tuple<std::int64_t, std::int64_t, Name>;
  using PqCmp = decltype([](PqElt const& x, PqElt const& y) {
                return std::get<0>(x) > std::get<0>(y); });
  std::priority_queue<PqElt, std::vector<PqElt>, PqCmp> todo;

  todo.emplace(0, 0, "-AA");

  while(!todo.empty()) {
    auto const& [steps, depth, name] = todo.top();
    todo.pop();
    if (name == "-ZZ") { return steps; }
    if (seen.emplace(depth, name).second) {
      if (auto const it = distances.find(name); it != distances.end()) {
        for (auto const& [next, cost] : it->second) {
          if (next == "-ZZ") {
            if (depth == 0) todo.emplace(steps + cost, depth, "-ZZ");
          } else {
            auto const depth_ = depth + (recursive ? (next[0]=='+' ? 1 : -1) : 0);
            if (depth_ >= 0) todo.emplace(steps + cost + 1, depth_, OtherName(next));
          }
        }
      }
    }
  }
  throw std::runtime_error{"no solution"};
}

} // namespace

auto main(int argc, char** argv) -> int {
  auto map = Grid::Parse(Startup(argc, argv));
  auto portals = FindPortals(map);
  auto distances = FindDistances(map, portals);

  std::cout << "Part 1: " << SolveMaze(distances, false) << std::endl;
  std::cout << "Part 2: " << SolveMaze(distances, true) << std::endl;
}