let load_file fname =
  match In_channel.with_open_text fname In_channel.input_line with
  | Some x -> x
  | None -> failwith "load_file"

(** [disk_size disk_str] returns the size of the disk represented by the string
    [disk_str]. See AoC 2024 day 9 for description of string format. *)
let disk_size disk_str =
  let rec impl acc disk_str =
    match String.length disk_str with
    | 0 -> acc
    | len ->
        let h = int_of_string (String.sub disk_str 0 1) in
        let t = String.sub disk_str 1 (len - 1) in
        impl (acc + h) t
  in
  impl 0 disk_str

(** [disk_init disk_str] returns a disk which represents the description given
    by [disk_str]. The returned disk is a mutable array with elements being [-1]
    for free space and [id >= 0] for file with the given ID. *)
let disk_init disk_str =
  let size = disk_size disk_str in
  let disk = Array.make size (-1) in
  let rec add_id offset id = function
    | 0 -> offset
    | x ->
        disk.(offset) <- id;
        add_id (offset + 1) id (x - 1)
  in
  let rec impl offset id disk_str =
    match String.length disk_str with
    | 0 -> ()
    | 1 ->
        let _ = add_id offset id (int_of_string (String.sub disk_str 0 1)) in
        ()
    | str_len ->
        let len = int_of_string (String.sub disk_str 0 1) in
        impl
          (add_id offset id len + int_of_string (String.sub disk_str 1 1))
          (id + 1)
          (String.sub disk_str 2 (str_len - 2))
  in
  impl 0 0 disk_str;
  disk

(** [disk_defrag disk] defrags [disk] given 2024/09/part 1 rules. Updates disk
    in place. *)
let disk_defrag disk =
  let rec impl front back =
    if front >= back then ()
    else if disk.(front) <> -1 then impl (front + 1) back
    else if disk.(back) = -1 then impl front (back - 1)
    else (
      disk.(front) <- disk.(back);
      disk.(back) <- -1;
      impl (front + 1) (back - 1))
  in
  impl 0 (Array.length disk - 1);
  disk

(** [find_largest_id disk] returns the largest ID on a fragmented disk [disk].
*)
let find_largest_id disk =
  let rec impl pos =
    if pos = -1 then failwith "find_largest_id.impl"
    else if disk.(pos) = -1 then impl (pos - 1)
    else disk.(pos)
  in
  impl (Array.length disk - 1)

(** [find_id disk id search_pos] returns a pair [(start, length)] of the file
    [id] on [disk]. We only look backwards from [search_pos]. *)
let rec find_id disk id search_pos =
  let rec count_length len pos =
    if pos < 0 then (pos, len)
    else if disk.(pos) <> id then (pos + 1, len)
    else count_length (len + 1) (pos - 1)
  in
  if search_pos < 0 then failwith "find_id"
  else if disk.(search_pos) <> id then find_id disk id (search_pos - 1)
  else count_length 0 search_pos

(** [find_space disk len] finds [len] elements of free space on [disk] starting
    from 0. Returns [None] if no free-space found, or [Some pos] giving the
    position of the start of the found space. *)
let find_space disk len =
  let rec free_length acc pos =
    if pos >= Array.length disk then acc
    else if disk.(pos) <> -1 then acc
    else free_length (acc + 1) (pos + 1)
  in
  let rec impl pos =
    if pos >= Array.length disk then None
    else if disk.(pos) <> -1 then impl (pos + 1)
    else if free_length 0 pos < len then impl (pos + 1)
    else Some pos
  in
  impl 0

(** [move_file disk id src dest] moves the file [id] on [disk] from [src] to
    [dest].*)
let rec move_file disk id src dest =
  if src >= Array.length disk then ()
  else if dest >= Array.length disk then ()
  else if disk.(src) <> id then ()
  else if disk.(dest) <> -1 then failwith "move_block"
  else (
    disk.(dest) <- disk.(src);
    disk.(src) <- -1;
    move_file disk id (src + 1) (dest + 1);
    ())

(** [file_defrag disk id search_pos] Locates the file [id] on [disk] and defrags
    it if possible. We start searching down from [search_pos]. *)
let file_defrag disk id search_pos =
  let pos, len = find_id disk id search_pos in
  match find_space disk len with
  | None -> pos - 1
  | Some x ->
      if pos >= x then move_file disk id pos x;
      pos - 1

(** Defrag a whole disk according to 2024/09/part 2 rules. *)
let disk_defrag_whole disk =
  let max_id = find_largest_id disk in
  let rec impl id pos =
    if id = 0 then disk else impl (id - 1) (file_defrag disk id pos)
  in
  impl max_id (Array.length disk - 1)

(** [disk_checksum disk] Calculates the checksum for [disk]. *)
let disk_checksum disk =
  let rec impl acc idx =
    if idx = Array.length disk then acc
    else if disk.(idx) = -1 then impl acc (idx + 1)
    else impl (acc + (disk.(idx) * idx)) (idx + 1)
  in
  impl 0 0

(** [part algo str] defrags the disk represented by [str] using algorithm
    [algo]. *)
let part algo str = disk_init str |> algo |> disk_checksum

let _ =
  Aoc.main load_file
    [
      (string_of_int, part disk_defrag); (string_of_int, part disk_defrag_whole);
    ]