#load "str.cma";;

(* 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

let explode s = (* convert a string to a list of chars *)
  let rec exp i l =
   if i < 0 then l else exp (i - 1) (s.[i] :: l) in
  exp (String.length s - 1) []

let implode l = String.of_seq (List.to_seq l)

let (<<) f g x = f(g(x))
let (>>) f g x = g(f(x))

let rec find_idx nth x = function (* find nth index of x in a list *)
  | [] -> raise (Failure "Not Found")
  | h :: t ->
      if x = h
      then if nth = 0 then 0 else 1 + find_idx (nth-1) x t
      else 1 + find_idx nth x t

let rec take n l = (* take the first n elements *)
  match n with
  | 0 -> []
  | n ->
      match l with
      | [] -> []
      | x :: xs -> x :: (take (n-1) xs)

let rec drop n l = (* drop the first n elemnts *)
  match n with
  | 0 -> l
  | n -> drop (n-1) @@ List.tl l

let string_of_char = String.make 1

(* --- *)

type rule = Char of char | List of int list | Pair of rule * rule
type rule' = Char' of char | List' of rule' list | Pair' of rule' * rule' (* deref'd *)

let rec string_of_rule' = function
  | Char' x -> string_of_char x
  | List' l -> List.fold_left (fun acc x -> acc ^ (string_of_rule' x)) "" l
  | Pair' (a, b) -> "\\(" ^ (string_of_rule' a) ^ "\\|" ^ (string_of_rule' b) ^ "\\)"

let rec rule_of_string str =
  let str' = String.trim str in
  if String.contains str' '|'
  then let [l; r] = String.split_on_char '|' str' in
    Pair (rule_of_string l, rule_of_string r)
  else if String.contains str' '"' then
    Char (String.get str' 1)
  else
    List (List.map int_of_string (String.split_on_char ' ' str'))

let parse_rule_line line =
  let [key; rule_string] = String.split_on_char ':' line in
  let rule = rule_of_string rule_string in
  (int_of_string key, rule)

let read_chunks lines =
  let blank_idx = find_idx 0 "" lines in
  let rules = take blank_idx lines |> List.map parse_rule_line in
  let messages = drop blank_idx lines in
  (rules, messages)

let rec deref rules depth rule =
  if depth > 20 (* CHEATER MODE *)
  then List' []
  else
    match rule with
    | Char x -> Char' x
    | Pair (a, b) -> Pair' (deref rules (depth+1) a, deref rules (depth+1) b)
    | List l -> List' (List.map (fun key ->
        let rule = List.assoc key rules in
        deref rules depth rule) l)

let () =
  let (rules, messages) = read_file "day19.txt" |> read_chunks in

  let rule0 = deref rules 0 (List.assoc 0 rules) in
  print_string (string_of_rule' rule0);

  let rule0_str = string_of_rule' rule0 in
  let r = Str.regexp ("^" ^ rule0_str ^ "$") in

  let good = List.filter (fun m -> Str.string_match r m 0) messages in
  Printf.printf "\nnum good: %i" (List.length good);