open Base
open Ast
open Predef

(* Check if symbol(k) exists in map(map), add it if it does not. *)
let add_symbol map k d = 
  if Map.mem map k then map 
  else Map.add_exn map ~key:k ~data:d
;;

(* First Pass. FIXME *)
let rec first_pass exprs st loc =
  match exprs with
  | [] -> st
  | (Ginstr s)::t -> first_pass t (add_symbol st s loc) loc
  | (Ainstr _ | Cinstr _ )::t -> first_pass t st (loc + 1)
  | _::t -> first_pass t st loc
;;

(* Second Pass. FIXME *)
let rec second_pass exprs st loc =
  let is_new s = 
    try Int.of_string s |> ignore; false
    with Failure _ -> not (Map.mem st s) in
  match exprs with
  | [] -> st
  | (Ainstr a)::t when (is_new a) -> second_pass t (add_symbol st a loc) (loc + 1)
  | _::t -> second_pass t st loc 
  ;;

(* Decimal to Binary . FIXME *)
let translate_ainstr addr st =
  let to_int addr = 
    try Map.find_exn st addr
    with Not_found_s _ -> Int.of_string addr in
  let pad binary = 
    let length = 16 - String.length binary in
    let prefix = String.init length ~f:(fun _ -> '0') in
    String.concat [prefix; binary] in  
  let rec to_binary a =
    match a with
    | 0 -> ""
    | _ ->
      let rem = a % 2 in
      match rem with
      | 0 -> to_binary (a / 2) ^ "0"
      | _ -> to_binary (a / 2) ^ "1" in  
  pad (to_binary (to_int addr))
;;

let translate_cinstr (d, c, j) = 
  String.concat ["111"; comp c; dest d; jump j]
;;

let rec _translate exprs st tt = 
  match exprs with
  | [] -> tt
  | (Ainstr a)::t -> translate_ainstr a st :: _translate t st tt
  | (Cinstr (d, c, j))::t -> translate_cinstr (d, c, j) :: _translate t st tt
  | _::t -> _translate t st tt
;;

let parse s =
  let lexbuf = Lexing.from_string s in
  let ast = Parser.prog Lexer.read lexbuf in 
  ast
;;

let generate_exprs lines = 
  List.map lines ~f:parse
;;

let generate_st exprs = 
  let _st = first_pass exprs symbols 0 in
  second_pass exprs _st 16
;;

let translate lines = 
  let exprs = generate_exprs lines in
  let st = generate_st exprs in
  _translate exprs st []
;;