#load "str.cma" ;;
#use "lib/utils.ml" ;;

type crane_model = CrateMover9000 | CrateMover9001 ;;

type move = {cnt: int; from: int; into: int} ;;

let parse_stacks txt =
  let lines = String.split_on_char '\n' txt in
  let charw = 4 in
  let nstacks = ((String.length (List.nth lines 0)) + 1) / charw in
  let lines = List.rev lines
  |> List.tl in
  let parse_line line =
    List.init nstacks (fun i -> i + 1)
    |> List.map (fun i -> (String.get line ((i - 1) * charw + 1)))
  in
  List.map parse_line lines
  |> transpose
  |> List.map (List.filter (fun c -> c <> ' '))
in

let parse_moves txt =
  let lines = String.split_on_char '\n' txt
  |> List.filter (fun l -> String.trim l <> "")
  in
  let parse_line line =
    let tokens = String.split_on_char ' ' line in
    { cnt = int_of_string (List.nth tokens 1);
      from = int_of_string (List.nth tokens 3) - 1;
      into = int_of_string (List.nth tokens 5) - 1;
    }
  in
  List.map parse_line lines
in

let format_stacks stacks =
  stacks |> List.mapi (fun i stack ->
    let boxes = List.map (fun c -> Printf.sprintf "[%c]" c) stack in
    let tokens = (string_of_int (i + 1)) :: boxes in
    String.concat " " tokens
  ) |> String.concat "\n"
in

let debug_move stacks move =
  Printf.printf "move %d from %d to %d\n" move.cnt (move.from+1) (move.into+1);
  print_endline (format_stacks stacks)
in

let solution_of stacks =
  let last_char stack =
    let len = List.length stack in
    let c = List.nth stack (len - 1) in
    Printf.sprintf "%c" c
  in
  String.concat "" (List.map last_char stacks)
in

let move_boxes stacks moves crane =
  moves |> List.fold_left (fun stacks move ->
    let new_stack = List.nth stacks (move.from)
    |> List.rev
    |> List.to_seq
    in
    let buf = Seq.take move.cnt new_stack
    |> List.of_seq
    in
    let buf = match crane with
    | CrateMover9000 -> buf
    | CrateMover9001 -> List.rev buf
    in
    let new_stack = Seq.drop move.cnt new_stack
    |> List.of_seq
    |> List.rev
    in
    let stacks = stacks |> List.mapi (fun i stack ->
      if i = move.from then new_stack
      else if i = move.into then stack @ buf
      else stack
    ) in
    (* let _ = debug_move stacks move in *)
    stacks
  ) stacks
in

(* Silence warnings *)
let _ = debug_move in

let contents = read_file "inputs/5.txt" in
let parts = Str.split (Str.regexp "\n\n") contents in
let stacks = parse_stacks (List.nth parts 0) in
let moves = parse_moves (List.nth parts 1) in
(* let _ = print_endline (format_stacks stacks) in *)
let res1 = move_boxes stacks moves CrateMover9000 in
let res2 = move_boxes stacks moves CrateMover9001 in
Printf.printf "Solution 1: %s\n" (solution_of res1);
Printf.printf "Solution 2: %s\n" (solution_of res2)