type grid = { grid : char array; width : int }
(** Grid type. Not our normal type as we want to use it in a mutable way. *)

(** [grid_map fn strs] Returns a grid built up from the list of strings [strs].
    [fn] takes a single string of all map characters and returns a string with
    the appropriate replacements done. *)
let grid_make posmap_fn strs =
  let grid =
    List.fold_left String.cat "" strs
    |> posmap_fn |> String.to_seq |> Array.of_seq
  in
  let height = List.length strs in
  let width = Array.length grid / height in
  { grid; width }

(** [grid_print grid] prints [grid] to standard output. *)
let[@warning "-32"] grid_print grid =
  Array.iteri
    (fun i c ->
      print_char c;
      if i mod grid.width = grid.width - 1 then print_newline ())
    grid.grid

(** [find_robot_idx grid] returns the index of the robot in grid. *)
let find_robot_idx grid =
  match Array.find_index (( = ) '@') grid.grid with
  | None -> failwith "find_robot_idx"
  | Some x -> x

(** [instrs_of_file fname] reads from the file [fname] and returns a
    [(grid, instrs)] pair. [grid] is a list of strings describing the initial
    grid state. [instrs] is a list of characters giving movement instructions.
*)
let instrs_of_file fname =
  let strs = Aoc.strings_of_file fname in
  let rec impl acc = function
    | [] | "" :: [] -> failwith "instrs_of_file.impl"
    | "" :: t -> (List.rev acc, t)
    | h :: t -> impl (h :: acc) t
  in
  let grid, moves = impl [] strs in
  (grid, List.fold_left String.cat "" moves |> String.to_seq |> List.of_seq)

(** [move_x grid i di] moves the object in [grid] at [i] to [i + di] if
    possible, returning the new location of the object (either [i] or [i + di]).
    This assumes a horizontal movement of one step. *)
let rec move_x grid i di =
  assert (di == -1 || di == 1);
  assert (i >= 0 && i <= Array.length grid.grid);
  match grid.grid.(i + di) with
  | '#' -> i
  | 'O' | '[' | ']' ->
      if move_x grid (i + di) di = i + di then i else move_x grid i di
  | '.' ->
      grid.grid.(i + di) <- grid.grid.(i);
      grid.grid.(i) <- '.';
      i + di
  | _ -> failwith "move_x"

(** [can_move_y grid i di] returns true if and only if it is possible to move
    whatever is at [i] to [i + di] in [grid]. If [i + di] is full it recursively
    checks to see if that can be moved as well. [di] must be a vertical
    movement. *)
let rec can_move_y grid i di =
  assert (i >= 0 && i < Array.length grid.grid);
  assert (di = grid.width || di = -grid.width);
  match grid.grid.(i + di) with
  | '#' -> false
  | '.' -> true
  | 'O' -> can_move_y grid (i + di) di
  | '[' -> can_move_y grid (i + di) di && can_move_y grid (i + di + 1) di
  | ']' -> can_move_y grid (i + di - 1) di && can_move_y grid (i + di) di
  | _ -> failwith "can_move_y"

(** [move_x grid i di] moves the object in [grid] at [i] to [i + di] if
    possible, returning the new location of the object (either [i] or [i + di]).
    This assumes a vertical movement of one step. *)
let move_y grid i di =
  assert (i >= 0 && i < Array.length grid.grid);
  assert (di = grid.width || di = -grid.width);
  let rec do_move i =
    match grid.grid.(i + di) with
    | '#' -> failwith "move_y.do_move #"
    | '.' ->
        grid.grid.(i + di) <- grid.grid.(i);
        grid.grid.(i) <- '.'
    | 'O' ->
        do_move (i + di);
        do_move i
    | '[' ->
        do_move (i + di);
        do_move (i + di + 1);
        do_move i
    | ']' ->
        do_move (i + di - 1);
        do_move (i + di);
        do_move i
    | _ -> failwith "move_y.do_move"
  in
  if can_move_y grid i di then (
    do_move i;
    i + di)
  else i

(** [process_move grid robot dir] attempts to move robot at idx [robot] in
    [grid] in direction. Returns location of robot after attempt. *)
let process_move grid robot dir =
  match dir with
  | '^' -> move_y grid robot (-grid.width)
  | 'v' -> move_y grid robot grid.width
  | '<' -> move_x grid robot ~-1
  | '>' -> move_x grid robot 1
  | _ -> failwith "process_move"

(** [process_moves grid robot lst] Attempts each move in [lst] in turn given a
    [grid] and initial starting position of the robot [robot]. Returns the
    [grid]. *)
let rec process_moves grid robot = function
  | [] -> ()
  | h :: t -> process_moves grid (process_move grid robot h) t

(** [calc_score grid] returns the score for [grid]. *)
let calc_score grid =
  Array.mapi
    (fun idx c ->
      if c = 'O' || c = '[' then
        (idx mod grid.width) + (100 * (idx / grid.width))
      else 0)
    grid.grid
  |> Array.fold_left ( + ) 0

(** [expand_grid str] given an input [str] returns the output grid string for
    part 2. *)
let expand_grid str =
  let b = Buffer.create (String.length str * 2) in
  let add_c = function
    | '#' -> Buffer.add_string b "##"
    | 'O' -> Buffer.add_string b "[]"
    | '.' -> Buffer.add_string b ".."
    | '@' -> Buffer.add_string b "@."
    | _ -> failwith "expand_grid"
  in
  String.iter add_c str;
  Buffer.contents b

(** [part posmap_fn (grid, moves)] executes the [moves] in [grid] having first
    applied [posmap_fn] to grid. *)
let part posmap_fn (grid, moves) =
  let grid = grid_make posmap_fn grid in
  let robot = find_robot_idx grid in
  process_moves grid robot moves;
  (*grid_print grid;*)
  calc_score grid

let _ =
  Aoc.main instrs_of_file
    [ (string_of_int, part Fun.id); (string_of_int, part expand_grid) ]