#include <algorithm>
#include <cstdint>
#include <iostream>
#include <fstream>
#include <iterator>
#include <vector>
#include <map>
#include <sstream>

#include <doctest.h>

#include <aocpp/Startup.hpp>

namespace {

auto Parse(std::istream & in) {
  std::map<std::string, std::string> parents;
  std::string line;
  while (std::getline(in, line)) {
    auto it = line.find(')');
    if (it == std::string::npos) throw std::runtime_error{"bad input entry"};
    parents.insert({line.substr(it+1), line.substr(0, it)});
  }

  return parents;
}

auto Path(std::map<std::string, std::string> const& parents, std::string const& start)
{
  std::vector<std::string> path;
  std::string const* name = &start;
  while (*name != "COM") {
    name = &parents.at(*name);
    path.push_back(*name);
  }
  return path;
}

auto Part1(std::map<std::string, std::string> const& parents) {
  std::map<std::string, std::size_t> depths {{"COM", 0}};
  std::size_t part1 {0};

  for (auto & [k, _] : parents) {

    std::string const* cursor = &k;
    std::vector<std::string const*> todo;
  
    decltype(depths)::iterator it;
    while (depths.end() == (it = depths.find(*cursor))) {
      todo.push_back(cursor);
      cursor = &parents.at(*cursor);
    }
    auto n = it->second;

    for (; !todo.empty(); todo.pop_back()) {
      depths[*todo.back()] = ++n;
    }

    part1 += n;
  }
  
    
  return part1;
}

auto Part2(std::map<std::string, std::string> const& parents) {
  auto p1 = Path(parents, "SAN");
  auto p2 = Path(parents, "YOU");

  while (p1.back() == p2.back()) {
    p1.pop_back();
    p2.pop_back();
  }

  return p1.size() + p2.size();
}

} // namespace

TEST_SUITE("documented examples") {
  TEST_CASE("part 1") {
    std::istringstream in {"COM)B\nB)C\nC)D\nD)E\nE)F\nB)G\nG)H\nD)I\nE)J\nJ)K\nK)L\n"};
    REQUIRE(Part1(Parse(in)) == 42);
  }
  TEST_CASE("part 2") {
    std::istringstream in {"COM)B\nB)C\nC)D\nD)E\nE)F\nB)G\nG)H\nD)I\nE)J\nJ)K\nK)L\nK)YOU\nI)SAN\n"};
    REQUIRE(Part2(Parse(in)) == 4);
  }
}

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