#include <algorithm>
#include <deque>
#include <fstream>
#include <iostream>
#include <sstream>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using GraphMap = std::unordered_map<std::string, std::vector<std::string>>;

GraphMap getMap(const std::string &file_name) {
    GraphMap map{};

    std::ifstream file(file_name);
    std::string str;
    while (std::getline(file, str)) {
        auto deliminer = str.find_first_of('-');
        auto start = str.substr(0, deliminer);
        auto end = str.substr(deliminer + 1);
        map[start].push_back(end);
        map[end].push_back(start);
    }
    return map;
}

uint64_t getPaths(const GraphMap &map, const std::string &start,
                  std::vector<std::string> &visited,
#ifdef DEBUG
                  std::vector<std::string> &path,
#endif
                  bool can_twice = false, bool has_twice = false) {
    uint64_t paths = 0;
    for (const auto &dest : map.at(start)) {
        if (dest == "start") {
            continue;
        }
        bool is_twice = false;
        if (std::find(visited.begin(), visited.end(), dest) != visited.end()) {
            if (!can_twice || (can_twice && has_twice)) {
                continue;
            }
            has_twice = true;
            is_twice = true;
        }
        if (dest == "end") {
            paths++;
#ifdef DEBUG
            for (auto &thingy : path) {
                std::cout << thingy << "->";
            }
            std::cout << "end" << std::endl;
#endif
            continue;
        }
        if (dest[0] >= 'a' && dest[0] <= 'z' && !is_twice) {
            visited.push_back(dest);
        }
#ifdef DEBUG
        path.push_back(dest);
        paths += getPaths(map, dest, visited, path, can_twice, has_twice);
        path.pop_back();
#endif
        paths += getPaths(map, dest, visited, can_twice, has_twice);
        if (dest[0] >= 'a' && dest[0] <= 'z' && !is_twice) {
            visited.pop_back();
        }
        if (is_twice) {
            has_twice = false;
        }
    }
    return paths;
}

uint64_t part1(const GraphMap &map) {
    std::vector<std::string> visited{ "start" };
#ifdef DEBUG
    std::vector<std::string> path{ "start" };
    return getPaths(map, "start", visited, path);
#endif
    return getPaths(map, "start", visited);
}

uint64_t part2(const GraphMap &map) {
    std::vector<std::string> visited{ "start" };
#ifdef DEBUG
    std::vector<std::string> path{ "start" };
    return getPaths(map, "start", visited, path, true);
#endif
    return getPaths(map, "start", visited, true);
}

int main(int argc, char **argv) {
    if (argc < 2) {
        std::cerr << "You must provide input file!" << std::endl;
        return 1;
    }
    auto map = getMap(argv[1]);
    std::cout << "There are \033[91;1m" << part1(map)
              << "\033[0m paths in the cave system." << std::endl;
    std::cout << "There are \033[91;1m" << part2(map)
              << "\033[0m paths in the cave system with new rules."
              << std::endl;
    return 0;
}