(* Common util. Should probably put this in a module or something.. *)
let get_one_line file =
  try Some (input_line file) with
    End_of_file -> None

let get_lines file =
  let rec input lines =
    match get_one_line file with
      Some line -> input (line :: lines)
    | None -> List.rev lines
  in input []

let read_file file =
  let channel = open_in(file) in
  get_lines channel

(* --- *)

type operation = Nop | Acc | Jmp
type instruction = { op: operation; arg: int }
let op_of_string = function | "nop" -> Nop | "acc" -> Acc | "jmp" -> Jmp

let read_instruction = function |
  op :: arg :: [] -> { op = op_of_string op; arg = int_of_string arg }

let rec walk line_num acc lines_touched instructions =
  if List.length instructions < line_num + 1 (* at the end of the list *)
  then string_of_int acc (* success! *)
  else let { op; arg } = List.nth instructions line_num in (* match on the instruction *)
    if List.mem line_num lines_touched (* if we've already been on this line *)
      then raise (Failure (string_of_int acc)) (* raise a Failure of acc *)
      else match op with (* otherwise match on the op *)
        | Nop -> walk (line_num + 1) acc (line_num :: lines_touched) instructions
        | Acc -> walk (line_num + 1) (acc + arg) (line_num :: lines_touched) instructions
        | Jmp -> walk (line_num + arg) acc (line_num :: lines_touched) instructions

let swap_instruction ins line_num = (* swap the Nop and Jmp at the given line number *)
  List.mapi (fun i x ->
     if line_num = i then match x with (* if this is our line *)
       | { op = Jmp; arg } -> { op = Nop; arg }
       | { op = Nop; arg } -> { op = Jmp; arg }
       | a -> a (* if it is any other op, don't change it *)
     else x (* if it isn't our line number, don't change it *)
  ) ins

let rec walk_and_fix line_num instructions =
  let new_instructions = swap_instruction instructions line_num in (* swap nop and jmp *)
  try walk 0 0 [] new_instructions (* just return the acc this returns, if it doesn't raise *)
  with Failure acc -> walk_and_fix (line_num + 1) instructions (* Infinite loop detected. recurse! *)

let () =
  let ins = read_file "day8.txt"
    |> List.map @@ String.split_on_char ' '
    |> List.map read_instruction
  in

  try walk 0 0 [] ins
    |> Printf.printf "Complete! Wait this shouldn't have completed.. %s";
  with Failure acc -> (* catch infinite loop *)
    Printf.printf "Infinite loop detected. Last acc: %s\n" acc;

  walk_and_fix 0 ins
    |> Printf.printf "Infinite loop fixed! Last acc: %s\n";