open Angstrom
open Base
open Core
module Hashtbl = Base.Hashtbl

type rangemap = { dst_start : int; src_start : int; length : int }
[@@deriving sexp]

let compare_rangemap { dst_start = dst1; _ } { dst_start = dst2; _ } =
  Int.compare dst1 dst2

type input = { seeds : int list; maps : (string, rangemap list) Hashtbl.t }

let sexp_of_input { seeds; maps } =
  Sexp.List
    [
      Sexp.List [ Sexp.Atom "seeds"; List.sexp_of_t Int.sexp_of_t seeds ];
      Sexp.List
        [
          Sexp.Atom "maps";
          Hashtbl.sexp_of_t String.sexp_of_t
            (List.sexp_of_t sexp_of_rangemap)
            maps;
        ];
    ]

module Parse = struct
  let maybe p = option () (p >>= fun _ -> return ())

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

  let int_listP = skip_many (char ' ') *> sep_by1 (skip_many1 (char ' ')) intP

  let rangeP =
    int_listP >>= function
    | dst_start :: src_start :: length :: _rest ->
        return { dst_start; src_start; length }
    | _ -> assert false

  let map_nameP = take_while1 (fun c -> not (Char.is_whitespace c))
  let map_headerP = map_nameP <* string " map:\n"
  let rangesP = many1 (rangeP <* char '\n') <* maybe (char '\n')

  let mapP =
    let open Angstrom.Let_syntax in
    let%bind mapname = map_headerP in
    let%bind ranges = rangesP in
    return (mapname, List.sort ~compare:compare_rangemap ranges)

  let seedsP = string "seeds: " *> int_listP <* string "\n\n"

  let inputP =
    let open Angstrom.Let_syntax in
    let%bind seeds = seedsP in
    let%bind maps = many1 mapP in
    return (seeds, maps)

  exception Parse_exn of string

  let parse input =
    match parse_string ~consume:All inputP input with
    | Ok (seeds, maps) ->
        let maps = Hashtbl.of_alist_exn (module String) maps in
        { seeds; maps }
    | Error e -> raise (Parse_exn e)

  let _test_parse input =
    let inp = parse input in
    Out_channel.print_endline (Sexp.to_string_hum (sexp_of_input inp))
end

module Part1 = struct
  let steps_names =
    [
      "seed-to-soil";
      "soil-to-fertilizer";
      "fertilizer-to-water";
      "water-to-light";
      "light-to-temperature";
      "temperature-to-humidity";
      "humidity-to-location";
    ]

  (* for a single step, resolve the number using the provided range maps. Return None if no mapping was found. *)
  let resolve_step_maybe { dst_start; src_start; length } number =
    if src_start <= number && number <= src_start + length then
      Some (dst_start + (number - src_start))
    else None

  (* For all maps in a step, determine the final next number. *)
  let rec resolve_step maps number =
    match maps with
    | map :: maps' -> (
        match resolve_step_maybe map number with
        | Some ok -> ok
        | None -> resolve_step maps' number)
    | [] -> number

  (* Traverse steps for one seed *)
  let rec traverse_for_seed seed = function
    | step :: steps ->
        let next = resolve_step step seed in
        traverse_for_seed next steps
    | [] -> seed

  (* optain list of steps from hashmap *)
  let maps_list ?(steps = steps_names) maps_hm =
    let f = Hashtbl.find_exn maps_hm in
    List.map steps ~f

  (* For all seeds, find resulting location (end number). *)
  let traverse_all { seeds; maps } =
    let maps = maps_list maps in
    let results = List.map seeds ~f:(fun seed -> traverse_for_seed seed maps) in
    results

  (* Read, parse, and process input from channel. *)
  let _process ch =
    let input = In_channel.input_all ch in
    let input' = Parse.parse input in
    let result = traverse_all input' in
    Out_channel.print_endline (Sexp.to_string_hum (sexp_of_input input'));
    Out_channel.print_endline
      (Sexp.to_string_hum (List.sexp_of_t Int.sexp_of_t result));
    Out_channel.printf "Result is %d\n"
      (List.fold ~init:Int.max_value ~f:Int.min result)
end

module Part2 = struct
  type range = { from : int; length : int } [@@deriving sexp]

  let _intersect { from = from1; length = length1 }
      { from = from2; length = length2 } =
    let f1, t1 = (from1, from1 + length1) in
    let f2, t2 = (from2, from2 + length2) in
    let newfrom = Int.max f1 f2 in
    let newto = Int.min t1 t2 in
    match (newfrom, newto) with
    | a, b when a < b -> Some { from = a; length = b - a + 1 }
    | _ -> None

  let _all_ranges maps =
    let first = List.hd_exn maps in
    let last = List.last_exn maps in
    let range1 = { dst_start = 0; src_start = 0; length = first.dst_start } in
    let range_last =
      {
        src_start = last.dst_start + last.length;
        dst_start = last.dst_start + last.length;
        length = Int.max_value - last.dst_start - last.length;
      }
    in
    let maps = if range1.length > 0 then range1 :: maps else maps in
    if range_last.length > 0 then List.append maps [ range_last ] else maps

  type new_input = {
    (* seed_ranges : range list; *)
    _maps : (string, rangemap list) Hashtbl.t;
  }

  let _convert_new_seeds_to_old ({ maps = _maps; _ } : input) =
    { (*seed_ranges = convert_list seeds;*) _maps }
end

let () = Part1._process In_channel.stdin