open Angstrom
open Base
open Core

type initseq = string list

module Parse = struct
  let chomp s = String.chop_suffix_if_exists ~suffix:"\n" s

  let parse_initseq s : initseq =
    let s = chomp s in
    String.split s ~on:','
end

module Part1 = struct
  let hash s =
    let l = String.to_list s in
    let step st c = Int.(17 * (st + Char.to_int c) % 256) in
    List.fold l ~init:0 ~f:step

  let hash_all steps = List.map steps ~f:hash
  let hash_initseq s = Parse.chomp s |> Parse.parse_initseq |> hash_all
  let hash_sum s = hash_initseq s |> List.fold ~init:0 ~f:( + )
end

module Part2 = struct
  (* A lens has a label and a focal length *)
  type lens = { label : string; f : int } [@@deriving show]

  let equal_lens_label { label; _ } { label = label2; _ } =
    String.(label = label2)

  (* An op is insertion or removal of a lens *)
  type op = Remove of string | Insert of lens [@@deriving show]

  (* A hash map is an array of buckets *)
  type bucket = lens list [@@deriving show]
  type hm = bucket array

  (* Make a hash map *)
  let make_hm () : hm = Array.create ~len:256 []

  (* Get a bucket by lens label. *)
  let get_hm hm (lbl : string) =
    let ix = Part1.hash lbl in
    Array.get hm ix

  (* Set a bucket by lens label. *)
  let set_hm hm lbl list =
    let ix = Part1.hash lbl in
    Array.set hm ix list

  (* Remove a lens by label. *)
  let remove_hm hm lbl =
    let l = get_hm hm lbl in
    let l' = List.filter l ~f:(fun l2 -> not String.(lbl = l2.label)) in
    set_hm hm lbl l'

  (* Update a list of lenses: either replace an existing lens or add a new one *)
  let update_list lbl lens =
    let replace (ok, a) e =
      if (not ok) && equal_lens_label e lens then (true, lens :: a)
      else (ok, e :: a)
    in
    let ok, l' = List.fold ~f:replace ~init:(false, []) lbl in
    let l_new = List.rev (if ok then l' else lens :: l') in
    l_new

  (* Insert or replace a lens into a hash map *)
  let insert_hm hm lens =
    let l = get_hm hm lens.label in
    let l' = update_list l lens in
    set_hm hm lens.label l'

  (* Parsing for Part 2 *)
  module Parse = struct
    let parse_digit =
      satisfy Char.is_digit >>| fun c -> String.of_char c |> Int.of_string

    let parse_op =
      let open Angstrom in
      let label = take_while1 Char.is_alpha in
      let remove = label <* char '-' >>| fun l -> Remove l in
      let insert =
        label <* char '=' >>= fun l ->
        parse_digit >>= fun f -> return (Insert { label = l; f })
      in
      remove <|> insert

    let parse_ops s =
      let open Angstrom in
      let parse s =
        parse_string ~consume:All parse_op s |> Result.ok_or_failwith
      in
      let parts = Parse.parse_initseq s in
      let ops = List.map parts ~f:parse in
      ops
  end

  (* Process a single op *)
  let process_one hm = function
    | Remove label -> remove_hm hm label
    | Insert l -> insert_hm hm l

  (* Process a list of ops *)
  let process_all hm ops = List.iter ops ~f:(process_one hm)

  (* Assign a score to a bucket *)
  let score_bucket ix l =
    let a = ix + 1 and f i acc l = acc + ((i + 1) * l.f) in
    let bc = List.foldi l ~init:0 ~f in
    let score = a * bc in
    score

  (* Assign a score to the whole hash map *)
  let score hm =
    let f ix acc l = acc + score_bucket ix l in
    Array.foldi hm ~init:0 ~f

  (* Process an input string *)
  let process s =
    let ops = Parse.parse_ops s in
    let hm = make_hm () in
    process_all hm ops;
    let score = score hm in
    score
end

let () =
  let inp = In_channel.(input_all stdin) in
  let s = Part1.hash_sum inp in
  printf "Part 1: %d\n" s;
  let score = Part2.process inp in
  printf "Part 2: %d\n" score