formater.lang

import Lexer from "./lexer"
import Parser from "./parser"
import fs from "fs"
import {
  IdLookup,
  NamedLet,
  NumExpr,
  FunctionCall,
  CommandExpr,
  JsOpExpr,
  FunctionDef,
  ReturnExpr,
  DataClassDef,
  NewExpr,
  DotAccess,
  ClassDef,
  ClassInstanceEntry,
  ClassGetterExpr,
  PrefixDotLookup,
  StrExpr,
  NotExpr,
  ArrayLiteral,
  IfStatement,
  NodeAssignment,
  NodePlusAssignment,
  WhileStatement,
  RegexNode,
  ContinueStatement,
  BreakStatement,
  IfBranch,
  PrefixBindLookup,
  ElseIfBranch,
  ElseBranch,
  PropertyLookup,
  ExportDefault,
  ExportStatement,
  SpreadExpr,
  SimpleArg,
  SpreadArg,
  ArrowFn,
  IsOperator,
  BoundFunctionDef,
  ForLoop,
  IsNotOperator,
  ParenExpr,
  LetObjectDeconstruction,
  RegularObjectProperty,
  RenamedProperty,
  ImportStatement,
  DefaultImport,
  LetArrDeconstruction,
  ArrNameEntry,
  ArrComma,
  DefaultObjClassArg,
  NamedClassArg,
  ObjClassArg,
  SimpleDefaultArg,
  ObjLit,
  SimpleObjEntry,
  ObjClassArgEntry,
} from "./parser"
class Formatter(ast, { indentation = 0, parents = [] })
  get padding = Array(.indentation + 1).join(" ")

  def indent() = .indentation += 2

  def dedent() = .indentation += -2

  def space_for(count, max = 5)
    let space = " "
    if count > max
      .indentation += 2
      space = "\n" + .padding
      .indentation += -2
    end
    space
  end

  def format
    let output = ""
    let i = 0
    for let node of .ast do
      output += .padding
      output += .format_node(node, i === .ast.length - 1)
      output += "\n"
      i += 1
    end
    output
  end

  get is_in_if = .parents.at(-1) === IfBranch || .parents.at(-1) === ElseIfBranch || .parents.at(-1) === ElseBranch

  def format_node(node, is_last = false)
    if node is FunctionDef
      return .format_function_def(node)
    else if node is ReturnExpr
      return .format_return_expr(node, is_last && !.is_in_if)
    else if node is NumExpr
      return .format_num_expr(node)
    else if node is NamedLet
      return .format_named_let(node)
    else if node is JsOpExpr
      return .format_js_op_expr(node)
    else if node is StrExpr
      return "\"" + node.value + "\""
    else if node is ForLoop
      return .format_for_loop(node)
    else if node is IdLookup
      return node.name
    else if node is ArrayLiteral
      return .format_array_literal(node)
    else if node is DefaultImport
      return .format_default_import(node)
    else if node is ImportStatement
      return .format_import_statement(node)
    else if node is PrefixDotLookup
      return "." + node.name
    else if node is NodePlusAssignment
      return .format_node_plus_assignment(node)
    else if node is FunctionCall
      return .format_function_call(node)
    else if node is DotAccess
      return .format_dot_access(node)
    else if node is ClassDef
      return .format_class_def(node)
    else if node is IfStatement
      return .format_if_statement(node)
    else if node is NodeAssignment
      return .format_node_assignment(node)
    else if node is IsOperator
      return .format_is_operator(node)
    else if node is CommandExpr
      return .format_command_expr(node)
    else if node is PrefixBindLookup
      return .format_prefix_bind_lookup(node)
    else if node is PropertyLookup
      return .format_property_lookup(node)
    else if node is NewExpr
      return .format_new_expr(node)
    else if node is ObjLit
      return .format_obj_lit(node)
    else if node is NotExpr
      return "!" + .format_node(node.expr)
    else if node is ArrowFn
      return .format_arrow_fn(node)
    else if node is LetArrDeconstruction
      return .format_let_arr_deconstruction(node)
    else if node is LetObjectDeconstruction
      return .format_let_object_deconstruction(node)
    else
      assert_not_reached! "Format not implemented for " + node.constructor.name
    end
  end

  def format_arrow_fn({ arg_name, return_expr })
    arg_name + " => " + .format_node(return_expr)
  end

  def format_let_arr_entry(entry)
    if entry is ArrComma
      return ""
    else if entry is ArrNameEntry
      return entry.name
    else
      assert_not_reached! "Let arr unknown " + entry.constructor.name
    end
  end

  def format_let_arr_deconstruction({ entries, rhs })
    "let [" + entries.map(::format_let_arr_entry).join(", ") + "] = " + .format_node(rhs)
  end

  def format_let_object_entry(entry)
    if entry is RegularObjectProperty
      return entry.name
    else if entry is RenamedProperty
      return entry.old_name + ": " + entry.new_name
    else
      assert_not_reached! "Let arr unknown " + entry.constructor.name
    end
  end

  def format_let_object_deconstruction({ entries, rhs })
    "let { " + entries.map(::format_let_object_entry).join(", ") + " } = " + .format_node(rhs)
  end

  def format_obj_lit_entry(entry)
    if entry is SimpleObjEntry
      return entry.name + ": " + .format_node(entry.expr)
    else
      assert_not_reached! "Object literal entry not found " + entry.constructor.name
    end
  end

  def format_obj_lit({ entries })
    let entries_js = entries.map(::format_obj_lit_entry)
    if entries_js.map(x => x.length).sum() > 35
      .indent()
      let obj_lit = "{\n" + .padding + entries.map(::format_obj_lit_entry).join(",\n" + .padding) + ",\n"
      .dedent()
      return obj_lit + .padding + "}"
    else
      return "{ " + entries_js.join(", ") + " }"
    end
  end

  def format_new_expr({ expr }) = "new " + .format_node(expr)

  def format_property_lookup({ lhs, property })
    .format_node(lhs) + "[" + .format_node(property) + "]"
  end

  def format_prefix_bind_lookup({ name }) = "::" + name

  def format_command_expr({ name, expr }) = name + " " + .format_node(expr)

  def format_is_operator({ lhs, rhs })
    .format_node(lhs) + " is " + .format_node(rhs)
  end

  def format_if_branch({ test_expr, body })
    let i = "if " + .format_node(test_expr) + "\n"
    i += .format_body(body, IfBranch)
    i
  end

  def format_else_if_branch({ test_expr, body })
    let i = "else if " + .format_node(test_expr) + "\n"
    i += .format_body(body, ElseIfBranch)
    i
  end

  def format_else_branch({ body }) = "else\n" + .format_body(body, ElseBranch)

  def format_branch(branch)
    if branch is IfBranch
      return .format_if_branch(branch)
    else if branch is ElseIfBranch
      return .format_else_if_branch(branch)
    else if branch is ElseBranch
      return .format_else_branch(branch)
    else
      assert_not_reached! "not implemented if branch"
    end
  end

  def is_one_line_return(body) = body.length === 1 && body[0] is ReturnExpr

  def is_early_return(branches)
    branches.length === 1 && branches[0] is IfBranch && .is_one_line_return(branches[0].body)
  end

  def format_if_statement({ branches })
    if .is_early_return(branches)
      let { test_expr, body } = branches[0]
      return .format_node(body[0]) + " if " + .format_node(test_expr)
    else
      return branches.map(::format_branch).join(.padding) + .padding + "end"
    end
  end

  def format_class_entry(entry)
    if entry is ClassGetterExpr
      return "get " + entry.name + " = " + .format_node(entry.expr)
    else if entry is FunctionDef
      return .format_function_def(entry)
    else
      assert_not_reached! "class entry: " + entry.constructor.name
    end
  end

  def format_class_obj_arg(node)
    if node is ObjClassArgEntry
      return node.name
    else if node is DefaultObjClassArg
      return node.name + " = " + .format_node(node.expr)
    else
      assert_not_reached! "obj arg: " + node.constructor.name
    end
  end

  def format_class_arg(arg)
    if arg is NamedClassArg
      return arg.name
    else if arg is ObjClassArg
      return "{ " + arg.entries.map(::format_class_obj_arg).join(", ") + " }"
    else
      assert_not_reached! "class arg: " + arg.constructor.name
    end
  end

  def format_class_def({ name, properties, entries })
    let c = "class " + name
    if properties
      c += "(" + properties.map(::format_class_arg).join(", ") + ")"
    end
    c += "\n"
    .indent()
    c += .padding + entries.map(::format_class_entry).join("\n\n" + .padding)
    .dedent()
    c += "\n"
    c += "end"
  end

  def format_dot_access({ lhs, property }) = .format_node(lhs) + "." + property

  def format_function_call({ lhs_expr, args })
    let start = ""
    let separator = ", "
    let ending = ""
    if args.length > 2
      .indentation += 2
      start = "\n" + .padding
      separator = ",\n" + .padding
      .indentation += -2
      ending = ",\n" + .padding
    end
    let args_output = start + args.map(::format_node).join(separator)
    args_output += ending
    .format_node(lhs_expr) + "(" + args_output + ")"
  end

  def format_node_assignment({ lhs_expr, rhs_expr })
    .format_node(lhs_expr) + " = " + .format_node(rhs_expr)
  end

  def format_node_plus_assignment({ lhs_expr, rhs_expr })
    .format_node(lhs_expr) + " += " + .format_node(rhs_expr)
  end

  def format_for_loop({ iter_name, iterable_expr, body })
    let f = "for let " + iter_name + " of " + .format_node(iterable_expr) + " do\n"
    f += .format_body(body, ForLoop)
    f += .padding + "end"
    f
  end

  def format_import_statement({ imports, path })
    let space = .space_for(imports.length)
    let i = "import {" + space + imports.join("," + space) + "," + space[0] + "}"
    i += " from " + "\"" + path + "\""
    i
  end

  def format_default_import({ name, path })
    "import " + name + " from " + "\"" + path + "\""
  end

  def format_array_literal({ elements })
    return "[]" if elements.length === 0
    let space = .space_for(elements.length)
    "[" + space + elements.map(::format_node).join("," + space) + "," + space[0] + "]"
  end

  def format_js_op_expr({ lhs, type, rhs })
    .format_node(lhs) + " " + type + " " + .format_node(rhs)
  end

  def format_named_let({ name, expr }) = "let " + name + " = " + .format_node(expr)

  def format_num_expr({ value }) = value

  def format_return_expr({ expr }, is_last)
    return .format_node(expr) if is_last
    "return " + .format_node(expr)
  end

  def format_body(body, parent, indent_by = 2)
    new Formatter(body, {
      indentation: .indentation + indent_by,
      parents: .parents.concat(parent),
    }).format().trimEnd() + "\n"
  end

  def format_arg(arg_node)
    if arg_node is SimpleArg
      return arg_node.name
    else if arg_node is SimpleDefaultArg
      return arg_node.name + " = " + .format_node(arg_node.expr)
    else if arg_node is ObjClassArg
      return .format_class_arg(arg_node)
    else
      console.log(arg_node)
      assert_not_reached! "Arg format not implemented for " + node.constructor.name
    end
  end

  def too_long_for_one_liner(expr)
    return false if !expr
    let str = .format_node(expr)
    str.includes("\n") || str.length > 50
  end

  def format_function_def({ name, args, body })
    if body.length > 1 || body[0] is IfStatement || .too_long_for_one_liner(body[0])
      let f = "def " + name
      if args.length > 0
        let args_f = args.map(::format_arg).join(", ")
        f += "(" + args_f + ")"
      end
      f += "\n"
      f += .format_body(body, FunctionDef)
      f += .padding + "end"
      return f
    else
      return "def " + name + "(" + args.map(::format_arg).join(", ") + ") = " + .format_node(body[0], true)
    end
  end
end
let [, , file_name] = process.argv
let str = fs.readFileSync(file_name).toString()
let tokens = new Lexer(str).tokenize()
let ast = new Parser(tokens).parse()
let output = new Formatter(ast).format()
fs.writeFileSync(file_name, output)