open Core
open Types
open PPrint

let rec escape_chlist = function
  | [] -> []
  | '$' :: '{' :: l' -> '\\' :: '$' :: '{' :: escape_chlist l'
  | '\n' :: l' -> '\\' :: 'n' :: escape_chlist l'
  | '\r' :: l' -> '\\' :: 'r' :: escape_chlist l'
  | '\t' :: l' -> '\\' :: 't' :: escape_chlist l'
  | '\\' :: l' -> '\\' :: '\\' :: escape_chlist l'
  | '"' :: l' -> '\\' :: '"' :: escape_chlist l'
  | c :: l' -> c :: escape_chlist l'

let escape_string s = s |> String.to_list |> escape_chlist |> String.of_list
let out_width = ref 80
let set_width i = out_width := i
let indent = ref 2
let set_indent i = indent := i

let rec doc_of_expr = function
  | BinaryOp (op, lhs, rhs) ->
      let lhs_doc = maybe_parens_bop op `Left lhs
      and rhs_doc = maybe_parens_bop op `Right rhs in
      infix !indent 1 (doc_of_bop op) lhs_doc rhs_doc
  | UnaryOp (op, e) -> precede (doc_of_uop op) (maybe_parens (prec_of_uop op) e)
  | Cond (e1, e2, e3) ->
      surround !indent 1
        (soft_surround !indent 1 (string "if") (doc_of_expr e1) (string "then"))
        (doc_of_expr e2)
        (string "else" ^^ nest !indent (break 1 ^^ doc_of_expr e3))
  | With (e1, e2) ->
      flow (break 1) [ string "with"; doc_of_expr e1 ^^ semi; doc_of_expr e2 ]
  | Assert (e1, e2) ->
      flow (break 1) [ string "assert"; doc_of_expr e1 ^^ semi; doc_of_expr e2 ]
  | Test (e, path) ->
      maybe_parens 4 e ^^ space ^^ string "?"
      ^^ group (break 1 ^^ separate_map dot doc_of_expr path)
  | SetLet bs ->
      surround !indent 1
        (string "let " ^^ lbrace)
        (group (separate_map (break 1) doc_of_binding bs))
        rbrace
  | Let (bs, e) ->
      surround !indent 1 (string "let")
        (separate_map (break 1) doc_of_binding bs)
        (prefix !indent 1 (string "in") (doc_of_expr e))
  | Val v -> doc_of_val v
  | Id id -> string id
  | Select (e, path, oe) ->
      maybe_parens 1 e ^^ dot ^^ doc_of_attpath path
      ^^ optional
           (fun e ->
             space ^^ string "or" ^^ nest !indent (break 1 ^^ maybe_parens 1 e))
           oe
  | Apply (e1, e2) -> prefix !indent 1 (maybe_parens 2 e1) (maybe_parens 2 e2)
  | Aquote e -> surround !indent 0 (string "${") (doc_of_expr e) (string "}")

and maybe_parens lvl e =
  if prec_of_expr e >= lvl then surround !indent 0 lparen (doc_of_expr e) rparen
  else doc_of_expr e

and maybe_parens_bop op (loc : [ `Left | `Right ]) e =
  match (loc, assoc_of_bop op) with
  | (`Left, Some Left | `Right, Some Right)
    when prec_of_expr e >= prec_of_bop op ->
      doc_of_expr e
  | _, _ -> maybe_parens (prec_of_bop op) e

and doc_of_attpath path = separate_map dot doc_of_expr path

and doc_of_paramset (params, kind) =
  let ps =
    List.map ~f:doc_of_param params
    @ if Poly.(kind = Loose) then [ string "..." ] else []
  in
  surround !indent 0 lbrace (separate (comma ^^ break 1) ps) rbrace

and doc_of_param (id, oe) =
  string id ^^ optional (fun e -> qmark ^^ space ^^ doc_of_expr e) oe

and doc_of_binding = function
  | AttrPath (path, e) ->
      doc_of_attpath path ^^ space ^^ equals ^^ space ^^ doc_of_expr e ^^ semi
  | Inherit (oe, ids) ->
      let id_docs =
        List.map
          ~f:(function
            | Id x | Val (Str (x, [])) -> string x | _ -> assert false)
          ids
      in
      let xs =
        flow (break 1)
          (match oe with
          | Some e -> parens (doc_of_expr e) :: id_docs
          | None -> id_docs)
      in
      soft_surround !indent 0 (string "inherit" ^^ space) xs semi

and doc_of_bop = function
  | Plus -> plus
  | Minus -> minus
  | Mult -> star
  | Div -> slash
  | Gt -> rangle
  | Lt -> langle
  | Lte -> string "<="
  | Gte -> string ">="
  | Eq -> string "=="
  | Neq -> string "!="
  | Or -> string "||"
  | And -> string "&&"
  | Impl -> string "->"
  | Merge -> string "//"
  | Concat -> string "++"

and doc_of_uop = function Negate -> minus | Not -> bang

and doc_of_val = function
  | Str (start, xs) ->
      dquotes
        (string (escape_string start)
        ^^ concat
             (List.map
                ~f:(fun (e, s) ->
                  surround !indent 0 (string "${") (doc_of_expr e)
                    (string "}" ^^ string (escape_string s)))
                xs))
  | IStr (i, start, xs) ->
      let qq = string "''" in
      let str s =
        String.split ~on:'\n' s
        |> List.map ~f:(fun s ->
               let len = String.length s in
               let s' =
                 if len >= i then String.sub s ~pos:i ~len:(len - i) else s
               in
               string s')
        |> separate hardline
      in
      enclose qq qq
        (str start
        ^^ concat
             (List.map
                ~f:(fun (e, s) ->
                  enclose (string "${") rbrace (doc_of_expr e) ^^ str s)
                xs))
  | Int x | Float x | Path x | SPath x | HPath x | Uri x -> string x
  | Lambda (pattern, body) ->
      let pat =
        match pattern with
        | Alias id -> string id
        | ParamSet (None, ps) -> doc_of_paramset ps
        | ParamSet (Some id, ps) ->
            doc_of_paramset ps ^^ group (break 1 ^^ at ^^ break 1 ^^ string id)
      in
      flow (break 1) [ pat ^^ colon; doc_of_expr body ]
  | List [] -> lbracket ^^ space ^^ rbracket
  | List es ->
      surround !indent 1 lbracket
        (separate_map (break 1) (maybe_parens 2) es)
        rbracket
  | AttSet (Nonrec, []) -> lbrace ^^ space ^^ rbrace
  | AttSet (Nonrec, bs) ->
      surround !indent 1 lbrace
        (group (separate_map (break 1) doc_of_binding bs))
        rbrace
  | AttSet (Rec, bs) ->
      string "rec" ^^ space ^^ doc_of_val (AttSet (Nonrec, bs))

let print chan expr = ToChannel.pretty 0.7 !out_width chan (doc_of_expr expr)

let to_string expr =
  let buf = Stdlib.Buffer.create 0 in
  ToBuffer.pretty 0.7 !out_width buf (doc_of_expr expr);
  Stdlib.Buffer.contents buf