open Util

(* AST nodes *)

type ast = block

and block = stmt list

and stmt =
  | Return of value option
  | Comment of string
  | Expr_stmt of value

and value =
  | Fun of value list * block * bool
  | Call of value * value list
  | Object of (string * value) list
  | Array of value list
  | Index of value * value
  | Range of value option * value option * bool
  | Try of block * block * block
  | Throw of value
  | New of value * value list
  | Class of string option * value option * class_body
  | Binop of binop * value * value
  | Monop of monop * value
  | Chainop of value * (binop * value) list
  | If of value * block * block
  | For of value * value option * binop * value * block
  | While of value * block
  | Switch of value * (value list * block) list * block
  | Assign of value * value
  | Property of value * string

  | Regex of string * string
  | Number of float
  | String of string
  | Var of string
  | Js of string
  | This
  | Null
  | Undefined
  | True
  | False

  | Param of string
  | Default of value * value
  | Splat of value option

and class_body = class_stmt list

and class_stmt =
  | Unnamed_expr of value
  | Prototype_property of string * value

and binop =
  | Eq | Ne | Gt | Ge | Lt | Le
  | And | Or
  | In | Of
  | Plus | Minus | Times | Div | Mod
  | Div_trunc | Pow | Mod_math

and monop =
  | Not | Exists
  | Neg | Pos

(* Stringification *)

let rec string_of_ast tree = string_of_block tree ^ "\n"

and string_of_block block = cat "\n" string_of_stmt block

and string_of_stmt = function
  | Return None -> "return"
  | Return (Some value) -> "return " ^ string_of_value value
  | Comment comment -> prefix_lines "# " comment
  | Expr_stmt value -> string_of_value value

and string_of_value = function
  | Fun (params, body, is_bound) ->
    let params =
      match params with
      | [] -> ""
      | _ -> "(" ^ cat ", " string_of_value params ^ ") "
    in
    let glyph = if is_bound then "=>" else "->" in
    params ^ glyph ^ inline_block "" body

  | Call (fn, args) ->
    string_of_value fn ^ "(" ^ cat ", " string_of_value args ^ ")"

  | Object key_value_pairs ->
    let string_of_pair (key, value) = key ^ string_of_value value in
    "{" ^ cat ", " string_of_pair key_value_pairs ^ "}"

  | Array elements ->
    "[" ^ cat ", " string_of_value elements ^ "]"

  | Index (arr, (Range _ as range)) ->
    string_of_value arr ^ string_of_value range
  | Index (arr, idx) ->
    string_of_value arr ^ "[" ^ string_of_value idx ^ "]"

  | Range (lbnd, ubnd, inclusive) ->
    let splat = if inclusive then ".." else "..." in
    let s = function None -> "" | Some v -> string_of_value v in
    "[" ^ s lbnd ^ splat ^ s ubnd ^ "]"

  | Try (body, catch, finally) ->
    let s pre = function [] -> "" | b -> pre ^ indent_block b in
    "try" ^ indent_block body ^ "\n" ^
    s "catch" catch ^ "\n" ^
    s "finally" finally

  | Throw value ->
    "throw " ^ string_of_value value

  | New (cls, args) ->
    "new " ^ string_of_value cls ^ "(" ^ cat ", " string_of_value args ^ ")"

  | Class (name, parent, body) ->
    let cls =
      match name with
      | None -> "class"
      | Some name -> "class " ^ name
    in
    let extends =
      match parent with
      | None -> ""
      | Some parent -> " extends " ^ string_of_value parent
    in
    let body =
      let string_of_class_stmt = function
        | Unnamed_expr value ->

          string_of_value value
        | Prototype_property (name, value) ->
          name ^ ": " ^ string_of_value value
      in
      match body with
      | [] -> failwith "empty body"
      | body -> "\n" ^ indent (cat "\n" string_of_class_stmt body)
    in
    cls ^ extends ^ body

  | Binop (op, left, right) ->
    "(" ^ string_of_value left ^
    " " ^ string_of_binop op ^ " " ^
    string_of_value right ^ ")"

  | Monop (Exists, opnd) ->
    string_of_value opnd ^ "?"
  | Monop (op, opnd) ->
    let op =
      match op with
      | Not -> "not"
      | Neg -> "-"
      | Pos -> "+"
      | _ -> failwith "this cannot happen"
    in
    op ^ string_of_value opnd

  | Chainop (hd, tl) ->
    let rec s = function
      | [] ->
        failwith "invalid chained operator"
      | [(op, opnd)] ->
        string_of_binop op ^ " " ^ string_of_value opnd
      | (op, opnd) :: tl ->
        string_of_binop op ^ " " ^ string_of_value opnd ^ " " ^ s tl
    in
    string_of_value hd ^ " " ^ s tl

  | If (cond, [], _) ->
    failwith "empty if-body"
  (*
  | If (cond, [if_expr], []) ->
    "if " ^ string_of_value cond ^ " then " ^ string_of_value if_expr
  | If (cond, [if_expr], [else_expr]) ->
    "if " ^ string_of_value cond ^ " then " ^ string_of_value if_expr ^
    " else " ^ string_of_value else_expr
  *)
  | If (cond, if_body, []) ->
    "if " ^ string_of_value cond ^ indent_block if_body
  | If (cond, if_body, else_body) ->
    "if " ^ string_of_value cond ^ indent_block if_body ^
    "\nelse" ^ indent_block else_body

  | For (value, index, op, arr, body) ->
    let preamble =
      let iterator =
        let value = string_of_value value in
        match index with
        | None -> value
        | Some index -> value ^ ", " ^ string_of_value index
      in
      "for " ^ iterator ^ " " ^ string_of_binop op ^ " " ^ string_of_value arr
    in
    begin
      match body with
      | [] -> failwith "empty loop body"
      | [stmt] -> string_of_stmt stmt ^ preamble
      | _ -> preamble ^ indent_block body
    end

  | While (cond, body) ->
    "while " ^ string_of_value cond ^ indent_block body

  | Switch (value, when_cases, default_case) ->
    let string_of_when (match_values, body) =
      "\nwhen " ^ cat ", " string_of_value match_values ^
      inline_block "then" body
    in
    "switch " ^ string_of_value value ^
    indent (cat "" string_of_when when_cases) ^
    inline_block "else" default_case

  | Assign (left, right) ->
    string_of_value left ^ " = " ^ string_of_value right

  | Property (obj, prop) ->
    string_of_value obj ^ "." ^ prop

  | Regex (pattern, modifiers) ->
    "/" ^ pattern ^ "/" ^ modifiers

  | Number n   -> string_of_number n
  | String s   -> "\"" ^ s ^ "\""
  | Var name   -> name
  | Js js      -> "`" ^ js ^ "`"
  | This       -> "this"
  | Null       -> "null"
  | Undefined  -> "undefined"
  | True       -> "true"
  | False      -> "false"
  | Param name -> name

  | Default (Param name, value) ->
    name ^ "=" ^ string_of_value value
  | Default _ ->
    failwith "Default can only have Param as first child"

  | Splat None ->
    "..."
  | Splat (Some value) ->
    string_of_value value ^ "..."

and string_of_number n = string_of_float n  (* FIXME *)

and string_of_binop = function
  | Eq        -> "is"
  | Ne        -> "isnt"
  | Gt        -> ">"
  | Ge        -> ">="
  | Lt        -> "<"
  | Le        -> "<="
  | And       -> "and"
  | Or        -> "or"
  | In        -> "in"
  | Of        -> "of"
  | Plus      -> "+"
  | Minus     -> "-"
  | Times     -> "*"
  | Div       -> "/"
  | Mod       -> "%"
  | Div_trunc -> "//"
  | Pow       -> "**"
  | Mod_math  -> "%%"

and indent_block = function
  | [] -> failwith "empty block"
  | stmts -> "\n" ^ indent (string_of_block stmts)

and inline_block prefix = function
  | [] -> ""
  | [stmt] -> prefix ^ " " ^ string_of_stmt stmt
  | stmts -> indent_block stmts