(* Common util. Should probably put this in a module or something.. *)
let get_one_line file =
  try Some (input_line file) with
    End_of_file -> None

let get_lines file =
  let rec input lines =
    match get_one_line file with
      Some line -> input (line :: lines)
    | None -> List.rev lines
  in input []

let read_file file =
  let channel = open_in(file) in
  get_lines channel

(* --- *)

type bag = { color: string; contents: bag list; count: int }

let read_contents = function (* takes the last part of the string and turns it into a "list" of "bags" *)
  | "no other bags." -> []
  | str -> String.split_on_char ',' str
    |> List.map String.trim
    |> List.map @@ String.split_on_char ' '
    |> List.map (function | count :: modifier :: color :: _rest -> (* pattern match on bag description strings *)
        { color = modifier ^ " " ^ color; contents = []; count = int_of_string count })

let read_bags line =
  let parts = String.split_on_char ' ' line in
  let modifier :: color :: _bags :: _contain :: rest = parts in
  let contents = read_contents (String.concat " " rest) in (* re-concat the rest of theh string back together *)
  { color = modifier ^ " " ^ color; contents; count = 1 }

let bag_contains_color color bag =
  List.exists (fun x -> String.equal x.color color) bag.contents

let get_bag_from_color color all_bags = (* get the top-level bag from it's color string *)
  List.find (fun x -> String.equal x.color color) all_bags

let rec has_bag color all_bags bag =
  if bag_contains_color color bag (* if this bag directly contains that color *)
  then true (* return true *)
  else List.exists (fun x -> (* otherwise, check it's contents *)
    let bag' = get_bag_from_color x.color all_bags in (* for each contained bag, find it in the top level, and call has_bag on it *)
    has_bag color all_bags bag')
  bag.contents

let rec num_children all_bags color = (* find the number of child bags for a given color *)
  let bag = get_bag_from_color color all_bags in
  match bag.contents with
  | [] -> 0
  | contents -> List.fold_left (fun acc x -> (* fold over each child bag *)
      acc + (x.count * (1 + num_children all_bags x.color))) 0 contents

let () =
  let all_bags = read_file "day7.txt" |> List.map read_bags in

  let contains_shiny_gold = List.filter (has_bag "shiny gold" all_bags) all_bags in
  let sum_of_bags_containing_shiny_gold = List.length contains_shiny_gold in
  Printf.printf "# of bags containing a \"shiny gold\" bag: %i\n" sum_of_bags_containing_shiny_gold;

  let child_count = num_children all_bags "shiny gold" in
  Printf.printf "# of bags *inside* a \"shiny gold\" bag: %i\n" child_count;