type vm = {
  a : int;
  b : int;
  c : int;
  ip : int;
  code : int array;
  out : int list;
}
(** Virtual Machine definition

    Consists of three general-purpose registers: A, B, C, an instruction pointer
    (ip), and array of code, and a list of output. The output list is in reverse
    order of output. *)

(** Generate a VM from a list of strings *)
let vm_of_strings lst =
  let re_val = Str.regexp {|Register [ABC]: \([0-9]+\)|} in
  let re_prog = Str.regexp {|Program: \([0-9,]+\)|} in
  let get_val = function
    | [] -> failwith "vm_of_strings.get_val"
    | h :: t ->
        let _ = Str.search_forward re_val h 0 in
        (int_of_string (Str.matched_group 1 h), t)
  in
  let skip_line = function
    | "" :: t -> t
    | _ -> failwith "vm_of_strings.skip_line"
  in
  let get_code = function
    | [] -> failwith "vm_of_strings.get_prog"
    | h :: t ->
        let _ = Str.search_forward re_prog h 0 in
        let nums =
          Str.matched_group 1 h |> Aoc.ints_of_string ~sep:"," |> Array.of_list
        in
        (nums, t)
  in
  let a, lst = get_val lst in
  let b, lst = get_val lst in
  let c, lst = get_val lst in
  let lst = skip_line lst in
  let code, lst = get_code lst in
  let ip = 0 in
  let out = [] in
  assert (List.is_empty lst);
  { a; b; c; ip; code; out }

(** Generate a VM from a file *)
let vm_of_file fname = Aoc.strings_of_file fname |> vm_of_strings

(** [is_halted vm] returns true if the VM is halted. *)
let is_halted vm = vm.ip >= Array.length vm.code || vm.ip < 0

(** [get_literal vm] returns the Literal value from the current IP + 1 in [vm].
*)
let get_literal vm = vm.code.(vm.ip + 1)

(** [print_combo vm] prints the combo operand at the current IP + 1 in [vm]. *)
let[@warning "-32"] print_combo vm =
  match vm.code.(vm.ip + 1) with
  | 0 | 1 | 2 | 3 -> print_int vm.code.(vm.ip + 1)
  | 4 -> Printf.printf "A (=%d)" vm.a
  | 5 -> Printf.printf "B (=%d)" vm.a
  | 6 -> Printf.printf "C (=%d)" vm.a
  | 7 -> failwith "print_combo reserved"
  | _ -> failwith "print_combo not 3-bit"

(** [get_combo vm] returns the value of interpreting the combo operant at IP + 1
    in [vm]. *)
let get_combo vm =
  match vm.code.(vm.ip + 1) with
  | 0 | 1 | 2 | 3 -> vm.code.(vm.ip + 1)
  | 4 -> vm.a
  | 5 -> vm.b
  | 6 -> vm.c
  | 7 -> failwith "get_combo reserved"
  | _ -> failwith "get_combo not 3-bit"

(** [print_insn vm] prints the current instruction at VM. *)
let[@warning "-32"] print_insn vm =
  Printf.printf "%d: " vm.ip;
  if is_halted vm then print_endline "Halted"
  else Printf.printf "%d " vm.code.(vm.ip);
  (match vm.code.(vm.ip) with
  | 0 ->
      print_string "adv ";
      print_combo vm
  | 1 -> Printf.printf "bxl %d" (get_literal vm)
  | 2 ->
      print_string "bst ";
      print_combo vm
  | 3 -> Printf.printf "jnz %d" (get_literal vm)
  | 4 -> print_string "bxc"
  | 5 ->
      print_string "out ";
      print_combo vm
  | 6 ->
      print_string "bdv ";
      print_combo vm
  | 7 ->
      print_string "cdv ";
      print_combo vm
  | _ -> failwith "print_insn");
  Printf.printf " A=%d B=%d C=%d\n" vm.a vm.b vm.c

(** [execute_insn vm] executes the current instruction in [vm] returns an
    updated VM. *)
let execute_insn vm =
  if is_halted vm then vm
  else
    (*print_insn vm;*)
    match vm.code.(vm.ip) with
    | 0 -> { vm with a = vm.a / (1 lsl get_combo vm); ip = vm.ip + 2 }
    | 1 -> { vm with b = vm.b lxor get_literal vm; ip = vm.ip + 2 }
    | 2 -> { vm with b = get_combo vm land 7; ip = vm.ip + 2 }
    | 3 ->
        if vm.a <> 0 then { vm with ip = get_literal vm }
        else { vm with ip = vm.ip + 2 }
    | 4 -> { vm with b = vm.b lxor vm.c; ip = vm.ip + 2 }
    | 5 -> { vm with out = (get_combo vm land 7) :: vm.out; ip = vm.ip + 2 }
    | 6 -> { vm with b = vm.a / (1 lsl get_combo vm); ip = vm.ip + 2 }
    | 7 -> { vm with c = vm.a / (1 lsl get_combo vm); ip = vm.ip + 2 }
    | _ -> failwith "execute_insn"

(** [execute_until_halted vm] executes a VM until it runs out of steam. *)
let rec execute_until_halted vm =
  match is_halted vm with
  | true -> vm
  | false -> execute_until_halted (execute_insn vm)

(** [string_of_ouput vm] gives the output of [vm]. *)
let string_of_output vm =
  List.rev vm.out |> List.map string_of_int |> String.concat ","

(** [scan_digit acc ip vm] updates the acc for A so that the output of running
    VM agrees in the last [ip] digits with the input program. *)
let scan_digit acc ip vm =
  let rec impl v =
    let acc = acc + (v lsl (3 * ip)) in
    let vm = { vm with a = acc } in
    let vm = execute_until_halted vm in
    let out = List.rev vm.out in
    if List.length out <= ip then impl (v + 1)
    else if
      List.of_seq (Seq.drop ip (List.to_seq out))
      = List.of_seq (Seq.drop ip (Array.to_seq vm.code))
    then acc
    else impl (v + 1)
  in
  impl 0

(** [scan_all vm] generates the input A which means the output of executing [vm]
    is the same as the input program. *)
let scan_all vm =
  let rec impl acc ip =
    if ip < 0 || ip >= Array.length vm.code then { vm with a = acc }
    else impl (scan_digit acc ip vm) (ip - 1)
  in
  impl 0 (Array.length vm.code - 1)

(** [string_of_a vm] returns the A register of [vm]. *)
let string_of_a vm = string_of_int vm.a

let _ =
  Aoc.main vm_of_file
    [ (string_of_output, execute_until_halted); (string_of_a, scan_all) ]