open Printf;;

let rec list_of_lines in_file =
  try
    let line = input_line in_file in
    line :: list_of_lines(in_file)
  with End_of_file ->
    close_in in_file;
    []

let rec split_on_empty_line lines =
  match lines with
  | head :: tail ->
    if String.length head = 0 then
      ([], tail)
    else
      let rest_of_list, remaining_lines = split_on_empty_line tail in
      (head :: rest_of_list, remaining_lines)
  | _ -> ([], lines)

let encode_wire_id str =
  (Char.code str.[0]) 
  |> Int.logor (Int.shift_left (Char.code str.[1]) 8)
  |> Int.logor (Int.shift_left (Char.code str.[2]) 16)

let decode_wire_id id =
  [Char.chr (Int.logand id 0xff);
   Char.chr (Int.logand (Int.shift_right id 8) 0xff);
   Char.chr (Int.shift_right id 16)]
  |> List.to_seq |> String.of_seq

let parse_init str =
  (encode_wire_id (String.sub str 0 3), (String.sub str 5 1) |> int_of_string |> (fun x -> x = 1))

type gate_type =
  GATE_OR
  | GATE_AND
  | GATE_XOR

let gate_str gate =
  match gate with
  | GATE_OR -> "or"
  | GATE_AND -> "and"
  | GATE_XOR -> "xor"

let parse_gate str =
  let in1 = encode_wire_id (String.sub str 0 3) in
  let gate_str = (String.sub str 4 3) in
  let gate = if gate_str = "XOR" then GATE_XOR else if gate_str = "AND" then GATE_AND else GATE_OR in
  let in2_pos = if gate = GATE_OR then 7 else 8 in
  let in2 = encode_wire_id (String.sub str in2_pos 3) in
  let out_pos = in2_pos + 7 in
  let out = encode_wire_id (String.sub str out_pos 3) in
  (gate, in1, in2, out)

module IntPair =
struct
  type t = int * int
  let compare x0 x1 =
    Stdlib.compare x0 x1
end
module IntPairMap = Map.Make(IntPair)
module IntMap = Map.Make(Int)

let make_wire_state_map init_list =
  List.fold_left (fun acc (id, state) -> IntMap.add id state acc) IntMap.empty init_list

let make_wire_gate_mapping_fold acc_map gate_data =
  let gate, in1, in2, out = gate_data in
  let update_fun = (fun other_in prev_opt ->
    let new_entry = (gate, other_in, out) in
    match prev_opt with
    | Some prev -> Some (new_entry :: prev)
    | None -> Some [new_entry]
    ) in
  acc_map
  |> IntMap.update in1 (update_fun in2)
  |> IntMap.update in2 (update_fun in1)

let make_wire_gate_mapping gate_data_list =
  List.fold_left make_wire_gate_mapping_fold IntMap.empty gate_data_list

let rec simulate_circuit_inner wire_state_mapping wire_states wire_state_queue =
  match wire_state_queue with
  | (id, state) :: queue_tail -> begin
    let new_states = IntMap.add id state wire_states in
    match IntMap.find_opt id wire_state_mapping with
    | Some connection_list ->
        let new_queue = connection_list
        |> List.fold_left (fun queue (gate, other_in, out) ->
          match IntMap.find_opt other_in new_states with
          | Some other_state -> begin
              match gate with
              | GATE_OR -> (out, state || other_state) :: queue
              | GATE_AND -> (out, state && other_state) :: queue
              | GATE_XOR -> (out, state <> other_state) :: queue
          end
          | None -> queue
          ) queue_tail in
        simulate_circuit_inner wire_state_mapping new_states new_queue
    | None -> simulate_circuit_inner wire_state_mapping new_states queue_tail
  end
  | _ -> wire_states

let simulate_circuit wire_gate_mapping wire_states =
  simulate_circuit_inner wire_gate_mapping IntMap.empty (wire_states |> IntMap.to_seq |> List.of_seq)

let generate_prefixed_id prefix_c n =
  let n_str = string_of_int n in
  let prefix_str = String.make 1 prefix_c in
  if String.length n_str = 1 then
    String.cat (String.cat prefix_str "0") n_str
  else
    String.cat prefix_str n_str

let rec combine_z_output_inner output_states n out_num =
  match IntMap.find_opt (encode_wire_id @@ generate_prefixed_id 'z' n) output_states with
  | Some bit_bool ->
    let bit_num = if bit_bool then 1 else 0 in
    combine_z_output_inner output_states (n + 1) (Int.logor out_num (Int.shift_left bit_num n))
  | None -> out_num

let combine_z_output output_states =
  combine_z_output_inner output_states 0 0

let find_gate_out wire_gate_mapping in1 in2 gate_type =
  wire_gate_mapping
  |> IntMap.find in1
  |> List.find_opt (fun (other_gate, other_in, _) -> other_gate = gate_type && other_in = in2)
  |> Option.map (fun (_, _, out) -> out)

let search_gate_out_incomplete_in wire_gate_mapping in0 gate_type =
  wire_gate_mapping
  |> IntMap.find in0
  |> List.find_opt (fun (other_gate, _, _) -> other_gate = gate_type)
  |> Option.map (fun (_, other_in, out) -> printf "Found missing second input %s for input %s\n" (decode_wire_id other_in) (decode_wire_id in0); out)

let find_gate_out_opt wire_gate_mapping in1_opt in2_opt gate_type =
  match (in1_opt, in2_opt) with
  | Some in1, Some in2 ->
      let result1 = find_gate_out wire_gate_mapping in1 in2 gate_type in
      if Option.is_none result1 then begin
        printf "Did not find expected gate with inputs %s and %s\n" (decode_wire_id in1) (decode_wire_id in2);
        let new_result1 = search_gate_out_incomplete_in wire_gate_mapping in1 gate_type in
        let new_result2 = search_gate_out_incomplete_in wire_gate_mapping in2 gate_type in
        if Option.is_some new_result1 then
          new_result1
        else
          new_result2
      end else
        result1
  | Some in1, None -> search_gate_out_incomplete_in wire_gate_mapping in1 gate_type
  | None, Some in2 -> search_gate_out_incomplete_in wire_gate_mapping in2 gate_type
  | None, None -> printf "Trying to find a gate with no given inputs\n"; None

let expect_wires_equal level expected_opt actual_opt =
  match expected_opt with
  | None -> begin
    match actual_opt with
    | None -> printf "No idea what was expected and got uncertain at level %d\n" level
    | Some actual -> printf "No idea what was expected but got %s at level %d\n" (decode_wire_id actual) level
  end
  | Some expected ->
    match actual_opt with
    | None -> printf "Expected wire %s, got uncertain at level %d\n" (decode_wire_id expected) level
    | Some actual ->
        if expected <> actual then
          printf "Expected wire %s, got %s at level %d\n" (decode_wire_id expected) (decode_wire_id actual) level

let is_z_wire id =
  let str = decode_wire_id id in
  str.[0] = 'z' && int_of_string_opt (String.sub str 1 2) |> Option.is_some
let expect_wire_not_z id_opt =
  match id_opt with
  | Some id ->
      if is_z_wire id then begin
        printf "Expected non-z wire, got %s\n" (decode_wire_id id); None end
      else Some id
  | None -> None

let rec fix_adder_structure_inner wire_gate_mapping prev_carry_id level =
  let x_id = encode_wire_id @@ generate_prefixed_id 'x' level in
  let y_id = encode_wire_id @@ generate_prefixed_id 'y' level in
  let z_id = encode_wire_id @@ generate_prefixed_id 'z' level in
  if IntMap.find_opt x_id wire_gate_mapping |> Option.is_none then begin
    expect_wires_equal level (Some z_id) prev_carry_id;
  end else begin
    let partial_sum_id = find_gate_out wire_gate_mapping x_id y_id GATE_XOR |> expect_wire_not_z in
    let partial_carry_id = find_gate_out wire_gate_mapping x_id y_id GATE_AND |> expect_wire_not_z in
    let sum_id = find_gate_out_opt wire_gate_mapping partial_sum_id prev_carry_id GATE_XOR in
    let second_partial_carry_id = find_gate_out_opt wire_gate_mapping partial_sum_id prev_carry_id GATE_AND |> expect_wire_not_z  in
    let is_last_level = IntMap.find_opt (encode_wire_id @@ generate_prefixed_id 'x' (level + 1)) wire_gate_mapping |> Option.is_none in
    let carry_id = find_gate_out_opt wire_gate_mapping partial_carry_id second_partial_carry_id GATE_OR in
    if is_last_level then
      fix_adder_structure_inner wire_gate_mapping carry_id (level + 1)
    else
      let checked_carry_id = expect_wire_not_z carry_id in
      expect_wires_equal level (Some z_id) sum_id;
      fix_adder_structure_inner wire_gate_mapping checked_carry_id (level + 1)
  end

let fix_adder_structure wire_gate_mapping =
  let x0_id = encode_wire_id @@ generate_prefixed_id 'x' 0 in
  let y0_id = encode_wire_id @@ generate_prefixed_id 'y' 0 in
  let z0_id = encode_wire_id @@ generate_prefixed_id 'z' 0 in
  let sum_id = find_gate_out wire_gate_mapping x0_id y0_id GATE_XOR in
  let carry_id = find_gate_out wire_gate_mapping x0_id y0_id GATE_AND in
  expect_wires_equal 0 (Some z0_id) sum_id;
  fix_adder_structure_inner wire_gate_mapping carry_id 1

let () =
  Printexc.record_backtrace true;
  let f = open_in "input.txt" in
  let init_strs, connection_strs = list_of_lines f |> split_on_empty_line in
  let init_list = List.map parse_init init_strs in
  let connection_list = List.map parse_gate connection_strs in
  let wire_states_init = make_wire_state_map init_list in
  let wire_gate_mapping = make_wire_gate_mapping connection_list in
  let end_states = simulate_circuit wire_gate_mapping wire_states_init in
  let result = combine_z_output end_states in
  fix_adder_structure wire_gate_mapping;
  printf "Part 1: %d\n" result