open Angstrom
open Base
open Core

(* A spring is either working, broken, or unknown. *)
type spring = Working | Broken | Unknown [@@deriving show, sexp]

(* A record is a row of springs, and a list of groups of broken springs. *)
type record = { row : spring array; damaged_groups : int list }
[@@deriving show, sexp]

  (* The input to the puzzle is a list of records. *)
type input = record list [@@deriving show, sexp]

module Parse = struct
  let parse_spring =
    choice
      [
        char '?' *> return Unknown;
        char '#' *> return Broken;
        char '.' *> return Working;
      ]

  let parse_row = many1 parse_spring

  let parse_record =
    let open Angstrom.Let_syntax in
    let%bind row = parse_row <* char ' ' >>| Array.of_list in
    let%bind damaged_groups =
      sep_by1 (char ',') (take_while1 Char.is_digit >>| Int.of_string)
    in
    return { row; damaged_groups }

  let parse_all = sep_by1 (char '\n') parse_record <* choice [ char '\n' >>| fun _ -> (); end_of_input ]

  let parse_input (s:string) : input =
    parse_string ~consume:All parse_all s |> Result.ok_or_failwith
end

(* for each record, find the number of possible ways to group
   the damaged springs according to the group specification.

   The desired result is the sum of all records' combinations. *)
module Part1 = struct
  (* A group (during traversal) is either a not-yet-started group of n broken springs,
     or a group of n broken springs that has already been entered.
     In the latter case, it first must be finished before allowing the next working spring. *)
  type group = Whole of int | Entered of int [@@deriving show, sexp]

  (* count combinations for a single record: row is the array of springs, ix is the current index,
     and groups is the current list of groups. *)
  let rec count row ix (groups : group list) =
    let len = Array.length row in
    match ix with
    | l when l = len -> ( match groups with [] | [ Entered 0 ] -> 1 | _ -> 0)
    | ix -> (
        match (row.(ix), groups) with
        | Broken, (Entered c | Whole c) :: cs when c > 0 ->
            count row (ix + 1) (Entered (c - 1) :: cs)
        | Broken, Entered 0 :: _ -> 0
        | Broken, _ -> 0
        | Working, Entered 0 :: cs -> count row (ix + 1) cs
        | Working, Entered _ :: _ -> 0 (* we're in a non-finished group of broken springs *)
        | Working, cs -> count row (ix + 1) cs
        | Unknown, [] -> (* assume ? = working *) count row (ix + 1) []
        | Unknown, Entered 0 :: cs ->
            (* assume ? = working because previous group is separated by working spring *)
            count row (ix + 1) cs
        | Unknown, Whole c :: cs ->
            assert (c > 0);
            (* first assume ? = broken *)
            let with_broken = count row (ix + 1) (Entered (c - 1) :: cs) in
            (* then assume ? = working; i.e. skip group *)
            let with_working = count row (ix + 1) (Whole c :: cs) in
            Out_channel.printf "broken: %d working: %d (ix %d)\n" with_broken
              with_working ix;
            with_broken + with_working
        | Unknown, Entered c :: cs ->
            assert (c > 0);
            (* forced assumption: ? = broken because we are in a group. *)
            count row (ix + 1) (Entered (c - 1) :: cs))

  (* count combinations for a single record *)
  let count_combinations { row; damaged_groups } =
    Out_channel.printf "\n";
    let groups = List.map ~f:(fun c -> Whole c) damaged_groups in
    count row 0 groups

      (* a list of counts of combinations per record. *)
  type combinations_counts = int list [@@deriving show]

    (* count combinations for each record *)
  let count_all_combinations records : combinations_counts =
    List.map ~f:count_combinations records
end

let () =
  let input = In_channel.(input_all stdin) in
  let parsed = Parse.parse_input input in
  let combos = Part1.count_all_combinations parsed in
  let sum = List.fold combos ~init:0 ~f:Int.(+) in
  Out_channel.(printf "%s\nsum: %d\n" (Part1.show_combinations_counts combos) sum)