(* Module [Zipper]: functional linear zippers *)

open Base

type !+'a zipper = { pos : int; before : 'a Sequence.t; after : 'a Sequence.t }
type !+'a t = 'a zipper

let empty = { pos = 0; before = Sequence.empty; after = Sequence.empty }
let before z = z.before
let after z = z.after
let focus z = after z |> Sequence.next |> Option.map ~f:fst
let focus_or ~default z = Option.value ~default (focus z)
let is_far_left z = before z |> Sequence.is_empty

let is_far_right ?(by_one = false) z =
  after z |> Sequence.length_is_bounded_by ~max:(if by_one then 1 else 0)

let is_empty z = is_far_left z && is_far_right z
let left_length z = z.pos
let right_length z = after z |> Sequence.length
let length z = left_length z + right_length z

let left z =
  match Sequence.next z.before with
  | None -> z
  | Some (h, t) ->
      { pos = z.pos - 1; before = t; after = Sequence.shift_right z.after h }

let rec left_while f z =
  if (not (is_far_left z)) && Option.(focus z |> map ~f |> value ~default:false)
  then left z |> left_while f
  else z

let rec far_left z = if is_far_left z then z else z |> left |> far_left

let right ?(by_one = false) z =
  match Sequence.next z.after with
  | None -> z
  | Some (_, t) when by_one && Sequence.is_empty t -> z
  | Some (h, t) ->
      { pos = z.pos + 1; before = Sequence.shift_right z.before h; after = t }

let rec right_while ?(by_one = false) f z =
  if
    (not (is_far_right ~by_one z))
    && Option.(focus z |> map ~f |> value ~default:false)
  then right z |> right_while ~by_one f
  else z

let rec far_right ?(by_one = false) z =
  if is_far_right ~by_one z then z else z |> right |> far_right ~by_one

let goto ?(by_one = false) n z =
  let n = n - z.pos in
  let step = if n < 0 then left else right ~by_one in
  Fn.apply_n_times ~n:(abs n) step z

let pop ?(n = 1) z =
  let a, b = Sequence.split_n z.after n in
  (Sequence.of_list a, { z with after = b })

let pop_before ?(n = 1) =
  let rec aux m z =
    if m < n && not (is_far_left z) then aux (m + 1) (left z) else pop ~n:m z
  in
  aux 0

let pop_after ?(n = 1) z =
  if right_length z < 2 then (Sequence.empty, z)
  else
    let a, b = right z |> pop ~n in
    (a, left b)

let push x z = { z with after = Sequence.shift_right z.after x }
let push_seq s z = { z with after = Sequence.append s z.after }
let push_after x z = right z |> push x |> left
let push_after_seq s z = right z |> push_seq s |> left

let push_before x z =
  { z with pos = z.pos + 1; before = Sequence.shift_right z.before x }

let push_before_seq s z =
  let f a e = Sequence.shift_right a e in
  {
    z with
    pos = z.pos + Sequence.length s;
    before = Sequence.fold ~init:z.before ~f s;
  }

let split z =
  ( { z with after = Sequence.empty },
    { z with pos = 0; before = Sequence.empty } )

let join z1 ~z2 =
  let z1 = far_right z1 and z2 = far_left z2 in
  { z1 with after = z2.after }

let iter_left f z = Sequence.iter ~f z.before
let iter_right f z = Sequence.iter ~f z.after

let iter f z =
  iter_left f z;
  iter_right f z

let for_all f z = Sequence.(for_all ~f z.before && for_all ~f z.after)
let exists f z = Sequence.(exists ~f z.before || exists ~f z.after)
let find_left f z = Sequence.find ~f z.before
let find_right f z = Sequence.find ~f z.after
let apply_focus f z = Option.map ~f (focus z)
let apply_focus_or ~default f z = Option.value ~default (apply_focus f z)
let map_left f z = { z with before = Sequence.map ~f z.before }
let map_right f z = { z with after = Sequence.map ~f z.after }

let map_focus f z =
  match Sequence.next z.after with
  | None -> z
  | Some (h, t) -> { z with after = Sequence.shift_right t (f h) }

let map_focus_or ~default f z =
  match Sequence.next z.after with
  | None -> { z with after = Sequence.singleton default }
  | Some (h, t) -> { z with after = Sequence.shift_right t (f h) }

let map f z =
  { z with before = Sequence.map ~f z.before; after = Sequence.map ~f z.after }

let mapi_left f z = { z with before = Sequence.mapi ~f z.before }
let mapi_right f z = { z with after = Sequence.mapi ~f z.after }

let mapi f z =
  {
    z with
    before = Sequence.mapi ~f z.before;
    after = Sequence.mapi ~f z.after;
  }

let filter_left f z = { z with before = Sequence.filter ~f z.before }
let filter_right f z = { z with after = Sequence.filter ~f z.after }
let filter p z = z |> filter_left p |> filter_right p
let context_left n z = { z with before = Sequence.take z.before n }
let context_right n z = { z with after = Sequence.take z.after n }
let context ~l ?(r = l) z = z |> context_left l |> context_right r
let swap_focus a = map_focus_or ~default:a (Fn.const a)
let of_seq s = { empty with after = s }
let to_seq z = z |> far_left |> after
let window ~from ~len z = goto from z |> context_right len |> after