open! Ast
open! Base

let addr a = "@" ^ a

let push = [ "@SP"; "A=M"; "M=D"; "@SP"; "M=M+1" ]
let pop = [ "@R13"; "M=D"; "@SP"; "AM=M-1"; "D=M"; "@R13"; "A=M"; "M=D" ]

let common = function
  | Local -> "@LCL"
  | This -> "@THIS"
  | That -> "@THAT"
  | Argument -> "@ARG"
  | _ -> failwith "Invalid Segment"

let pointer = function
  | "0" -> "@THIS"
  | "1" -> "@THAT"
  | _ -> failwith "Invalid Pointer"

let push_constant a = [ addr a; "D=A" ] @ push

let push_pointer a = [ pointer a; "D=M" ] @ push

let push_static offset class_name =
  [ addr (class_name ^ ":" ^ offset); "D=M" ] @ push

let push_temp a = [ addr a; "D=A"; "@5"; "A=D+A"; "D=M" ] @ push

let pop_pointer a =
  [ "@SP"; "AM=M-1"; "D=M"; pointer a; "M=D" ]

let pop_static offset class_name =
  [ "@SP"; "AM=M-1"; "D=M"; addr (class_name ^ ":" ^ offset); "M=D" ]

let push_common segment offset =
  [ addr offset ; "D=A" ; common segment ; "A=D+M" ; "D=M" ] @ push

let pop_temp a =
[ addr a; "D=A"; "@5" ; "D=D+A" ] @ pop

let pop_common segment offset = [
  addr offset
  ; "D=A"
  ; common segment
  ; "D=D+M" ] @ pop

let process_push segment offset class_name =
  match segment with
  | Local | Argument | This | That -> push_common segment offset
  | Constant -> push_constant offset
  | Temp -> push_temp offset
  | Static -> push_static class_name offset
  | Pointer -> push_pointer offset

let process_pop segment offset class_name =
  match segment with
  | Local | Argument | This | That -> pop_common segment offset
  | Temp -> pop_temp offset
  | Static -> pop_static class_name offset
  | Pointer -> pop_pointer offset
  | Constant -> failwith "Invalid Segment"

let translate command segment offset class_name =
  let offset = Int.to_string offset in
    match command with
    | Pop -> process_pop segment offset class_name
    | Push -> process_push segment offset class_name