(** 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. *)

(* A line can come from several forms — a statement and the test at its
   start — so it carries every line they started on. *)
let untag (text : string) : int list * string =
  if String.length text > 0 && text.[0] = '\001' then
    match String.index_opt text '\002' with
    | Some k ->
      (List.filter_map int_of_string_opt
         (String.split_on_char ',' (String.sub text 1 (k - 1))),
       String.sub text (k + 1) (String.length text - k - 1))
    | None -> ([], text)
  else ([], text)

let tag (line : int) (text : string) =
  let tags, body = untag text in
  let tags = if line > 0 && not (List.mem line tags) then line :: tags else tags in
  if tags = [] then body
  else "\001" ^ String.concat "," (List.map string_of_int tags) ^ "\002" ^ body

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

(** Where every form in [fs] starts, as (line, col), in source order. *)
let form_starts (fs : Form.t list) : (int * int) list =
  let out = ref [] in
  let rec walk (f : Form.t) =
    out := (f.loc.Loc.line, f.loc.Loc.col) :: !out;
    match f.v with
    | Form.List l | Form.Vec l | Form.Map l -> List.iter walk l
    | _ -> ()
  in
  List.iter walk fs;
  List.sort_uniq compare !out

let weave ?(starts = []) (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 all = Array.map fst lines and body = Array.map snd lines in
  (* For ordering, a line is as early as the earliest form on it. *)
  let tags =
    Array.map (function [] -> None | l -> Some (List.fold_left min max_int l)) all
  in
  let before = Array.make (n + 1) [] and trailing = Array.make n [] in
  List.iter
    (fun c ->
       (* The line printed from exactly [line], when one was: a printer
          may reorder forms — handler-bind's clauses go after its body — so a
          comment goes with the form, not with whatever line follows. *)
       let exact line =
         let r = ref (-1) in
         Array.iteri (fun i t -> if List.mem line t && !r < 0 then r := i) all;
         !r
       in
       if c.own_line then begin
         (* The form it precedes: the first to start after it. *)
         let owner =
           List.find_opt (fun (l, _) -> l > c.line) starts |> Option.map fst
         in
         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 =
           match owner with
           | Some l when exact l >= 0 -> exact l
           | _ -> find 0
         in
         before.(i) <- c :: before.(i)
       end
       else begin
         (* The line the code before it went to: the one printed from its own
            line, or else the latest line tagged at or before it. *)
         let last_exact =
           let r = ref (-1) in
           Array.iteri (fun i t -> if List.mem c.line t then r := i) all;
           !r
         in
         if last_exact >= 0 then trailing.(last_exact) <- c :: trailing.(last_exact)
         else begin
           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
       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