type vm = {
  a : int;
  b : int;
  c : int;
  ip : int;
  code : int array;
  out : int list;
}

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 }

let vm_of_file fname = Aoc.strings_of_file fname |> vm_of_strings
let get_literal vm = vm.code.(vm.ip + 1)

let 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"

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"

let is_halted vm = vm.ip >= Array.length vm.code || vm.ip < 0

let 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

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")

let rec execute_until_halted vm =
  match is_halted vm with
  | true -> vm
  | false -> execute_until_halted (execute_insn vm)

let string_of_output vm =
  List.rev vm.out |> List.map string_of_int |> String.concat ","

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