type op = And | Or | Xor
type gate = { in1 : string; in2 : string; op : op; out : string }

module StringMap = Map.Make (String)

let get_wire_value str =
  let re = Str.regexp {|\(.+\): \([01]\)|} in
  let _ = Str.search_forward re str 0 in
  let v = if Str.matched_group 2 str = "0" then 0 else 1 in
  (Str.matched_group 1 str, v)

let get_gate_op = function
  | "AND" -> And
  | "OR" -> Or
  | "XOR" -> Xor
  | _ -> failwith "get_gate_op"

let[@warning "-32"] string_of_op = function
  | And -> "AND"
  | Or -> "OR"
  | Xor -> "XOR"

let get_gate_config str =
  let re = Str.regexp {|\(.+\) \(AND\|OR\|XOR\) \(.+\) -> \(.+\)|} in
  let _ = Str.search_forward re str 0 in
  let in1 = Str.matched_group 1 str in
  let op = get_gate_op (Str.matched_group 2 str) in
  let in2 = Str.matched_group 3 str in
  let out = Str.matched_group 4 str in
  { in1; in2; op; out }

let initial_wires_of_strings =
  let rec impl acc = function
    | "" :: t -> (acc, t)
    | h :: t ->
        let wire, v = get_wire_value h in
        impl (StringMap.add wire v acc) t
    | _ -> failwith "initial_wires_of_strings"
  in
  impl StringMap.empty

let gates_from_strings =
  let rec impl acc = function
    | [] -> acc
    | h :: t -> impl (get_gate_config h :: acc) t
  in
  impl []

let config_of_file fname =
  let lst = Aoc.strings_of_file fname in
  let wires, lst = initial_wires_of_strings lst in
  let gates = gates_from_strings lst in
  (wires, gates)

let process_gate wires gate =
  match
    ( gate.op,
      StringMap.find_opt gate.in1 wires,
      StringMap.find_opt gate.in2 wires )
  with
  | And, Some a, Some b -> if a = 1 && b = 1 then Some 1 else Some 0
  | Or, Some a, Some b -> if a = 1 || b = 1 then Some 1 else Some 0
  | Xor, Some a, Some b -> if a <> b then Some 1 else Some 0
  | _, _, _ -> None

let process_gates wires =
  let rec impl wires acc = function
    | [] -> (wires, acc)
    | h :: t -> begin
        match process_gate wires h with
        | None -> impl wires (h :: acc) t
        | Some x ->
            let wires = StringMap.add h.out x wires in
            impl wires acc t
      end
  in
  impl wires []

let rec repeat_to_end wires gates =
  let old_len = List.length gates in
  let wires, gates = process_gates wires gates in
  if gates = [] then Some wires
  else if old_len = List.length gates then begin
    Printf.printf "Loop detected: %d\n" (List.length gates);
    None
  end
  else begin
    repeat_to_end wires gates
  end

let calc_value = Fun.flip (List.fold_right (fun x acc -> x + (2 * acc))) 0

let k_wires wires x =
  StringMap.filter (fun k _ -> k.[0] = x) wires
  |> StringMap.bindings |> List.map snd |> calc_value

let wires_set wires x v' =
  let set_v k v =
    if k.[0] = x then
      let idx = int_of_string (String.sub k 1 (String.length k - 1)) in
      (v' lsr idx) land 1
    else v
  in
  StringMap.mapi set_v wires

let part1 (wires, gates) =
  match repeat_to_end wires gates with
  | None -> failwith "part1"
  | Some wires -> k_wires wires 'z'

let part2 (wires, gates) =
  let run_test x y =
    let wires = wires_set wires 'x' x in
    let wires = wires_set wires 'y' y in
    Printf.printf "%d + %d = " (k_wires wires 'x') (k_wires wires 'y');
    match repeat_to_end wires gates with
    | None -> print_endline "(infinite loop)"
    | Some wires ->
        let z = k_wires wires 'z' in
        print_int z;
        if z <> x + y then print_string " (wrong answer)";
        print_newline ()
  in
  let tst n =
    Printf.printf "Test for n = %d\n" n;
    run_test (1 lsl n) 0;
    run_test 0 (1 lsl n);
    run_test (1 lsl n) (1 lsl n)
  in
  Seq.ints 0 |> Seq.take 45 |> Seq.iter tst;
  0

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