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

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
  string_of_int (List.length rings)

let rec search_candidate max_set connections current candidates =
  match StringSet.choose_opt candidates with
  | None ->
      if List.length current > List.length max_set then current else max_set
  | Some h ->
      let current' = h :: current in
      let candidates' =
        StringSet.inter candidates (StringMap.find h connections)
      in
      let max_set = search_candidate max_set connections current' candidates' in
      search_candidate max_set connections current
        (StringSet.remove h candidates)

let find_max_set connections =
  let rec impl max_set = function
    | [] -> max_set
    | (k, v) :: t -> impl (search_candidate max_set connections [ k ] v) t
  in
  impl [] (StringMap.to_list connections)

let part2 connections =
  let max_set = find_max_set connections in
  List.sort compare max_set |> String.concat ","

let _ = Aoc.main load_file [ (Fun.id, part1); (Fun.id, part2) ]