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

(*
  let disk_print disk =
    let prnt c = if c = -1 then print_char '.' else print_int c in
    Array.iter prnt disk;
    print_newline();
*)
let disk_size disk_str =
  let rec impl acc disk_str =
    let len = String.length disk_str in
    match len with
    | 0 -> acc
    | len ->
        let h = int_of_string (String.sub disk_str 0 1) in
        flush_all ();
        let t = String.sub disk_str 1 (len - 1) in
        flush_all ();
        impl (acc + h) t
  in
  impl 0 disk_str

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

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

let calc_checksum disk =
  let rec impl acc idx =
    if disk.(idx) = -1 then acc else impl (acc + (disk.(idx) * idx)) (idx + 1)
  in
  impl 0 0

let part1 str = disk_init str |> disk_defrag |> calc_checksum
let _ = Aoc.main load_file [ (string_of_int, part1) ]