[@@@warning "-32"] (* pretty-printers *)

type vec = { off : int; len : int }
type box = Box | TBox of int | BBox

let pp_box ppf = function
  | Box -> Fmt.string ppf "box"
  | TBox tab -> Fmt.pf ppf "(TBox %d)" tab
  | BBox -> Fmt.string ppf "bbox"

type value =
  | String of vec * string
  | Bytes of vec * bytes
  | Bigstring of vec * Bigstringaf.t

let pp_value ppf = function
  | String (vec, v) -> Fmt.pf ppf "%S" (String.sub v vec.off vec.len)
  | Bytes (vec, v) -> Fmt.pf ppf "%S" (Bytes.sub_string v vec.off vec.len)
  | Bigstring _ -> Fmt.pf ppf "#bigstring"

let split_value len x =
  assert (len > 0);

  match x with
  | String (vec, v) ->
      let v0 = String.sub v 0 len in
      let v1 = String.sub v len (vec.len - len) in
      ( String ({ off = 0; len }, v0),
        String ({ off = 0; len = vec.len - len }, v1) )
  | Bytes (vec, v) ->
      let v0 = Bytes.sub v 0 len in
      let v1 = Bytes.sub v len (vec.len - len) in
      ( Bytes ({ off = 0; len }, v0),
        Bytes ({ off = 0; len = vec.len - len }, v1) )
  | Bigstring (vec, v) ->
      ( Bigstring ({ off = vec.off; len }, v),
        Bigstring ({ off = vec.off + len; len = vec.len - len }, v) )

let length_of_value = function
  | String (vec, _) | Bytes (vec, _) | Bigstring (vec, _) -> vec.len

type atom =
  | Breakable of value
  | Unbreakable of value
  | Break of { len : int; indent : int }
  | New_line
  | Open of box
  | Close

let pp_atom ppf = function
  | Breakable v -> Fmt.pf ppf "<breakable:%a>" pp_value v
  | Unbreakable v -> Fmt.pf ppf "<unbreakable:%a>" pp_value v
  | Break { len; indent } -> Fmt.pf ppf "<break:len= %d, indent= %d>" len indent
  | New_line -> Fmt.pf ppf "<new-line>"
  | Open box -> Fmt.pf ppf "(box %a" pp_box box
  | Close -> Fmt.pf ppf ")"

let box = Box
let tbox indent = TBox indent
let bbox = BBox
let o box = Open box
let new_line = New_line
let close = Close
let fws = Break { len = 1; indent = 1 }
let spaces len = Break { len; indent = 0 }
let break ~len ~indent = Break { len; indent }

let v ~breakable =
  match breakable with
  | true -> fun x -> Breakable x
  | false -> fun x -> Unbreakable x

let string ?(breakable = false) ?(off = 0) ?len x =
  let len = match len with Some len -> len | None -> String.length x - off in
  let value = String ({ off; len }, x) in
  v ~breakable value

let bytes ?(breakable = false) ?(off = 0) ?len x =
  let len = match len with Some len -> len | None -> Bytes.length x - off in
  let value = Bytes ({ off; len }, x) in
  v ~breakable value

let bigstring ?(breakable = false) ?(off = 0) ?len x =
  let len =
    match len with Some len -> len | None -> Bigstringaf.length x - off
  in
  let value = Bigstring ({ off; len }, x) in
  v ~breakable value

type token =
  | TValue of value
  | TBreak of int
  | TBox of [ `Root | `Box | `Indent of int ]
  | TClose

let length_of_token = function
  | TValue value -> length_of_value value
  | TBreak len -> len
  | TBox _ -> 0
  | TClose -> 0

module Option = struct
  let bind x f = match x with Some x -> f x | None -> None let ( >>= ) = bind
end

module Stack : sig
  type +'a t

  val empty : _ t
  val push : 'a -> 'a t -> 'a t
  val swap_exn : ('a -> 'a) -> 'a t -> 'a t
  val pop : 'a t -> ('a * 'a t) option
  val fold : ('a -> 'b -> 'a) -> 'a -> 'b t -> 'a
  val tail_exn : 'a t -> 'a t
  val map : ('a -> 'b) -> 'a t -> 'b t
  val pp : 'a Fmt.t -> 'a t Fmt.t
end = struct
  type 'a t = 'a list

  let empty = []
  let push x t = x :: t

  exception Empty

  let swap_exn f = function [] -> raise Empty | x :: r -> f x :: r
  let pop = function [] -> None | x :: r -> Some (x, r)
  let fold = List.fold_left
  let tail_exn = function _ :: r -> r | [] -> raise Empty
  let map f l = List.map f l
  let pp = Fmt.Dump.list
end

module Queue = Ke.Fke

type t = {
  boxes : [ `Root | `Box | `Indent of int ] Stack.t;
  breaks : [ `Indent of int ] list Stack.t;
  inner : int list Stack.t;
  indent : int;
  margin : int;
  new_line : string;
  queue : token Queue.t;
  encoder : Enclosure.encoder;
}

let pp_box ppf = function
  | `Root -> Fmt.string ppf "`Root"
  | `Box -> Fmt.string ppf "`Box"
  | `Indent n -> Fmt.pf ppf "(`Indent %d)" n

let pp_break ppf (`Indent n) = Fmt.pf ppf "(`Indent %d)" n

let pp_token ppf = function
  | TValue (String ({ off; len }, x)) -> Fmt.pf ppf "%S" (String.sub x off len)
  | TValue (Bytes ({ off; len }, x)) ->
      Fmt.pf ppf "%S" (Bytes.sub_string x off len)
  | TValue (Bigstring ({ off; len }, x)) ->
      Fmt.pf ppf "%S" (Bigstringaf.substring x ~off ~len)
  | TBreak len -> Fmt.pf ppf "<%S>" (String.make len ' ')
  | TBox `Box -> Fmt.pf ppf "["
  | TBox (`Indent n) -> Fmt.pf ppf "[<%d>" n
  | TBox `Root -> Fmt.pf ppf "[<root>"
  | TClose -> Fmt.pf ppf "]"

let pp ppf t =
  Fmt.pf ppf
    "{ @[<hov>boxes= @[<hov>%a@];@ breaks= @[<hov>%a@];@ inner= @[<hov>%a@];@ \
     indent= %d;@ margin= %d;@ new_line= %S;@ queue= @[<hov>%a@];@ encoder= \
     @[<hov>%a@];@] }"
    (Stack.pp pp_box) t.boxes
    (Stack.pp (Fmt.Dump.list pp_break))
    t.breaks
    (Stack.pp Fmt.(Dump.list int))
    t.inner t.indent t.margin t.new_line (Queue.pp pp_token) t.queue
    Enclosure.pp t.encoder

let ( <.> ) f g x = f (g x)
let flip f a b = f b a

let merge_breaks token (queue, x) =
  match (token, x) with
  | Break { len = len_x; _ }, TBreak len -> Some (queue, len_x + len)
  | _, _ -> None

let current_length_of_line t =
  t.indent + Queue.fold (flip (( + ) <.> length_of_token)) 0 t.queue

let emit_line k t =
  let rec go queue encoder =
    match Queue.pop queue with
    | Some (TValue (String ({ off; len }, v)), queue) ->
        Enclosure.kschedule_string (go queue) encoder ~off ~len v
    | Some (TValue (Bytes ({ off; len }, v)), queue) ->
        Enclosure.kschedule_bytes (go queue) encoder ~off ~len v
    | Some (TValue (Bigstring ({ off; len }, v)), queue) ->
        Enclosure.kschedule_bigstring (go queue) encoder ~off ~len v
    | Some (TBreak len, queue) ->
        Enclosure.kschedule_string (go queue) encoder ~len (String.make len ' ')
    | Some (TBox _, queue) | Some (TClose, queue) -> go queue encoder
    | None ->
        let k encoder = k { t with encoder; queue = Queue.empty } in
        let k encoder = Enclosure.flush k encoder in
        Enclosure.kschedule_string k encoder t.new_line
  in
  Enclosure.kschedule_string (go t.queue) t.encoder (String.make t.indent ' ')

let merge_indents k t =
  let indent_by_box =
    Stack.fold
      (fun a -> function `Box | `Root -> a | `Indent n -> a + n)
      0 t.boxes
  in
  let indent_by_break =
    let k a l =
      List.fold_left
        (fun a -> function `Indent n -> (max : int -> int -> int) a n)
        a l
    in
    Stack.fold k 0 t.breaks
  in
  k
    {
      t with
      indent = indent_by_box + indent_by_break;
      inner = Stack.map (fun _ -> []) t.boxes;
      breaks = Stack.map (fun _ -> []) t.boxes;
    }

let is_in_box t =
  let exception In in
  try
    Queue.rev_iter (function TBox _ -> raise In | _ -> ()) t.queue;
    false
  with In -> true

let without_last_box q =
  let rec go q =
    match Queue.tail q with
    | Some (q, TBox _) -> q
    | Some (q, _) -> go q
    | None -> Queue.empty
  in
  go q

let unroll_last_box_and_emit k value t =
  let rec unroll acc q =
    let q, x = Queue.tail_exn q in
    match x with TBox _ -> (List.rev (x :: acc), q) | x -> unroll (x :: acc) q
  in
  let box, q = unroll [ TValue value ] t.queue in
  (* assert (Queue.is_empty q = false); *)
  let k t =
    let rec roll acc q =
      match acc with [] -> q | x :: r -> roll r (Queue.push q x)
    in
    let q = roll (List.rev box) t.queue in
    k { t with queue = q }
  in
  emit_line (merge_indents k) { t with queue = q }

let rec kpush_breakable_value ~current_length_of_line k value t =
  if current_length_of_line >= t.margin then
    emit_line (merge_indents (kpush k (Breakable value))) t
  else if current_length_of_line + length_of_token (TValue value) > t.margin
  then
    if is_in_box t && not (Queue.is_empty (without_last_box t.queue)) then
      unroll_last_box_and_emit k value t
    else
      let len = t.margin - current_length_of_line in
      let value0, value1 = split_value len value in
      let token0 = TValue value0 in
      let token1 = v ~breakable:true value1 in
      emit_line
        (merge_indents (kpush k token1))
        { t with queue = Queue.push t.queue token0 }
  else
    k
      {
        t with
        queue = Queue.push t.queue (TValue value);
        inner =
          Stack.swap_exn (fun lenv -> length_of_value value :: lenv) t.inner;
      }

and kpush_unbreakable_value ~current_length_of_line k value t =
  if
    current_length_of_line + length_of_token (TValue value) > t.margin
    && not (Queue.is_empty t.queue)
  then
    if is_in_box t && not (Queue.is_empty (without_last_box t.queue)) then
      unroll_last_box_and_emit k value t
    else
      match Queue.tail t.queue with
      | Some (queue, TBreak _) ->
          emit_line
            (merge_indents (kpush k (Unbreakable value)))
            { t with queue }
      | Some _ -> emit_line (merge_indents (kpush k (Unbreakable value))) t
      | None -> emit_line (merge_indents (kpush k (Unbreakable value))) t
  else
    k
      {
        t with
        queue = Queue.push t.queue (TValue value);
        inner =
          Stack.swap_exn (fun lenv -> length_of_value value :: lenv) t.inner;
      }

and kpush k value t =
  let current_length_of_line = current_length_of_line t in

  let append stack len =
    match Stack.pop stack with
    | Some (lenv, stack) -> Stack.push (len :: lenv) stack
    | None -> assert false
  in

  match value with
  | New_line -> emit_line (merge_indents k) t
  | Open Box ->
      k
        {
          t with
          boxes = Stack.push `Box t.boxes;
          inner = Stack.push [] t.inner;
          breaks = Stack.push [] t.breaks;
          queue = Queue.push t.queue (TBox `Box);
        }
  | Open (TBox indent) ->
      k
        {
          t with
          boxes = Stack.push (`Indent indent) t.boxes;
          inner = Stack.push [] t.inner;
          breaks = Stack.push [] t.breaks;
          queue = Queue.push t.queue (TBox (`Indent indent));
        }
  | Open BBox ->
      let indent =
        Stack.fold (flip (( + ) <.> List.fold_left ( + ) 0)) 0 t.inner
      in
      k
        {
          t with
          boxes = Stack.push (`Indent indent) t.boxes;
          inner = Stack.push [] t.inner;
          breaks = Stack.push [] t.breaks;
          queue = Queue.push t.queue (TBox (`Indent indent));
        }
  | Close ->
      (* XXX(dinosaure): check [`Root] box. *)
      k
        {
          t with
          boxes = Stack.tail_exn t.boxes;
          inner = Stack.tail_exn t.inner;
          breaks = Stack.tail_exn t.breaks;
          queue = Queue.push t.queue TClose;
        }
  | Breakable value -> kpush_breakable_value ~current_length_of_line k value t
  | Unbreakable value ->
      kpush_unbreakable_value ~current_length_of_line k value t
  | Break { len; indent } as break -> (
      match
        let open Option in
        Queue.tail t.queue >>= merge_breaks break
      with
      | Some (queue, len) ->
          if current_length_of_line + length_of_token (TBreak len) >= t.margin
          then
            emit_line (merge_indents k)
              { t with queue; breaks = append t.breaks (`Indent indent) }
          else
            k
              {
                t with
                queue = Queue.push queue (TBreak len);
                inner = append t.inner len;
                breaks = append t.breaks (`Indent indent);
              }
      | None ->
          if current_length_of_line + length_of_token (TBreak len) > t.margin
          then
            emit_line (merge_indents k)
              { t with breaks = append t.breaks (`Indent indent) }
          else
            k
              {
                t with
                queue = Queue.push t.queue (TBreak len);
                inner = append t.inner len;
                breaks = append t.breaks (`Indent indent);
              })

let kflush k t =
  let rec go queue encoder =
    match Queue.pop queue with
    | Some (TValue (String ({ off; len }, v)), queue) ->
        Enclosure.kschedule_string (go queue) encoder ~off ~len v
    | Some (TValue (Bytes ({ off; len }, v)), queue) ->
        Enclosure.kschedule_bytes (go queue) encoder ~off ~len v
    | Some (TValue (Bigstring ({ off; len }, v)), queue) ->
        Enclosure.kschedule_bigstring (go queue) encoder ~off ~len v
    | Some (TBreak len, queue) ->
        Enclosure.kschedule_string (go queue) encoder ~len (String.make len ' ')
    | Some (TBox _, queue) | Some (TClose, queue) -> go queue encoder
    | None ->
        Enclosure.flush
          (fun encoder -> k { t with encoder; queue = Queue.empty })
          encoder
  in
  go t.queue t.encoder

external identity : 'a -> 'a = "%identity"

let push = kpush identity
let flush = kflush identity
let is_empty t = Queue.is_empty t.queue && Enclosure.is_empty t.encoder

let create ?(margin = 998) ?(new_line = "\r\n") ~emitter len =
  let encoder = Enclosure.create ~emitter len in
  {
    encoder;
    queue = Queue.empty;
    boxes = Stack.push `Root Stack.empty;
    breaks = Stack.push [] Stack.empty;
    inner = Stack.push [] Stack.empty;
    indent = 0;
    margin;
    new_line;
  }

 pretty-printers 
 assert (Queue.is_empty q = false); 
 XXX(dinosaure): check [`Root] box. 

type vec = { off : int; len : int }
type box = Box | TBox of int | BBox

let pp_box ppf = function
  | Box -> Fmt.string ppf "box"
  | TBox tab -> Fmt.pf ppf "(TBox %d)" tab
  | BBox -> Fmt.string ppf "bbox"

type value =
  | String of vec * string
  | Bytes of vec * bytes
  | Bigstring of vec * Bigstringaf.t

let pp_value ppf = function
  | String (vec, v) -> Fmt.pf ppf "%S" (String.sub v vec.off vec.len)
  | Bytes (vec, v) -> Fmt.pf ppf "%S" (Bytes.sub_string v vec.off vec.len)
  | Bigstring _ -> Fmt.pf ppf "#bigstring"

let split_value len x =
  assert (len > 0);

  match x with
  | String (vec, v) ->
      let v0 = String.sub v 0 len in
      let v1 = String.sub v len (vec.len - len) in
      ( String ({ off = 0; len }, v0),
        String ({ off = 0; len = vec.len - len }, v1) )
  | Bytes (vec, v) ->
      let v0 = Bytes.sub v 0 len in
      let v1 = Bytes.sub v len (vec.len - len) in
      ( Bytes ({ off = 0; len }, v0),
        Bytes ({ off = 0; len = vec.len - len }, v1) )
  | Bigstring (vec, v) ->
      ( Bigstring ({ off = vec.off; len }, v),
        Bigstring ({ off = vec.off + len; len = vec.len - len }, v) )

let length_of_value = function
  | String (vec, _) | Bytes (vec, _) | Bigstring (vec, _) -> vec.len

type atom =
  | Breakable of value
  | Unbreakable of value
  | Break of { len : int; indent : int }
  | New_line
  | Open of box
  | Close

let pp_atom ppf = function
  | Breakable v -> Fmt.pf ppf "<breakable:%a>" pp_value v
  | Unbreakable v -> Fmt.pf ppf "<unbreakable:%a>" pp_value v
  | Break { len; indent } -> Fmt.pf ppf "<break:len= %d, indent= %d>" len indent
  | New_line -> Fmt.pf ppf "<new-line>"
  | Open box -> Fmt.pf ppf "(box %a" pp_box box
  | Close -> Fmt.pf ppf ")"

let box = Box
let tbox indent = TBox indent
let bbox = BBox
let o box = Open box
let new_line = New_line
let close = Close
let fws = Break { len = 1; indent = 1 }
let spaces len = Break { len; indent = 0 }
let break ~len ~indent = Break { len; indent }

let v ~breakable =
  match breakable with
  | true -> fun x -> Breakable x
  | false -> fun x -> Unbreakable x

let string ?(breakable = false) ?(off = 0) ?len x =
  let len = match len with Some len -> len | None -> String.length x - off in
  let value = String ({ off; len }, x) in
  v ~breakable value

let bytes ?(breakable = false) ?(off = 0) ?len x =
  let len = match len with Some len -> len | None -> Bytes.length x - off in
  let value = Bytes ({ off; len }, x) in
  v ~breakable value

let bigstring ?(breakable = false) ?(off = 0) ?len x =
  let len =
    match len with Some len -> len | None -> Bigstringaf.length x - off
  in
  let value = Bigstring ({ off; len }, x) in
  v ~breakable value

type token =
  | TValue of value
  | TBreak of int
  | TBox of [ `Root | `Box | `Indent of int ]
  | TClose

let length_of_token = function
  | TValue value -> length_of_value value
  | TBreak len -> len
  | TBox _ -> 0
  | TClose -> 0

module Option = struct
  let bind x f = match x with Some x -> f x | None -> None let ( >>= ) = bind
end

module Stack : sig
  type +'a t

  val empty : _ t
  val push : 'a -> 'a t -> 'a t
  val swap_exn : ('a -> 'a) -> 'a t -> 'a t
  val pop : 'a t -> ('a * 'a t) option
  val fold : ('a -> 'b -> 'a) -> 'a -> 'b t -> 'a
  val tail_exn : 'a t -> 'a t
  val map : ('a -> 'b) -> 'a t -> 'b t
  val pp : 'a Fmt.t -> 'a t Fmt.t
end = struct
  type 'a t = 'a list

  let empty = []
  let push x t = x :: t

  exception Empty

  let swap_exn f = function [] -> raise Empty | x :: r -> f x :: r
  let pop = function [] -> None | x :: r -> Some (x, r)
  let fold = List.fold_left
  let tail_exn = function _ :: r -> r | [] -> raise Empty
  let map f l = List.map f l
  let pp = Fmt.Dump.list
end

module Queue = Ke.Fke

type t = {
  boxes : [ `Root | `Box | `Indent of int ] Stack.t;
  breaks : [ `Indent of int ] list Stack.t;
  inner : int list Stack.t;
  indent : int;
  margin : int;
  new_line : string;
  queue : token Queue.t;
  encoder : Enclosure.encoder;
}

let pp_box ppf = function
  | `Root -> Fmt.string ppf "`Root"
  | `Box -> Fmt.string ppf "`Box"
  | `Indent n -> Fmt.pf ppf "(`Indent %d)" n

let pp_break ppf (`Indent n) = Fmt.pf ppf "(`Indent %d)" n

let pp_token ppf = function
  | TValue (String ({ off; len }, x)) -> Fmt.pf ppf "%S" (String.sub x off len)
  | TValue (Bytes ({ off; len }, x)) ->
      Fmt.pf ppf "%S" (Bytes.sub_string x off len)
  | TValue (Bigstring ({ off; len }, x)) ->
      Fmt.pf ppf "%S" (Bigstringaf.substring x ~off ~len)
  | TBreak len -> Fmt.pf ppf "<%S>" (String.make len ' ')
  | TBox `Box -> Fmt.pf ppf "["
  | TBox (`Indent n) -> Fmt.pf ppf "[<%d>" n
  | TBox `Root -> Fmt.pf ppf "[<root>"
  | TClose -> Fmt.pf ppf "]"

let pp ppf t =
  Fmt.pf ppf
    "{ @[<hov>boxes= @[<hov>%a@];@ breaks= @[<hov>%a@];@ inner= @[<hov>%a@];@ \
     indent= %d;@ margin= %d;@ new_line= %S;@ queue= @[<hov>%a@];@ encoder= \
     @[<hov>%a@];@] }"
    (Stack.pp pp_box) t.boxes
    (Stack.pp (Fmt.Dump.list pp_break))
    t.breaks
    (Stack.pp Fmt.(Dump.list int))
    t.inner t.indent t.margin t.new_line (Queue.pp pp_token) t.queue
    Enclosure.pp t.encoder

let ( <.> ) f g x = f (g x)
let flip f a b = f b a

let merge_breaks token (queue, x) =
  match (token, x) with
  | Break { len = len_x; _ }, TBreak len -> Some (queue, len_x + len)
  | _, _ -> None

let current_length_of_line t =
  t.indent + Queue.fold (flip (( + ) <.> length_of_token)) 0 t.queue

let emit_line k t =
  let rec go queue encoder =
    match Queue.pop queue with
    | Some (TValue (String ({ off; len }, v)), queue) ->
        Enclosure.kschedule_string (go queue) encoder ~off ~len v
    | Some (TValue (Bytes ({ off; len }, v)), queue) ->
        Enclosure.kschedule_bytes (go queue) encoder ~off ~len v
    | Some (TValue (Bigstring ({ off; len }, v)), queue) ->
        Enclosure.kschedule_bigstring (go queue) encoder ~off ~len v
    | Some (TBreak len, queue) ->
        Enclosure.kschedule_string (go queue) encoder ~len (String.make len ' ')
    | Some (TBox _, queue) | Some (TClose, queue) -> go queue encoder
    | None ->
        let k encoder = k { t with encoder; queue = Queue.empty } in
        let k encoder = Enclosure.flush k encoder in
        Enclosure.kschedule_string k encoder t.new_line
  in
  Enclosure.kschedule_string (go t.queue) t.encoder (String.make t.indent ' ')

let merge_indents k t =
  let indent_by_box =
    Stack.fold
      (fun a -> function `Box | `Root -> a | `Indent n -> a + n)
      0 t.boxes
  in
  let indent_by_break =
    let k a l =
      List.fold_left
        (fun a -> function `Indent n -> (max : int -> int -> int) a n)
        a l
    in
    Stack.fold k 0 t.breaks
  in
  k
    {
      t with
      indent = indent_by_box + indent_by_break;
      inner = Stack.map (fun _ -> []) t.boxes;
      breaks = Stack.map (fun _ -> []) t.boxes;
    }

let is_in_box t =
  let exception In in
  try
    Queue.rev_iter (function TBox _ -> raise In | _ -> ()) t.queue;
    false
  with In -> true

let without_last_box q =
  let rec go q =
    match Queue.tail q with
    | Some (q, TBox _) -> q
    | Some (q, _) -> go q
    | None -> Queue.empty
  in
  go q

let unroll_last_box_and_emit k value t =
  let rec unroll acc q =
    let q, x = Queue.tail_exn q in
    match x with TBox _ -> (List.rev (x :: acc), q) | x -> unroll (x :: acc) q
  in
  let box, q = unroll [ TValue value ] t.queue in
  let k t =
    let rec roll acc q =
      match acc with [] -> q | x :: r -> roll r (Queue.push q x)
    in
    let q = roll (List.rev box) t.queue in
    k { t with queue = q }
  in
  emit_line (merge_indents k) { t with queue = q }

let rec kpush_breakable_value ~current_length_of_line k value t =
  if current_length_of_line >= t.margin then
    emit_line (merge_indents (kpush k (Breakable value))) t
  else if current_length_of_line + length_of_token (TValue value) > t.margin
  then
    if is_in_box t && not (Queue.is_empty (without_last_box t.queue)) then
      unroll_last_box_and_emit k value t
    else
      let len = t.margin - current_length_of_line in
      let value0, value1 = split_value len value in
      let token0 = TValue value0 in
      let token1 = v ~breakable:true value1 in
      emit_line
        (merge_indents (kpush k token1))
        { t with queue = Queue.push t.queue token0 }
  else
    k
      {
        t with
        queue = Queue.push t.queue (TValue value);
        inner =
          Stack.swap_exn (fun lenv -> length_of_value value :: lenv) t.inner;
      }

and kpush_unbreakable_value ~current_length_of_line k value t =
  if
    current_length_of_line + length_of_token (TValue value) > t.margin
    && not (Queue.is_empty t.queue)
  then
    if is_in_box t && not (Queue.is_empty (without_last_box t.queue)) then
      unroll_last_box_and_emit k value t
    else
      match Queue.tail t.queue with
      | Some (queue, TBreak _) ->
          emit_line
            (merge_indents (kpush k (Unbreakable value)))
            { t with queue }
      | Some _ -> emit_line (merge_indents (kpush k (Unbreakable value))) t
      | None -> emit_line (merge_indents (kpush k (Unbreakable value))) t
  else
    k
      {
        t with
        queue = Queue.push t.queue (TValue value);
        inner =
          Stack.swap_exn (fun lenv -> length_of_value value :: lenv) t.inner;
      }

and kpush k value t =
  let current_length_of_line = current_length_of_line t in

  let append stack len =
    match Stack.pop stack with
    | Some (lenv, stack) -> Stack.push (len :: lenv) stack
    | None -> assert false
  in

  match value with
  | New_line -> emit_line (merge_indents k) t
  | Open Box ->
      k
        {
          t with
          boxes = Stack.push `Box t.boxes;
          inner = Stack.push [] t.inner;
          breaks = Stack.push [] t.breaks;
          queue = Queue.push t.queue (TBox `Box);
        }
  | Open (TBox indent) ->
      k
        {
          t with
          boxes = Stack.push (`Indent indent) t.boxes;
          inner = Stack.push [] t.inner;
          breaks = Stack.push [] t.breaks;
          queue = Queue.push t.queue (TBox (`Indent indent));
        }
  | Open BBox ->
      let indent =
        Stack.fold (flip (( + ) <.> List.fold_left ( + ) 0)) 0 t.inner
      in
      k
        {
          t with
          boxes = Stack.push (`Indent indent) t.boxes;
          inner = Stack.push [] t.inner;
          breaks = Stack.push [] t.breaks;
          queue = Queue.push t.queue (TBox (`Indent indent));
        }
  | Close ->
      k
        {
          t with
          boxes = Stack.tail_exn t.boxes;
          inner = Stack.tail_exn t.inner;
          breaks = Stack.tail_exn t.breaks;
          queue = Queue.push t.queue TClose;
        }
  | Breakable value -> kpush_breakable_value ~current_length_of_line k value t
  | Unbreakable value ->
      kpush_unbreakable_value ~current_length_of_line k value t
  | Break { len; indent } as break -> (
      match
        let open Option in
        Queue.tail t.queue >>= merge_breaks break
      with
      | Some (queue, len) ->
          if current_length_of_line + length_of_token (TBreak len) >= t.margin
          then
            emit_line (merge_indents k)
              { t with queue; breaks = append t.breaks (`Indent indent) }
          else
            k
              {
                t with
                queue = Queue.push queue (TBreak len);
                inner = append t.inner len;
                breaks = append t.breaks (`Indent indent);
              }
      | None ->
          if current_length_of_line + length_of_token (TBreak len) > t.margin
          then
            emit_line (merge_indents k)
              { t with breaks = append t.breaks (`Indent indent) }
          else
            k
              {
                t with
                queue = Queue.push t.queue (TBreak len);
                inner = append t.inner len;
                breaks = append t.breaks (`Indent indent);
              })

let kflush k t =
  let rec go queue encoder =
    match Queue.pop queue with
    | Some (TValue (String ({ off; len }, v)), queue) ->
        Enclosure.kschedule_string (go queue) encoder ~off ~len v
    | Some (TValue (Bytes ({ off; len }, v)), queue) ->
        Enclosure.kschedule_bytes (go queue) encoder ~off ~len v
    | Some (TValue (Bigstring ({ off; len }, v)), queue) ->
        Enclosure.kschedule_bigstring (go queue) encoder ~off ~len v
    | Some (TBreak len, queue) ->
        Enclosure.kschedule_string (go queue) encoder ~len (String.make len ' ')
    | Some (TBox _, queue) | Some (TClose, queue) -> go queue encoder
    | None ->
        Enclosure.flush
          (fun encoder -> k { t with encoder; queue = Queue.empty })
          encoder
  in
  go t.queue t.encoder

external identity : 'a -> 'a = "%identity"

let push = kpush identity
let flush = kflush identity
let is_empty t = Queue.is_empty t.queue && Enclosure.is_empty t.encoder

let create ?(margin = 998) ?(new_line = "\r\n") ~emitter len =
  let encoder = Enclosure.create ~emitter len in
  {
    encoder;
    queue = Queue.empty;
    boxes = Stack.push `Root Stack.empty;
    breaks = Stack.push [] Stack.empty;
    inner = Stack.push [] Stack.empty;
    indent = 0;
    margin;
    new_line;
  }
