open Angstrom
open Base
open Core

type color = Red | Green | Blue [@@deriving sexp]
type reveal = { red : int; green : int; blue : int } [@@deriving sexp]
type game = { id : int; reveals : reveal list } [@@deriving sexp]

module Part1 = struct
  module Parser = struct
    (* Parse a game like
       "Game 4: 1 green, 3 red, 6 blue; 3 green, 6 red; 3 green, 15 blue, 14 red"
    *)
    let intP =
      take_while1 (function '0' .. '9' -> true | _ -> false) >>| Int.of_string

    let game_idP = string "Game " *> intP <* string ": "

    let colorP =
      choice
        [
          string "red" *> return Red;
          string "blue" *> return Blue;
          string "green" *> return Green;
        ]

    let color_revealP =
      let%bind.Angstrom count = intP in
      let%bind.Angstrom () = skip_while Char.is_whitespace in
      let%bind.Angstrom color = colorP in
      return (count, color)

    let revealP =
      let f acc x =
        match x with
        | c, Red -> { acc with red = c }
        | c, Blue -> { acc with blue = c }
        | c, Green -> { acc with green = c }
      in
      let init = { red = 0; blue = 0; green = 0 } in
      let%bind.Angstrom colors = sep_by1 (string ", ") color_revealP in
      return @@ List.fold ~f ~init colors

    let gameP =
      let%bind.Angstrom id = game_idP in
      let%bind.Angstrom reveals = sep_by1 (string "; ") revealP in
      return { id; reveals }

    let parse_game = parse_string ~consume:All gameP
  end

  exception Game_parse_failed of string

  let parse_line line =
    match Parser.parse_game line with
    | Ok o -> o
    | Error e -> raise (Game_parse_failed e)

  let check_reveal contents reveal =
    let { red = cr; green = cg; blue = cb } = contents in
    let { red = rr; green = rg; blue = rb } = reveal in
    rr <= cr && rg <= cg && rb <= cb

  let check_game contents game =
    let f acc x = acc && check_reveal contents x in
    List.fold ~init:true ~f game

  let check_line contents line =
    let game = parse_line line in
    if check_game contents game.reveals then game.id else 0
end

module Part2 = struct
  let min_contents_fold acc rev =
    let { red = ar; green = ag; blue = ab } = acc in
    let { red = rr; green = rg; blue = rb } = rev in
    Int.{ red = max ar rr; green = max ag rg; blue = max ab rb }

  let min_contents =
    List.fold ~init:{ red = 0; green = 0; blue = 0 } ~f:min_contents_fold

  let power game =
    let mc = min_contents game.reveals in
    mc.red * mc.green * mc.blue
end

let _fold_lines1 contents =
  let f acc line = acc + Part1.check_line contents line in
  In_channel.fold_lines In_channel.stdin ~f ~init:0

let fold_lines2 () =
  let f acc line = acc + (Part2.power @@ Part1.parse_line line) in
  In_channel.fold_lines In_channel.stdin ~f ~init:0

let () =
  Out_channel.printf "%d\n"
    (*(fold_lines1 Part1.Parser.{ red = 12; green = 13; blue = 14 })*)
    (fold_lines2 ())