(** [towels_of_strings lst] returns a pair containing a list of available towels
    and a list of patterns wanted. *)
let towels_of_strings = function
  | h :: "" :: t ->
      let re = Str.regexp "[, ]+" in
      let h = Str.split re h in
      (h, t)
  | _ -> failwith "towels_of_strings"

(** [towels_of_file fname] returns the list of towels and patterns from the file
    [fname]. *)
let towels_of_file fname = Aoc.strings_of_file fname |> towels_of_strings

(** Memoizing hash table shared between parts 1 and 2. *)
let memo = Hashtbl.create 1000

(** [memoize memo f value] returns the result of [f value]. The hashtable [memo]
    is used to cache results, so repeated calls with the same [value] will not
    call [f] again. *)
let memoize memo f value =
  match Hashtbl.find_opt memo value with
  | Some x -> x
  | None ->
      let x = f value in
      Hashtbl.add memo value x;
      x

(** [count_hashes memo towels pattern] counts the number of ways of matching
    [pattern] using [towels]. [memo] is a hashtable used for memoizing results.
*)
let rec count_matches memo towels pattern =
  let pattern_len = String.length pattern in
  let rec count_matched = function
    | [] -> 0
    | h :: t ->
        let towel_len = String.length h in
        if String.starts_with ~prefix:h pattern then
          memoize memo
            (count_matches memo towels)
            (String.sub pattern towel_len (pattern_len - towel_len))
          + count_matched t
        else count_matched t
  in
  if pattern_len = 0 then 1 else count_matched towels

let part1 (towels, patterns) =
  List.map (count_matches memo towels) patterns
  |> List.filter (( > ) 0)
  |> List.length

let part2 (towels, patterns) =
  List.map (memoize memo (count_matches memo towels)) patterns
  |> List.fold_left ( + ) 0

let _ =
  Aoc.main towels_of_file [ (string_of_int, part1); (string_of_int, part2) ]