(** What a printer needs from the text a program was read from and a
    [Form.t] does not carry: the comments, and the spelling of each number.
    [flan convert] reads both here and puts them back (author's decision 83),
    so a converted file keeps its [;] notes and its [0xFFF00FFF]s.

    Both syntaxes share the lexical rules this depends on: a comment runs
    from [;] to the end of its line, a string is ["..."] with backslash
    escapes, and [\c] is a character — so [\;] is not a comment. *)

type comment = {
  line : int;        (* 1-based *)
  text : string;     (* from the [;] to the end of the line *)
  own_line : bool;   (* nothing but spaces before it on its line *)
  gap_after : bool;  (* the line after it is blank *)
}

let comments (src : string) : comment list =
  let n = String.length src in
  let out = ref [] in
  let line = ref 1 and line_start = ref 0 in
  let i = ref 0 in
  while !i < n do
    (match src.[!i] with
     | '\n' -> incr line; line_start := !i + 1; incr i
     | '"' ->
       incr i;
       while !i < n && src.[!i] <> '"' do
         if src.[!i] = '\\' then incr i;
         if !i < n && src.[!i] = '\n' then (incr line; line_start := !i + 1);
         incr i
       done;
       incr i
     | '\\' -> i := !i + 2
     | ';' ->
       let j = ref !i in
       while !j < n && src.[!j] <> '\n' do incr j done;
       let text = String.sub src !i (!j - !i) in
       let text =
         if text <> "" && text.[String.length text - 1] = '\r'
         then String.sub text 0 (String.length text - 1) else text
       in
       let before = String.sub src !line_start (!i - !line_start) in
       let k = ref (!j + 1) in
       while !k < n && (src.[!k] = ' ' || src.[!k] = '\r') do incr k done;
       let gap_after = !j < n && (!k >= n || src.[!k] = '\n') in
       out := { line = !line; text; own_line = String.trim before = ""; gap_after } :: !out;
       i := !j
     | _ -> incr i)
  done;
  List.rev !out

(** A number's text as written, when it reads back to the same value: the
    text under its span. [Form.Int] keeps only the value. *)
let spelling (src : string) : Form.t -> string option =
  let lines = Array.of_list (String.split_on_char '\n' src) in
  fun (f : Form.t) ->
    let l = f.loc in
    if l.Loc.line < 1 || l.Loc.line > Array.length lines || l.Loc.eline <> l.Loc.line
    then None
    else
      let text = lines.(l.Loc.line - 1) in
      let a = l.Loc.col - 1 and b = l.Loc.ecol - 1 in
      if a < 0 || b > String.length text || b <= a then None
      else
        let t = String.sub text a (b - a) in
        match f.v with
        | Form.Int i when Int64.of_string_opt t = Some i -> Some t
        | Form.Float x
          when String.exists (fun c -> c = '.' || c = 'e' || c = 'E') t
               && (match float_of_string_opt t with
                   | Some y -> Int64.equal (Int64.bits_of_float x) (Int64.bits_of_float y)
                   | None -> false) ->
          Some t
        | _ -> None

(* ── Lines tagged with where they came from ───────────────────────────

   A printer marks the first line of each form it lays out with the source
   line that form started on. [weave] reads the marks back out and uses them
   to put each comment where it was: an own-line comment above the first
   line that came from after it, a trailing comment at the end of the line
   the code it followed was printed on. *)

let tag (line : int) (text : string) =
  (* A tag already there is a form nested at the start of this one's first
     line; the outer form started no later, so it wins. *)
  let text =
    if String.length text > 0 && text.[0] = '\001' then
      match String.index_opt text '\002' with
      | Some k -> String.sub text (k + 1) (String.length text - k - 1)
      | None -> text
    else text
  in
  if line <= 0 then text
  else "\001" ^ string_of_int line ^ "\002" ^ text

let untag (text : string) : int option * string =
  if String.length text > 0 && text.[0] = '\001' then
    match String.index_opt text '\002' with
    | Some k ->
      (int_of_string_opt (String.sub text 1 (k - 1)),
       String.sub text (k + 1) (String.length text - k - 1))
    | None -> (None, text)
  else (None, text)

let indent_of s =
  let n = String.length s in
  let rec go i = if i < n && s.[i] = ' ' then go (i + 1) else i in
  go 0

let weave (cs : comment list) (text : string) : string =
  let lines = Array.of_list (List.map untag (String.split_on_char '\n' text)) in
  let n = Array.length lines in
  let tags = Array.map fst lines and body = Array.map snd lines in
  let before = Array.make (n + 1) [] and trailing = Array.make n [] in
  List.iter
    (fun c ->
       if c.own_line then begin
         (* The first line printed from code after the comment. *)
         let rec find i =
           if i >= n then n
           else match tags.(i) with Some t when t > c.line -> i | _ -> find (i + 1)
         in
         let i = find 0 in
         before.(i) <- c :: before.(i)
       end
       else begin
         (* The line the code before it went to: the latest line tagged at
            or before the comment's own line. *)
         let best = ref (-1) and best_tag = ref 0 in
         Array.iteri
           (fun i t ->
              match t with
              | Some t when t <= c.line && t >= !best_tag -> best := i; best_tag := t
              | _ -> ())
           tags;
         if !best < 0 then before.(0) <- c :: before.(0)
         else trailing.(!best) <- c :: trailing.(!best)
       end)
    cs;
  let b = Buffer.create (String.length text + 256) in
  let emit s = Buffer.add_string b s; Buffer.add_char b '\n' in
  for i = 0 to n do
    let ind =
      if i < n then indent_of body.(i)
      else 0
    in
    List.iter
      (fun c ->
         emit (String.make ind ' ' ^ c.text);
         (* A comment set apart from what follows it, at the top level, stays
            set apart: a file's header, a section rule. *)
         if c.gap_after && ind = 0 && i < n then emit "")
      (List.rev before.(i));
    if i < n then begin
      match List.rev trailing.(i) with
      | [] -> emit body.(i)
      | first :: more ->
        emit (body.(i) ^ "  " ^ first.text);
        (* A second trailing comment for the same printed line goes on its
           own line under it, which reads the same and keeps both. *)
        let ind = if i + 1 < n then indent_of body.(i + 1) else indent_of body.(i) in
        List.iter (fun c -> emit (String.make ind ' ' ^ c.text)) more
    end
  done;
  (* [text] ended in a newline, which split into a last empty line. *)
  let s = Buffer.contents b in
  let rec trim s =
    let k = String.length s in
    if k >= 2 && s.[k - 1] = '\n' && s.[k - 2] = '\n' then trim (String.sub s 0 (k - 1))
    else s
  in
  trim s