open Base
open Ast
open Predef

let add_symbol map k d = if Map.mem map k then map else Map.add_exn map ~key:k ~data:d

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

let rec second_pass exprs st loc =
  let is_new a = not (Map.mem st a) 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
;;

let translate_ainstr addr =
  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 addr)
;;

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

let parse s =
  let lexbuf = Lexing.from_string s in
  let ast = Parser.prog Lexer.read lexbuf in
  let eval_ainstr str =
    try Aconst (Int.of_string str) with
    | Failure _ -> Ainstr str
  in
  let _parse expr =
    match expr with
    | Ainstr str -> eval_ainstr str
    | _ -> expr
  in
  _parse ast
;;

let generate_exprs lines = List.map lines ~f:parse
let generate_st exprs = second_pass exprs (first_pass exprs symbols 0) 16

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

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