module StringMap = Map.Make (String)
module StringSet = Set.Make (String)

let[@warning "-32"] print_set set =
  StringSet.iter
    (fun x ->
      print_string x;
      print_char ' ')
    set;
  print_newline ()

let[@warning "-32"] print_sets = List.iter print_set

let[@warning "-32"] print_map =
  StringMap.iter (fun k v ->
      Printf.printf "%s: " k;
      print_set v)

let add_connection map (a, b) =
  let update s = function
    | None -> Some (StringSet.add s StringSet.empty)
    | Some set -> Some (StringSet.add s set)
  in
  let map = StringMap.update a (update b) map in
  let map = StringMap.update b (update a) map in
  map

let make_pairs = function
  | [ a; b ] -> (a, b)
  | _ -> raise (invalid_arg "make_pairs")

let load_file fname =
  Aoc.strings_of_file fname
  |> List.map (String.split_on_char '-')
  |> List.map make_pairs
  |> List.fold_left add_connection StringMap.empty

let is_three_ring connections a _ candidate =
  StringSet.mem candidate (StringMap.find a connections)

let rec find_third_member acc connections visited a b candidates =
  match candidates with
  | [] -> acc
  | h :: t ->
      if StringSet.mem h visited then
        find_third_member acc connections visited a b t
      else if not (is_three_ring connections a b h) then
        find_third_member acc connections visited a b t
      else
        let acc = StringSet.of_list [ a; b; h ] :: acc in
        find_third_member acc connections visited a b t

let rec find_second_member acc connections visited a candidates =
  match candidates with
  | [] -> acc
  | h :: t ->
      if StringSet.mem h visited then
        find_second_member acc connections visited a t
      else
        let visited = StringSet.add h visited in
        let acc =
          find_third_member acc connections visited a h
            (StringSet.to_list (StringMap.find h connections))
        in
        find_second_member acc connections visited a t

let rec find_rings acc visited connections computers =
  match computers with
  | [] -> acc
  | h :: t ->
      if StringSet.mem h visited then find_rings acc visited connections t
      else
        let visited = StringSet.add h visited in
        let acc =
          find_second_member acc connections visited h
            (StringSet.to_list (StringMap.find h connections))
        in
        find_rings acc visited connections t

let starts_with_t set = StringSet.exists (fun x -> x.[0] = 't') set

let part1 connections =
  let computers = StringMap.to_list connections |> List.map fst in
  let rings =
    find_rings [] StringSet.empty connections computers
    |> List.filter starts_with_t
  in
  print_sets rings;
  List.length rings

let _ = Aoc.main load_file [ (string_of_int, part1) ]