open Stdio
open Base

let input () = In_channel.stdin

(* Part 1 *)
module Part1 = struct
  let calibration_value_of_string s =
    let digits = String.filter s ~f:Char.is_digit in
    let first, last = (String.prefix digits 1, String.suffix digits 1) in
    Int.of_string (String.concat [ first; last ])

  let process ch =
    let f lines line = calibration_value_of_string line :: lines in
    List.rev @@ In_channel.fold_lines ch ~init:[] ~f

  let _part1 () =
    let r = process (input ()) in
    Out_channel.printf "%d\n" (List.fold ~init:0 ~f:( + ) r)
end

(* Part 2 *)

module Part2 = struct
  let folder ix acc x : (int * int) option =
    match (ix, acc, x) with
    | ix, None, d ->
        let i = Int.of_string_opt (String.of_char d) in
        Option.map i ~f:(fun i -> (ix, i))
    | _, Some d, _ -> Some d

  (* Return the (index, value) of the first digit in the string *)
  let find_first_digit s = String.foldi ~init:None ~f:folder s

  (* Return the (index, value) of the last digit in the string *)
  let find_last_digit s : (int * int) option =
    let%map.Option ix, v = find_first_digit (String.rev s) in
    (String.length s - ix - 1, v)

  let digits =
    [
      (1, "one");
      (2, "two");
      (3, "three");
      (4, "four");
      (5, "five");
      (6, "six");
      (7, "seven");
      (8, "eight");
      (9, "nine");
    ]

  (* generate a counting sequence *)
  let count_seq (from : int) (upto : int) : int Sequence.t =
    let f c =
      if Int.equal c upto then None
      else Some (c, if upto > from then c + 1 else c - 1)
    in
    Sequence.unfold ~init:from ~f

  (* try to find number string `number` with value `value` in string `s`, from the beginning if `fwd`. *)
  let find_number ~fwd s value number : int option * int =
    let ixs =
      if fwd then count_seq 0 (String.length s)
      else count_seq (String.length s - 1) 0
    in
    let check pos = String.is_substring_at s ~pos ~substring:number in
    (Sequence.find ixs ~f:check, value)

  (* find any number word in the string, from the beginning if `fwd` or from the end otherwise. *)
  let find_any_number ~fwd s : (int * int) option =
    let f (value, number) = find_number ~fwd s value number in
    let found = List.map digits ~f in
    let compare (ix1, _) (ix2, _) =
      match (ix1, ix2) with
      | Some a, Some b -> Int.compare a b * if fwd then 1 else -1
      | Some _, None -> -1
      | None, Some _ -> 1
      | None, None -> 0
    in
    (* sort by index *)
    let sorted = List.sort found ~compare in
    match List.hd_exn sorted with Some ix, d -> Some (ix, d) | None, _ -> None

  (* Process line `line`, and return numeric result. *)
  let handle line : int =
    let maybe_first_dig = find_first_digit line in
    let maybe_first_wrd = find_any_number ~fwd:true line in
    let maybe_last_dig = find_last_digit line in
    let maybe_last_wrd = find_any_number ~fwd:false line in
    let match_num ~fwd (dig : (int * int) option) (wrd : (int * int) option) =
      match (dig, wrd) with
      | Some (ix1, d1), Some (ix2, d2) ->
          if (fwd && ix1 < ix2) || ((not fwd) && ix1 > ix2) then d1 else d2
      | Some (_, d), None | None, Some (_, d) -> d
      | None, None -> assert false
    in
    let first_dig = match_num ~fwd:true maybe_first_dig maybe_first_wrd in
    let last_dig = match_num ~fwd:false maybe_last_dig maybe_last_wrd in
    (10 * first_dig) + last_dig

  let process ch =
    let f lines line = handle line :: lines in
    List.rev @@ In_channel.fold_lines ch ~init:[] ~f

  let part2 () =
    let r = process In_channel.stdin in
    Out_channel.printf "%d\n" (List.fold ~init:0 ~f:( + ) r)
end

let () = Part2.part2 ()