module IntSet = Set.Make (Int)

(** [add_rule a b m] adds the rule that [a] must appear before [b] to the rule
    map [m]. Returns the updated map [m] *)
let add_rule a b m =
  match Hashtbl.find_opt m a with
  | None -> Hashtbl.add m a (IntSet.singleton b)
  | Some s -> Hashtbl.replace m a (IntSet.add b s)

(** [find_rule a b m] returns [true] if the rule map [m] says that [a] should
    appear before [b]. *)
let find_rule a b m =
  match Hashtbl.find_opt m a with
  | Some s -> ( match IntSet.find_opt b s with Some _ -> true | None -> false)
  | None -> false

(** [compare m a b] is a total ordering on pages in the rule map [m]. Returns
    [-1] if [a] should appear before [b], [0] if [a = b], and [1] if [b] should
    appear before [a]. *)
let compare m a b =
  if a = b then 0
  else if find_rule a b m then -1
  else if find_rule b a m then 1
  else failwith "compare"

(** [sort m pages] Sorts [pages] into the order required by the rule map [m]. *)
let sort m = List.sort (compare m)

(** [is_page_order_valid m pages] returns [true] iff the page ordering given in
    [pages] is valid under the rule map [m]. *)
let rec is_page_order_valid m pages =
  let rec impl h = function
    | h' :: t -> find_rule h h' m && impl h t
    | [] -> true
  in
  match pages with h :: t -> impl h t && is_page_order_valid m t | [] -> true

(** [parse_rules lst] parses the rules in the list [lst] stopping when
    encountering an empty line. Returns a pair [(rule_map, tail)]. [tail] starts
    the line after the empty line. *)
let parse_rules =
  let m = Hashtbl.create 17 in
  let rec impl = function
    | "" :: t -> (m, t)
    | [] -> failwith "parse_rules.impl"
    | h :: t -> (
        match Aoc.ints_of_string ~sep:"|" h with
        | [ a; b ] ->
            add_rule a b m;
            impl t
        | _ -> failwith "parse_rules.impl")
  in
  impl

(** [parse_page_orders lst] parses a list of page orders. *)
let parse_page_orders =
  let rec impl acc = function
    | [] -> acc
    | h :: t -> impl (Aoc.ints_of_string ~sep:"," h :: acc) t
  in
  impl []

(** Read a rule map and list of pages from the file [fname]. *)
let read_file fname =
  let lst = Aoc.strings_of_file fname in
  let rule_map, t = parse_rules lst in
  let page_orders = parse_page_orders t in
  (rule_map, page_orders)

let middle_elt lst = List.nth lst (List.length lst / 2)

let part1 (rule_map, page_orders) =
  List.filter (is_page_order_valid rule_map) page_orders
  |> List.map middle_elt |> List.fold_left ( + ) 0

let part2 (rule_map, page_orders) =
  List.filter (fun lst -> not (is_page_order_valid rule_map lst)) page_orders
  |> List.map (sort rule_map)
  |> List.map middle_elt |> List.fold_left ( + ) 0

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