open! Ast
open! Base

let counter = ref 0

let count =
  fun () ->
    counter := (!counter) + 1;
    !counter

let addr a = "@" ^ a

let label name = "(" ^ name ^ ")"      

let function_definition name arg =
  let push_inner n = [
    addr n
    ; "D=A"
    ; "@LCL"
    ; "A=D+M"
    ; "M=0"
    ; "@SP"
    ; "M=M+1"
    ] in
  let locals = arg
    |> List.init ~f:Int.to_string
    |> List.concat_map ~f:push_inner in
  List.append [ label name ] locals

let function_call name arg =
  let i = Int.to_string (count ()) in
  let arg = Int.to_string arg in
  let ret = name ^ "_RETURN_" ^ i in
  [ addr ret
  ; "D=A"
  ; "@SP"
  ; "A=M"
  ; "M=D"
  ; "@SP"
  ; "M=M+1"
  ; "// push LCL "
  ; "@LCL"
  ; "D=M"
  ; "@SP"
  ; "A=M"
  ; "M=D"
  ; "@SP"
  ; "M=M+1"
  ; "// push ARG "
  ; "@ARG"
  ; "D=M"
  ; "@SP"
  ; "A=M"
  ; "M=D"
  ; "@SP"
  ; "M=M+1"
  ; "// push THIS "
  ; "@THIS"
  ; "D=M"
  ; "@SP"
  ; "A=M"
  ; "M=D"
  ; "@SP"
  ; "M=M+1"
  ; "// push THAT "
  ; "@THAT"
  ; "D=M"
  ; "@SP"
  ; "A=M"
  ; "M=D"
  ; "@SP"
  ; "M=M+1"
  ; "// reposition ARG"
  ; "D=M"
  ; addr arg
  ; "D=D-A"
  ; "@5"
  ; "D=D-A"
  ; "@ARG"
  ; "M=D"
  ; "// reposition LCL"
  ; "@SP"
  ; "D=M"
  ; "@LCL"
  ; "M=D"
  ; "// goto function"
  ; addr name
  ; "0;JMP"
  ; label ret ]

let translate command name arg =
  match command with
  | Function -> function_definition name arg
  | Call -> function_call name arg