(** [Form.t] to paren text, for [flan convert] of a [.fln] file: the other
    direction of [Indent_printer]. It keeps what [Form.pretty] cannot — the
    source's number spellings and comments, through [Source_text] — and lays
    a form out the way the corpus is written: flat when it fits, otherwise
    the head and the arguments that name the form on the first line and the
    rest one per line, two columns in. *)

let width = 80

(* The reader's prefix for a quoting form, as the corpus writes them:
   [`(do ~x ~@xs)] rather than [(quasiquote (do (unquote x) ...))]. *)
let sugar (f : Form.t) =
  match f.v with
  | Form.List [ { v = Form.Sym s; _ }; x ] ->
    (match s, x.v with
     | "quote", _ -> Some ("'", x)
     | "quasiquote", _ -> Some ("`", x)
     | "unquote-splicing", _ -> Some ("~@", x)
     (* [~@x] would read as a splice. *)
     | "unquote", Form.Sym n when String.length n > 0 && n.[0] = '@' -> None
     | "unquote", _ -> Some ("~", x)
     | _ -> None)
  | _ -> None

let rec flat spell (f : Form.t) =
  let seq l = String.concat " " (List.map (flat spell) l) in
  match f.v with
  | _ when sugar f <> None ->
    let p, x = Option.get (sugar f) in
    p ^ flat spell x
  | Form.Int _ | Form.Float _ ->
    (match spell f with Some t -> t | None -> Form.to_source f)
  | Form.List l -> "(" ^ seq l ^ ")"
  | Form.Vec l -> "[" ^ seq l ^ "]"
  | Form.Map l -> "{" ^ seq l ^ "}"
  | _ -> Form.to_source f

(* Heads whose arguments are statements or clauses rather than values: these
   break one argument to a line, never filled. *)
let statement_heads =
  [ "let"; "loop"; "set"; "if"; "when"; "unless"; "cond"; "while"; "until";
    "dotimes"; "match"; "handler-case"; "handler-bind"; "restart-case";
    "return"; "defer"; "do"; "fn"; "with-allocator"; "comment"; "quasiquote";
    "break"; "continue" ]

let form_heads =
  [ "defn"; "defn-"; "defmacro"; "def"; "defonce"; "defconst"; "defstruct";
    "defunion"; "defdata"; "defenum"; "import"; "defalias"; "defmethod";
    "defgeneric"; "defmulti"; "defclass" ]

(* How many arguments stay on the head's line when the form is broken. *)
let kept head =
  match head with
  | "defn" | "defn-" | "defmethod" -> 3
  | "defmacro" | "def" | "defonce" | "defconst" | "defstruct" | "defunion"
  | "defdata" | "defenum" | "import" | "defalias" -> 2
  | "do" | "cond" | "comment" | "restart-case" | "handler-case" -> 0
  | _ -> 1

(* Two at a time: [a b] on one line when it fits and no comment sits
   inside either, otherwise each laid out on its own lines. *)
let paired ~inside spell col placed items =
  let rec go = function
    | (a : Form.t) :: b :: more ->
      let line = flat spell a ^ " " ^ flat spell b in
      (* A comment on a line between the two would land beside the wrong
         one, so a pair more than a line apart keeps two lines. *)
      (if col + String.length line <= width && not (inside a) && not (inside b)
          && b.loc.Loc.line <= a.loc.Loc.eline + 1
       then [ Source_text.tag a.loc.Loc.line
                (Source_text.tag b.loc.Loc.line (String.make col ' ' ^ line)) ]
       else placed col a @ placed col b)
      @ go more
    | rest -> List.concat_map (placed col) rest
  in
  go items

(* [inside l] says whether a comment sits on a line of [f] before its last,
   where a flat [f] would leave it nowhere to go: such a form is broken. *)
let rec layout ?(inside = fun _ -> false) spell col (f : Form.t) : string list =
  let layout = layout ~inside in
  let one = flat spell f in
  let tagl (x : Form.t) = function
    | first :: rest -> Source_text.tag x.loc.Loc.line first :: rest
    | [] -> []
  in
  if col + String.length one <= width && not (inside f) then [ one ]
  else
    match sugar f with
    | Some (p, x) ->
      (match layout spell (col + String.length p) x with
       | first :: more ->
         let tags, body = Source_text.untag first in
         List.fold_left (fun l t -> Source_text.tag t l) (p ^ body) tags :: more
       | [] -> [ one ])
    | None ->
    (* A plain call broken across lines: its arguments fill each line,
       aligned under the first, and one that needs lines of its own gets
       them. *)
    let fill items =
      match items with
      | hd :: (_ :: _ as args) ->
        let lead = "(" ^ flat spell hd ^ " " in
        let acol = col + String.length lead in
        let pad = String.make acol ' ' in
        let vis l = String.length (snd (Source_text.untag l)) in
        (* [cur] is the line being filled, [fresh] while it holds no argument;
           the first line is built with its column and cut back after. *)
        let rec go cur fresh acc = function
          | [] -> List.rev (if fresh then acc else cur :: acc)
          | (a : Form.t) :: more ->
            let t = flat spell a in
            let sep = if fresh then "" else " " in
            let close = if more = [] then 1 else 0 in
            if (not (inside a)) && vis cur + String.length sep + String.length t + close <= width
            then go (Source_text.tag a.loc.Loc.line (cur ^ sep ^ t)) false acc more
            else if fresh then
              match layout spell acol a with
              | l1 :: ls ->
                let tags, b = Source_text.untag l1 in
                let l1 = List.fold_left (fun l t -> Source_text.tag t l) (cur ^ b) tags in
                let l1 = Source_text.tag a.loc.Loc.line l1 in
                go pad true (List.rev_append (l1 :: ls) acc) more
              | [] -> go cur fresh acc more
            else go pad true (cur :: acc) (a :: more)
        in
        let lines = go (String.make col ' ' ^ lead) true [] args in
        let lines =
          match lines with
          | l1 :: rest ->
            let tags, b = Source_text.untag l1 in
            List.fold_left (fun l t -> Source_text.tag t l) (String.sub b col (String.length b - col)) tags
            :: rest
          | [] -> []
        in
        let n = List.length lines in
        List.mapi (fun i l -> if i = n - 1 then l ^ ")" else l) lines
      | _ -> [ one ]
    in
    let bracket ?(pairs = false) o c items ~keep =
      let placed inner x =
        match layout spell inner x with
        | first :: more ->
          (* The pad goes after any tags, which lead the line. *)
          let tags, body = Source_text.untag first in
          let padded =
            List.fold_left (fun l t -> Source_text.tag t l) (String.make inner ' ' ^ body) tags
          in
          tagl x (padded :: more)
        | [] -> []
      in
      let lines =
        let head_len k =
          col + 1 + String.length
                      (String.concat " " (List.map (flat spell) (List.filteri (fun i _ -> i < k) items)))
        in
        let hang =
          (* [(Rule {.label "a"\n       .applies f})]: a literal the head
             names hangs on the head's line rather than dropping below it. *)
          keep > 1
          && (match List.nth_opt items (keep - 1) with
              | Some { Form.v = Form.Map _; _ } -> head_len (keep - 1) + 1 < width - 20
              | _ -> false)
        in
        let keep = if keep > 1 && head_len keep > width && not hang then 1 else keep in
        if hang && head_len keep > width then
          let first = List.filteri (fun i _ -> i < keep - 1) items in
          let lit = List.nth items (keep - 1) in
          let rest = List.filteri (fun i _ -> i >= keep) items in
          let lead = o ^ String.concat " " (List.map (flat spell) first) ^ " " in
          (match layout spell (col + String.length lead) lit with
           | l1 :: more ->
             let tags, b = Source_text.untag l1 in
             List.fold_left (fun l t -> Source_text.tag t l) (lead ^ b) tags :: more
           | [] -> [ lead ])
          @ List.concat_map (placed (col + 2)) rest
        else if keep > 0 then
          let first = List.filteri (fun i _ -> i < keep) items in
          let rest = List.filteri (fun i _ -> i >= keep) items in
          let before = List.filteri (fun i _ -> i < keep - 1) first in
          match List.nth first (keep - 1) with
          (* A binding vector with a comment inside it — [(let [a 1 ; first
             b 2] ...)] — goes a pair to a line, each line carrying its pair's
             source line, so a comment after a binding stays on it. *)
          | { Form.v = Form.Vec vs; _ } as v
            when keep > 1 && inside v && List.length vs mod 2 = 0 ->
            let lead = o ^ String.concat " " (List.map (flat spell) before) ^ " [" in
            let pad = String.make (col + String.length lead) ' ' in
            let rec pairs k = function
              | (a : Form.t) :: b :: more ->
                Source_text.tag a.loc.Loc.line
                  ((if k = 0 then lead else pad) ^ flat spell a ^ " " ^ flat spell b)
                :: pairs (k + 1) more
              | _ -> []
            in
            let ps = pairs 0 vs in
            let np = List.length ps in
            List.mapi (fun i l -> if i = np - 1 then l ^ "]" else l) ps
            @ List.concat_map (placed (col + 2)) rest
          | _ ->
            (o ^ String.concat " " (List.map (flat spell) first))
            :: (if pairs then paired ~inside spell (col + 2) placed rest
                else List.concat_map (placed (col + 2)) rest)
        else
          match items with
          | [] -> [ o ]
          | _ when pairs ->
            (match paired ~inside spell (col + 1) placed items with
             | first :: more ->
               let tags, body = Source_text.untag first in
               let body = String.sub body (col + 1) (String.length body - col - 1) in
               List.fold_left (fun l t -> Source_text.tag t l) (o ^ body) tags :: more
             | [] -> [ o ])
          | x :: xs ->
            (match layout spell (col + 1) x with
             | first :: more ->
               let tags, body = Source_text.untag first in
               List.fold_left (fun l t -> Source_text.tag t l) (o ^ body) tags :: more
             | [] -> [ o ])
            @ List.concat_map (placed (col + 1)) xs
      in
      let n = List.length lines in
      List.mapi (fun i l -> if i = n - 1 then l ^ c else l) lines
    in
    match f.v with
    (* A let's binding vector opens on the head's line, a pair to a line,
       each value hanging after its name, as the corpus writes them. *)
    | Form.List (({ v = Form.Sym (("let" | "loop") as h); _ } as hd)
                 :: ({ v = Form.Vec bs; _ } as v) :: body)
      when bs <> [] && List.length bs mod 2 = 0 && not (inside v) ->
      let lead = "(" ^ h ^ " [" in
      let vcol = col + String.length lead in
      let rec pairs k = function
        | (a : Form.t) :: b :: more ->
          let an = flat spell a in
          let pad = if k = 0 then lead else String.make vcol ' ' in
          let lines =
            match layout spell (vcol + String.length an + 1) b with
            | first :: rest ->
              let tags, bt = Source_text.untag first in
              List.fold_left (fun l t -> Source_text.tag t l)
                (Source_text.tag a.loc.Loc.line (pad ^ an ^ " " ^ bt)) tags
              :: rest
            | [] -> [ pad ^ an ]
          in
          lines @ pairs (k + 1) more
        | _ -> []
      in
      let flat_v = lead ^ String.concat " " (List.map (flat spell) bs) in
      let bl =
        if col + String.length flat_v + 1 <= width then [ flat_v ] else pairs 0 bs
      in
      let nb = List.length bl in
      let bl = List.mapi (fun i l -> if i = nb - 1 then l ^ "]" else l) bl in
      let placed x =
        match layout spell (col + 2) x with
        | first :: more ->
          let tags, b = Source_text.untag first in
          Source_text.tag x.loc.Loc.line
            (List.fold_left (fun l t -> Source_text.tag t l) (String.make (col + 2) ' ' ^ b) tags)
          :: more
        | [] -> []
      in
      let lines =
        (match bl with
         | first :: rest -> Source_text.tag hd.loc.Loc.line first :: rest
         | [] -> [])
        @ List.concat_map placed body
      in
      let n = List.length lines in
      List.mapi (fun i l -> if i = n - 1 then l ^ ")" else l) lines
    (* A match's arms and a cond's clauses go a pair to a line when the pair
       fits, the way the corpus writes them. *)
    | Form.List (({ v = Form.Sym "match"; _ }) :: _ :: rest as items)
      when List.length rest mod 2 = 0 ->
      bracket ~pairs:true "(" ")" items ~keep:2
    | Form.List (({ v = Form.Sym "cond"; _ }) :: rest as items)
      when List.length rest mod 2 = 0 ->
      bracket ~pairs:true "(" ")" items ~keep:1
    | Form.List (({ v = Form.Sym h; _ }) :: rest as items) ->
      (* A loop's label stays with its test: [(while :outer (< i n)]. *)
      let label =
        match h, rest with
        | ("while" | "until" | "dotimes"), { v = Form.Kw _; _ } :: _ :: _ -> 1
        | _ -> 0
      in
      let stmt (a : Form.t) =
        match a.v with
        | Form.List ({ v = Form.Sym x; _ } :: _) -> List.mem x statement_heads
        | _ -> false
      in
      let n = List.length rest in
      if List.mem h statement_heads || List.mem h form_heads || List.exists stmt rest
         || n < 2 || h.[0] = '.'
      then
        (* A body: the arguments before its first list stay on the head's
           line, [(repeat i 2\n  (set ...) ...)]. *)
        let lead =
          if List.mem h statement_heads || List.mem h form_heads then 0
          else
            let rec go k = function
              | { Form.v = Form.List _; _ } :: _ | [] -> k
              | _ :: more -> go (k + 1) more
            in
            go 0 rest
        in
        bracket "(" ")" items
          ~keep:(1 + label + max (min lead (n - 1)) (min (kept h) (n - label)))
      else fill items
    | Form.List items -> bracket "(" ")" items ~keep:0
    | Form.Vec items -> bracket "[" "]" items ~keep:0
    | Form.Map items when List.length items mod 2 = 0 ->
      bracket ~pairs:true "{" "}" items ~keep:0
    | Form.Map items -> bracket "{" "}" items ~keep:0
    | _ -> [ one ]

(** A whole file, with [source]'s comments and spellings when given. *)
let program ?source (fs : Form.t list) : string =
  let spell =
    match source with Some src -> Source_text.spelling src | None -> fun _ -> None
  in
  let cs = match source with Some src -> Source_text.comments src | None -> [] in
  let inside (f : Form.t) =
    List.exists
      (fun (c : Source_text.comment) ->
         f.loc.Loc.line <= c.line && c.line < f.loc.Loc.eline)
      cs
  in
  let text =
    Source_text.join_top
      (List.map
         (fun (f : Form.t) ->
            (f, String.concat "\n"
               (match layout ~inside spell 0 f with
                | first :: rest -> Source_text.tag f.loc.Loc.line first :: rest
                | [] -> [])))
         fs)
    ^ "\n"
  in
  Source_text.weave ~starts:(Source_text.form_starts fs) cs text