package main

import "core:math"
import "core:slice"
import "core:strconv"
import "core:strings"
import "core:bytes"
import "core:fmt"
import "core:os"

import "../util"

Color :: string

InnerBag :: struct {
    color: Color, 
    amount: i64,
}
Bags :: map[Color]Maybe([]InnerBag)

bag_contains_target :: proc(bags: ^Bags, current_color, target_color: Color) -> bool {
    maybe_inner_bags := bags[current_color]
    inner_bags, ok := maybe_inner_bags.?
    if !ok do return false

    for inner_bag in inner_bags {
        if inner_bag.color == target_color {
            return true
        } else if bag_contains_target(bags, inner_bag.color, target_color) {
            return true 
        }
    }
    
    return false
}

part_1 :: proc(lines: []string) {
    result: = 0
    bags := make(Bags)

    for line in lines {
        parts := strings.split(line[:len(line)-1], " contain ")
        outer_bag := strings.split(parts[0], " bag")[0]
        if strings.contains(line, "no other bags") {
            bags[outer_bag] = nil
            continue
        }
        
        inner_bags := strings.split(parts[1], ", ")
        bags[outer_bag] = slice.mapper(inner_bags, proc(bag: string) -> InnerBag {
            inner_bag := strings.split(bag, " bag")[0]
            bag_details := strings.split_n(inner_bag, " ", 2)

            color := bag_details[1]
            amount, ok := strconv.parse_i64(bag_details[0]); assert(ok)
            return InnerBag{
                color = color,
                amount = amount,
            }
        })
    }

    TARGET_COLOR :: "shiny gold"
    for color, _ in bags {
        if bag_contains_target(&bags, color, TARGET_COLOR) {
            result += 1
        }
    }

    fmt.printfln("Part 1: %d", result)
}

bag_sum_color :: proc(bags: ^Bags, sum_color: Color) -> i64 {
    maybe_inner_bags := bags[sum_color]
    inner_bags, ok := maybe_inner_bags.?
    if !ok do return 0

    amount: i64 = 0
    for inner_bag in inner_bags {
        amount += inner_bag.amount * (1+bag_sum_color(bags, inner_bag.color))
    }

    return amount
}

part_2 :: proc(lines: []string) {
    bags := make(Bags)

    for line in lines {
        parts := strings.split(line[:len(line)-1], " contain ")
        outer_bag := strings.split(parts[0], " bag")[0]
        if strings.contains(line, "no other bags") {
            bags[outer_bag] = nil
            continue
        }
        
        inner_bags := strings.split(parts[1], ", ")
        bags[outer_bag] = slice.mapper(inner_bags, proc(bag: string) -> InnerBag {
            inner_bag := strings.split(bag, " bag")[0]
            bag_details := strings.split_n(inner_bag, " ", 2)

            color := bag_details[1]
            amount, ok := strconv.parse_i64(bag_details[0]); assert(ok)
            return InnerBag{
                color = color,
                amount = amount,
            }
        })
    }

    TARGET_COLOR :: "shiny gold"
    result := bag_sum_color(&bags, TARGET_COLOR)

    fmt.printfln("Part 2: %d", result)
}

main :: proc() {
    context.allocator = context.temp_allocator
    defer free_all(context.temp_allocator)

    INPUT :: #load("input.txt", string)
    lines := strings.split(INPUT, "\n")
    lines = lines[:len(lines)-1]

    part_1(lines)
    part_2(lines)
}