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_in { src_start = src1; _ } { src_start = src2; _ } =
  Int.compare src1 src2

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

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

module Part1 = struct
  (* 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
  let amend_rangemap rm =
    let rm = List.sort ~compare:compare_rangemap_in rm in
    let lst = List.last_exn rm in
    let first_out =
      List.fold ~init:10000000000
        ~f:(fun min e -> Int.min min (e.src_start))
        rm
    in
    let last_out =
      List.fold ~init:0
        ~f:(fun max e -> Int.max max (e.dst_start + e.length))
        rm
    in
    ({ src_start = 0; dst_start = 0; length = first_out } :: rm)
    @ [
        {
          src_start = lst.src_start + lst.length;
          dst_start = last_out;
          length = 10000000000;
        };
      ]

  let last_rangemap inp =
    let last_step = List.last_exn steps_names in
    let last =
      List.sort ~compare:compare_rangemap_in
      @@ Hashtbl.find_exn inp.maps last_step
    in
    amend_rangemap last

  exception Range_overlap of (bool * bool * bool * bool * bool * bool)
  [@@deriving sexp]

  let translate ({ src_start = src1; dst_start = dst1; length = length1 } as rm)
      ({ src_start = src2; dst_start = dst2; length = length2 } as transl) =
    Out_channel.(
      printf "%s @ %s\n"
        (Sexp.to_string_hum (sexp_of_rangemap rm))
        (Sexp.to_string_hum (sexp_of_rangemap transl)));
    let starts_before = Int.(src1 < dst2) in
    let starts_within = Int.(src1 >= dst2 && src1 < dst2 + length2) in
    let starts_after = Int.(src1 >= dst2 + length2) in
    let ends_before = Int.(src1 + length1 < dst2) in
    let ends_within =
      Int.(src1 + length1 >= dst2 && src1 + length1 <= dst2 + length2)
    in
    let ends_after = Int.(src1 + length1 >= dst2 + length2) in
    if (starts_before && ends_before) || (starts_after && ends_after) then
      (None, starts_after && ends_after, false)
    else if starts_before && ends_within then (
      let length' = src1 + length1 - dst2 in
      let r =
        { src_start = src2; dst_start = dst1 + (dst2 - src1); length = length' }
      in
      assert (Int.(length' > 0));
      Out_channel.printf " => %s\n" (Sexp.to_string_hum (sexp_of_rangemap r));
      (Some r, false, true))
    else if starts_within && ends_within then (
      let length' = length1 in
      let r =
        { src_start = src2 + (src1 - dst2); dst_start = dst1; length = length' }
      in
      assert (Int.(length' > 0));
      Out_channel.printf " => %s\n" (Sexp.to_string_hum (sexp_of_rangemap r));
      (Some r, false, true))
    else if starts_within && ends_after then (
      let length' = dst2 + length2 - src1 in
      let r =
        { src_start = src2 + (src1 - dst2); dst_start = dst1; length = length' }
      in
      assert (Int.(length' > 0));
      Out_channel.printf " => %s\n" (Sexp.to_string_hum (sexp_of_rangemap r));
      (Some r, true, false))
    else if starts_before && ends_after then (
      let length' = length2 in
      let r =
        { src_start = src2; dst_start = dst1 + (dst2 - src1); length = length' }
      in
      assert (Int.(length' > 0));
      Out_channel.printf " => %s\n" (Sexp.to_string_hum (sexp_of_rangemap r));
      (Some r, true, false))
    else
      raise
        (Range_overlap
           ( starts_before,
             starts_within,
             starts_after,
             ends_before,
             ends_within,
             ends_after ))

  (* translate rangemap list rm into a new rangemap list using transl.
     The result is a rangemap incorporating the effect of
     both rangemaps. *)
  let rec rangemap_backwards rm transl =
    match (rm, transl) with
    | rm' :: rms, transl' :: transls -> (
        match translate rm' transl' with
        | Some r, false, true -> r :: rangemap_backwards rms (transl' :: transls)
        | Some r, true, false -> r :: rangemap_backwards (rm' :: rms) transls
        | None, false, false -> rangemap_backwards rms transls
        | None, true, false -> rangemap_backwards (rm' :: rms) transls
        | _, _, _ -> assert false)
    | [], _ | _, [] -> []

  let process ch =
    let input = In_channel.input_all ch in
    let input' = Parse.parse input in
    let final = last_rangemap input' in
    let maps =
      List.map
        ~f:(fun k -> amend_rangemap @@ Hashtbl.find_exn input'.maps k)
        steps_names
    in
    let transls = List.take maps (List.length maps - 1) in
    let translated =
      List.fold
        ~f:(fun current transl -> rangemap_backwards current transl)
        ~init:final (List.rev transls)
    in
    Out_channel.(
      print_endline
        (Sexp.to_string_hum (List.sexp_of_t sexp_of_rangemap translated)))
end

let () = Part2.process In_channel.stdin