open Utils

let initial_state =
    let ht = Hashtbl.create 12345                                     in
    let _ = Hashtbl.add ht 1 ["S"; "P"; "H"; "V"; "F"; "G"]           in
    let _ = Hashtbl.add ht 2 ["M"; "Z"; "D"; "V"; "B"; "F"; "J"; "G"] in
    let _ = Hashtbl.add ht 3 ["N"; "J"; "L"; "M"; "G"]                in
    let _ = Hashtbl.add ht 4 ["P"; "W"; "D"; "V"; "Z"; "G"; "N"]      in
    let _ = Hashtbl.add ht 5 ["B"; "C"; "R"; "V"]                     in
    let _ = Hashtbl.add ht 6 ["Z"; "L"; "W"; "P"; "M"; "S"; "R"; "V"] in
    let _ = Hashtbl.add ht 7 ["P"; "H"; "T"]                          in
    let _ = Hashtbl.add ht 8 ["V"; "Z"; "H"; "C"; "N"; "S"; "R"; "Q"] in
    let _ = Hashtbl.add ht 9 ["J"; "Q"; "V"; "P"; "G"; "L"; "F"]      in
    ht

exception InvalidMoveLine

let parse_move line =
    match String.split_on_char ' ' line with
        | "move" :: count :: "from" :: st_from :: "to" :: st_to :: [] ->
                (int_of_string count, int_of_string st_from, int_of_string st_to)
        | _ -> raise InvalidMoveLine

let move stacks f (count, idx_from, idx_to) =
    let from_stack = Hashtbl.find stacks idx_from in
    let to_stack = Hashtbl.find stacks idx_to in
    let rec aux count res_from append_to =
        if count = 0 then (res_from, append_to)
        else match res_from with
        | h :: t -> aux (count - 1) t (h :: append_to)
        | [] -> (res_from, append_to) in
    let (res_from, res_to) = aux count from_stack [] in
    let _ = Hashtbl.replace stacks idx_from res_from in
    let _ = Hashtbl.replace stacks idx_to (List.append (f res_to) to_stack) in
    stacks ;;

let parsed_moves = read_file "bin/d5/input.txt"
    |> List.filter_map
        (fun line ->
            try Some(parse_move line)
            with InvalidMoveLine -> let _ = print_endline ("invalid line " ^ line) in None) ;;

let solve f input =
    let state = Hashtbl.copy initial_state in
    let _ = input
        |> List.map (move state f) in
    let res = Hashtbl.length state
        |> range
        |> List.fold_left (fun acc x -> acc ^ (List.hd (Hashtbl.find state x))) "" in
    print_endline ("\n" ^ res) ;;

solve id parsed_moves ;;
solve List.rev parsed_moves ;;