#include <iostream>
#include <fstream>
#include <unordered_map>
#include <vector>
#include <variant>
#include <deque>

std::string getHeldBag(std::string &bag) {
    std::string colour = bag.substr(2);
    size_t bagPos = colour.find("bag.");
    size_t bagsPos = colour.find("bags");

    if (bag == "no other bags.") return "";

    if (bagsPos != std::variant_npos) colour = colour.substr(0, bagsPos + 3);
    else if (bagPos != std::variant_npos) colour = colour.substr(0, bagPos + 3);

    return colour;
}

std::vector<std::string> getHeldBags(std::string &bag) {
    std::vector<std::string> result = {};
    std::string colour = bag.substr(2);
    size_t bagsIndex = colour.find("bags");
    size_t bagIndex = colour.find("bag.");

    if (bag == "no other bags.") return result;
    uint32_t count = std::stoi(bag.substr(0, 1));

    if (bagsIndex != std::variant_npos) colour = colour.substr(0, bagsIndex + 3);
    else if (bagIndex != std::variant_npos) colour = colour.substr(0, bagIndex + 3);

    for (uint32_t i = 0; i < count; i++) result.push_back(colour);

    return result;
}

uint32_t countGoldBags(const std::string &start, const std::string &query,
                       std::unordered_map<std::string, std::vector<std::string>> &map) {
    uint32_t count = 0;
    std::deque<std::string> queue;

    if (start == query) return 0;
    queue.push_back(start);

    while (!queue.empty()) {
        std::string bag = queue.front();
        queue.pop_front();

        if (bag == query) count++;

        std::vector<std::string> heldBags = map.at(bag);
        for (std::string &subBag : heldBags) {
            queue.push_back(subBag);
        }
    }
    return count > 0;
}

uint32_t countBags(const std::string &query, std::unordered_map<std::string, std::vector<std::string>> &map) {
    uint32_t count = 0;
    std::deque<std::string> queue;

    queue.push_back(query);

    while (!queue.empty()) {
        std::string bag = queue.front();
        queue.pop_front();

        std::vector<std::string> heldBags = map.at(bag);
        count += heldBags.size();

        for (std::string &subBag : heldBags) queue.push_back(subBag);
    }

    return count;
}

int main() {
    std::ifstream input{"../input.txt"};
    std::unordered_map<std::string, std::vector<std::string>> shallowBagLookup = {};
    std::unordered_map<std::string, std::vector<std::string>> fullBagLookup = {};
    std::string rule;

    while (std::getline(input, rule)) {
        std::string colour = rule.substr(0, rule.find("s contain"));
        std::string rest = rule.substr(rule.find("contain ") + 8);
        std::vector<std::string> heldUnique = {};
        std::vector<std::string> held = {};

        size_t pos = 0;
        std::string tmp = rest;
        while ((pos = tmp.find(", ")) != std::variant_npos) {
            std::string str = tmp.substr(0, pos);

            std::string bag = getHeldBag(str);
            std::vector<std::string> bags = getHeldBags(str);
            if (!bag.empty()) heldUnique.push_back(bag);
            held.insert(held.end(), bags.begin(), bags.end());

            tmp = tmp.substr(pos + 2);
        }

        std::string bag = getHeldBag(tmp);
        std::vector<std::string> bags = getHeldBags(tmp);
        if (!bag.empty()) heldUnique.push_back(bag);
        held.insert(held.end(), bags.begin(), bags.end());

        shallowBagLookup.insert_or_assign(colour, heldUnique);
        fullBagLookup.insert_or_assign(colour, held);
    }

    const std::string query = "shiny gold bag";

    uint32_t goldBagCount = 0;
    for (auto &pair : shallowBagLookup) {
        goldBagCount += countGoldBags(pair.first, query, shallowBagLookup);
    }

    uint32_t bagsInGoldBagCount = countBags(query, fullBagLookup);

    std::cout << "[P1] Count: " << goldBagCount << std::endl;
    std::cout << "[P2] Count: " << bagsInGoldBagCount << std::endl;
    return 0;
}