open Base
open Angstrom
open Core

(* Core has its own different hashtbl. *)
module Hashtbl = Base.Hashtbl

(* typically: line x, character y *)
type position = { x : int; y : int } [@@deriving sexp]

module PositionKey : Hashtbl.Key.S with type t = position = struct
  type t = position

  let compare { x = ax; y = ay } { x = bx; y = by } =
    let cx = Int.compare ax bx in
    let cy = Int.compare ay by in
    if not (Int.equal cx 0) then cx else cy

  let sexp_of_t = sexp_of_position

  let hash { x; y } =
    let st = Hash.create () in
    let h1 = Hash.fold_int st x in
    let h2 = Hash.fold_int h1 y in
    Hash.get_hash_value h2
end

let create_position_tbl () : (position, 'a) Hashtbl.t =
  let position_hashtbl_key =
    (module PositionKey : Base.Hashtbl.Key.S with type t = position)
  in
  Hashtbl.create position_hashtbl_key

module Part1 = struct
  (* Parse lines like

      467..114..
      ...*......
      ..35..633.
      ......#...
      617*......
      .....+.58.
      ..592.....
      ......755.
      ...$.*....
      .664.598..
  *)
  type number = { start : position; value : int; digits : int }
  [@@deriving sexp]

  type symbol = { c : char; position : position } [@@deriving sexp]
  type item = Number of number | Symbol of symbol [@@deriving sexp]

  let syms = String.to_list "*#+$%/@=+-&_!?{}|[]"
  let is_sym c = List.mem syms c ~equal:Char.equal

  let intP =
    take_while1 (function '0' .. '9' -> true | _ -> false) >>| Int.of_string

  let numberP lineno =
    pos >>= fun p ->
    intP >>= fun i ->
    return
      {
        start = { x = lineno; y = p };
        value = i;
        digits = String.length (Int.to_string i);
      }

  let symP line =
    let to_sym l p c = return { c; position = { x = l; y = p } } in
    pos >>= fun p -> satisfy is_sym >>= to_sym line p

  let dotsP = skip_many (char '.')

  let number_or_symP lineno =
    choice
      [
        (numberP lineno >>= fun i -> return (Number i));
        (symP lineno >>= fun s -> return (Symbol s));
      ]

  let maybeP p = option () (p >>= fun _ -> return ())

  let lineP lineno =
    maybeP dotsP *> sep_by1 dotsP (number_or_symP lineno) <* maybeP dotsP

  exception ParseError of string

  let parse_line line lineno =
    let pr = parse_string ~consume:All (lineP lineno) line in
    match pr with Ok l -> l | Error e -> raise (ParseError e)

  let _debug_parse line lineno =
    let pr = parse_line line lineno in
    Sexp.to_string_hum (List.sexp_of_t sexp_of_item pr)

  (* fold over string lines and parse each line into `item list` *)
  let lines_folder (count, lines) line =
    let parsed = parse_line line count in
    (count + 1, parsed :: lines)

  (* read all lines from ch *)
  let read_lines ch =
    let _, all_lines = In_channel.fold_lines ch ~init:(0, []) ~f:lines_folder in
    all_lines

  (* build a hashmap of position -> symbol for all symbols in lines *)
  let build_symbol_map all_lines =
    let symmap = create_position_tbl () in
    let f = function
      | Symbol { c; position } ->
          ignore (Hashtbl.add symmap ~key:position ~data:c)
      | _ -> ()
    in
    let ff = List.iter ~f in
    List.iter ~f:ff all_lines;
    symmap

  (* return all number items from item list list (list of parsed lines) *)
  let all_numbers all_lines =
    let f = function Number _ -> true | _ -> false in
    let ff = List.filter ~f in
    List.concat (List.map ~f:ff all_lines)

  (* print symbol map for lines received on ch *)
  let _debug_symbol_map ch =
    let all_lines = read_lines ch in
    let symmap = build_symbol_map all_lines in
    let alist = Hashtbl.to_alist symmap in
    Sexp.to_string_hum
      (List.sexp_of_t
         (fun (a, b) -> List [ sexp_of_position a; Atom (Char.to_string b) ])
         alist)

  (* create 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

  (* generate adjacent positions for number at x, y with number of digits. *)
  let adjacent_positions { x; y } digits =
    let corners =
      [
        (x, y - 1);
        (x - 1, y - 1);
        (x + 1, y - 1);
        (x, y + digits);
        (x - 1, y + digits);
        (x + 1, y + digits);
      ]
    in
    let upper =
      Sequence.to_list
        (Sequence.map (count_seq y (y + digits)) ~f:(fun y' -> (x - 1, y')))
    in
    let lower =
      Sequence.to_list
        (Sequence.map (count_seq y (y + digits)) ~f:(fun y' -> (x + 1, y')))
    in
    List.concat [ corners; upper; lower ]

  (* check if there is a symbol at any adjacent position in the symbol map *)
  let rec check_adjacent_symbol symmap = function
    | [] -> false
    | (x, y) :: poss ->
        let ok = Hashtbl.mem symmap { x; y } in
        if ok then ok else check_adjacent_symbol symmap poss

  (* Using a symbol position map, check if a Number item has an adjacent symbol. *)
  let has_adjacent_symbol symmap = function
    | Number { start; digits; _ } ->
        let adjpos = adjacent_positions start digits in
        check_adjacent_symbol symmap adjpos
    | _ -> false

  (* from all parsed lines (item list list), return a flat list of all Number
     items with adjacent symbol. *)
  let filter_part_numbers all_lines =
    let symmap = build_symbol_map all_lines in
    let all_numbers = all_numbers all_lines in
    let part_numbers =
      List.filter ~f:(has_adjacent_symbol symmap) all_numbers
    in
    part_numbers

  (* Part 2 naturally depends on part 1: *)
  module Part2 = struct
    let build_gear_map all_lines =
      let symmap : (position, int * int) Hashtbl.t = create_position_tbl () in
      let f = function
        | Symbol { c = '*'; position } ->
            ignore (Hashtbl.add symmap ~key:position ~data:(0, 1))
        | Number _ | Symbol _ -> ()
      in
      let ff = List.iter ~f in
      List.iter ~f:ff all_lines;
      symmap

    (* part 2: *)
    let update_gearmap (gearmap : (position, int * int) Hashtbl.t) = function
      | Number { start; digits; value } ->
          let f (x, y) =
            (* only multiply two numbers *)
            let current = Hashtbl.find gearmap { x; y } in
            match current with
            | Some (count, value') when count < 2 ->
                Hashtbl.set gearmap ~key:{ x; y }
                  ~data:(count + 1, value * value')
            | _ -> ()
          in
          let adjpos = adjacent_positions start digits in
          List.iter ~f adjpos
      | _ -> ()

    (* part 2: *)
    let build_gear_table all_lines =
      let gearmap = build_gear_map all_lines in
      let all_numbers = all_numbers all_lines in
      let f n =
        if has_adjacent_symbol gearmap n then update_gearmap gearmap n else ()
      in
      List.iter ~f all_numbers;
      gearmap
  end

  let solve_parts_1_2 ch =
    let all_lines = read_lines ch in
    let pns = filter_part_numbers all_lines in
    let sum =
      List.fold ~init:0
        ~f:(fun acc -> function Number { value; _ } -> acc + value | _ -> acc)
        pns
    in
    Out_channel.printf "Sum is %d\n" sum;
    (* part 2: use already-parsed input *)
    let geartable = Part2.build_gear_table all_lines in
    let sexp_of_tuple (pos, (count, value)) =
      Sexp.List
        [
          sexp_of_position pos;
          Sexp.List [ Int.sexp_of_t count; Int.sexp_of_t value ];
        ]
    in
    let _gears_sexp =
      Sexp.to_string_hum
        (List.sexp_of_t sexp_of_tuple (Hashtbl.to_alist geartable))
    in
    let htfold ~key ~data:(count, value) acc =
      ignore key;
      if Int.equal count 2 then acc + value else acc
    in
    let gearratiosum = Hashtbl.fold geartable ~init:0 ~f:htfold in
    (*Out_channel.printf "Geartable is %s\n" gears_sexp*)
    Out_channel.printf "Gear ratio sum is %d\n" gearratiosum
  (*Sexp.to_string_hum (List.sexp_of_t sexp_of_item pns)*)
end

let () = Part1.solve_parts_1_2 In_channel.stdin