open Printf ;;
open Interp_ast ;;
open Interp ;;
open Interp_lib ;;
open Interp_interface ;;
open Interp_inter_imp ;;

open Big_int ;;

let lit_to_string = function
 | L_unit -> "unit"
 | L_zero -> "0b0"
 | L_one -> "0b1"
 | L_true -> "true"
 | L_false -> "false"
 | L_num n -> string_of_big_int n
 | L_hex s -> "0x"^s
 | L_bin s -> "0b"^s
 | L_undef -> "undefined"
 | L_string s -> "\"" ^ s ^ "\""
;;

let id_to_string = function
  | Id_aux(Id s,_) | Id_aux(DeIid s,_) -> s
;;

let loc_to_string = function
  | Unknown -> "location unknown"
  | Int(s,_) -> s
  | Range(s,fline,fchar,tline,tchar) -> 
      if fline = tline
      then sprintf "%s:%d:%d" s fline fchar
      else sprintf "%s:%d:%d-%d:%d" s fline fchar tline tchar
;;

let collapse_leading s =
  if String.length s <= 8 then s else
  let first_bit = s.[0] in
  let templ = sprintf "%c...%c" first_bit first_bit in
  let regexp = Str.regexp "^\\(000000*\\|111111*\\)" in
  Str.replace_first regexp templ s
;;

let bitvec_to_string l = "0b" ^ collapse_leading (String.concat "" (List.map (function
  | V_lit(L_aux(L_zero, _)) -> "0"
  | V_lit(L_aux(L_one, _)) -> "1"
  | _ -> assert false) l))
;;

(*let val_to_string v = match v with
  | Bitvector(bools, _, _) -> "0b" ^ collapse_leading (String.concat "" (List.map (function | true -> "1" | _ -> "0") bools))
  | Bytevector words->
    "0x" ^ (String.concat ""
	    (List.map (function 
	      | 10 -> "A"
	      | 11 -> "B"
	      | 12 -> "C"
	      | 13 -> "D"
	      | 14 -> "E"
	      | 15 -> "F"
	      | i  -> string_of_int i) words))		
  | Unknown0 -> "Unknown"*)

(*let reg_name_to_string = function
  | Reg0 s -> s
  | Reg_slice(s,(first,second)) -> s (*contemplate putting slice here*)
  | Reg_field(s,f,_) -> s ^ "." ^ f
  | Reg_f_slice(s,f,_,(first,second)) -> s ^ "." ^ f*)

let rec reg_to_string = function
  | Reg (id,_) -> id_to_string id
  | SubReg (id,r,_) -> sprintf "%s.%s" (reg_to_string r) (id_to_string id)
;;

let rec val_to_string_internal = function
 | V_boxref(n, t) -> sprintf "boxref %d" n
 | V_lit (L_aux(l,_)) -> sprintf "%s" (lit_to_string l)
 | V_tuple l ->
     let repr = String.concat ", " (List.map val_to_string_internal l) in
     sprintf "(%s)" repr
 | V_list l ->
     let repr = String.concat "; " (List.map val_to_string_internal l) in
     sprintf "[||%s||]" repr
 | V_vector (first_index, inc, l) ->
     let last_index = add_int_big_int (if inc then List.length l - 1 else 1 - List.length l) first_index  in
     let repr =
       try bitvec_to_string (* (if inc then l else List.rev l)*) l
       with Failure _ ->
         sprintf "[%s]" (String.concat "; " (List.map val_to_string_internal l)) in
     sprintf "%s [%s..%s]" repr (string_of_big_int first_index) (string_of_big_int last_index)
 | V_record(_, l) ->
     let pp (id, value) = sprintf "%s = %s" (id_to_string id) (val_to_string_internal value) in
     let repr = String.concat "; " (List.map  pp l) in
     sprintf "{%s}" repr
 | V_ctor (id,_, value) ->
     sprintf "%s %s" (id_to_string id) (val_to_string_internal value)
 | V_register r ->
     sprintf "reg-as-value %s" (reg_to_string r)
 | V_unknown -> "unknown"
;;

let rec top_frame_exp_state = function
  | Top -> raise (Invalid_argument "top_frame_exp")
  | Hole_frame(_, e, _, env, mem, Top) 
  | Thunk_frame(e, _, env, mem, Top) -> (e,(env,mem))
  | Thunk_frame(_, _, _, _, s)
  | Hole_frame(_, _, _, _, _, s) -> top_frame_exp_state s

let tunk = Unknown, None
let ldots = E_aux(E_id (Id_aux (Id "...", Unknown)), tunk)
let rec compact_exp (E_aux (e, l)) =
  let wrap e = E_aux (e, l) in
  match e with
 | E_block (e :: _) -> compact_exp e
 | E_nondet (e :: _) -> compact_exp e
 | E_if (e, _, _) ->
     wrap(E_if(compact_exp e, ldots, E_aux(E_block [], tunk)))
 | E_for (i, e1, e2, e3, o, e4) ->
    wrap(E_for(i, compact_exp e1, compact_exp e2, compact_exp e3, o, ldots))
 | E_case (e, _) ->
     wrap(E_case(compact_exp e, []))
 | E_let (bind, _) -> wrap(E_let(bind, ldots))
 | E_app (f, args) -> wrap(E_app(f, List.map compact_exp args))
 | E_app_infix (l, op, r) -> wrap(E_app_infix(compact_exp l, op, compact_exp r))
 | E_tuple exps -> wrap(E_tuple(List.map compact_exp exps))
 | E_vector exps -> wrap(E_vector(List.map compact_exp exps))
 | E_vector_access (e1, e2) ->
     wrap(E_vector_access(compact_exp e1, compact_exp e2))
 | E_vector_subrange (e1, e2, e3) ->
     wrap(E_vector_subrange(compact_exp e1, compact_exp e2, compact_exp e3))
 | E_vector_update (e1, e2, e3) ->
     wrap(E_vector_update(compact_exp e1, compact_exp e2, compact_exp e3))
 | E_vector_update_subrange (e1, e2, e3, e4) ->
     wrap(E_vector_update_subrange(compact_exp e1, compact_exp e2, compact_exp e3, compact_exp e4))
 | E_vector_append (e1, e2) ->
     wrap(E_vector_append(compact_exp e1, compact_exp e2))
 | E_list exps -> wrap(E_list(List.map compact_exp exps))
 | E_cons (e1, e2) ->
     wrap(E_cons(compact_exp e1, compact_exp e2))
 | E_record_update (e, fexps) ->
     wrap(E_record_update (compact_exp e, fexps))
 | E_field (e, id) ->
     wrap(E_field(compact_exp e, id))
 | E_assign (lexp, e) -> wrap(E_assign(lexp, compact_exp e))
 | E_block [] | E_nondet [] | E_cast (_, _) | E_internal_cast (_, _)
 | E_id _|E_lit _|E_vector_indexed (_, _)|E_record _|E_internal_exp _ ->
     wrap e

(* extract, compact and reverse expressions on the stack;
 * the top of the stack is the head of the returned list. *)
let rec compact_stack ?(acc=[]) = function
  | Top -> acc
  | Hole_frame(_,e,_,env,mem,s)
  | Thunk_frame(e,_,env,mem,s) -> compact_stack ~acc:(((compact_exp e),(env,mem)) :: acc) s
;;  

let sub_to_string = function None -> "" | Some (x, y) -> sprintf " (%s, %s)"
  (string_of_big_int x) (string_of_big_int y)
;;

let id_compare i1 i2 = 
  match (i1, i2) with 
    | (Id_aux(Id(i1),_),Id_aux(Id(i2),_)) 
    | (Id_aux(Id(i1),_),Id_aux(DeIid(i2),_)) 
    | (Id_aux(DeIid(i1),_),Id_aux(Id(i2),_))
    | (Id_aux(DeIid(i1),_),Id_aux(DeIid(i2),_)) -> compare i1 i2

module Reg = struct
  include Map.Make(struct type t = id let compare = id_compare end)
  let to_string id v =
    sprintf "%s -> %s" (id_to_string id) (val_to_string_internal v)
  let find id m =
(*    eprintf "reg_find called with %s\n" (id_to_string id);*)
    let v = find id m in
(*    eprintf "%s -> %s\n" (id_to_string id) (val_to_string v);*)
    v
end ;;
(* Old Mem, that used the id to map as well as the int... which seems wrong
module Mem = struct
  include Map.Make(struct
    type t = (id * big_int)
    let compare (i1, v1) (i2, v2) =
      (* optimize for common case: different addresses, same id *)
      match compare_big_int v1 v2 with
      | 0 -> id_compare i1 i2
      | n -> n
    end)
  (* debugging memory accesses
  let add (n, idx) v m =
    eprintf "%s[%s] <- %s\n" (id_to_string n) (string_of_big_int idx) (val_to_string v);
    add (n, idx) v m
  let find (n, idx) m =
    let v = find (n, idx) m in
    eprintf "%s[%s] -> %s\n" (id_to_string n) (string_of_big_int idx) (val_to_string v);
    v
  *)
  let to_string (n, idx) v =
    sprintf "%s[%s] -> %s" (id_to_string n) (string_of_big_int idx) (val_to_string v)
end ;;*)

module Mem = struct
  include Map.Make(struct
    type t = big_int
    let compare v1 v2 = compare_big_int v1 v2
  end)
  (* debugging memory accesses
  let add idx v m =
    eprintf "[%s] <- %s\n" (string_of_big_int idx) (val_to_string v);
    add idx v m
  let find idx m =
    let v = find idx m in
    eprintf "[%s] -> %s\n" (string_of_big_int idx) (val_to_string v);
    v
  *)
  let to_string idx v =
    sprintf "[%s] -> %s" (string_of_big_int idx) (val_to_string_internal v)
end ;;


let vconcat v v' = vec_concat (V_tuple [v; v']) ;;

let slice v = function
  | None -> v
  | Some (n, m) -> slice_vector v n m
;;

let rec slice_ir v = function
  | BF_single n -> slice_vector v n n
  | BF_range (n, m) -> slice_vector v n m
  | BF_concat (BF_aux (ir, _), BF_aux (ir', _)) -> vconcat (slice_ir v ir) (slice_ir v ir')
;;

let unit_lit = V_lit (L_aux(L_unit,Interp_ast.Unknown))

let rec perform_action ((reg, mem) as env) = function
 (* registers *)
 | Read_reg (Reg (id, _), sub) ->
     slice (Reg.find id reg) sub, env
 | Write_reg (Reg (id, _), None, value) ->
     unit_lit, (Reg.add id value reg, mem)
 | Write_reg (Reg (id, _), Some (start, stop), (V_vector _ as value)) ->
     let old_val = Reg.find id reg in
     let new_val = fupdate_vector_slice old_val value start stop in
     unit_lit, (Reg.add id new_val reg, mem)
 (* subregisters *)
 | Read_reg (SubReg (_, Reg (id, _), BF_aux (ir, _)), sub) ->
     slice (slice_ir (Reg.find id reg) ir) sub, env
 | Write_reg (SubReg (_, (Reg _ as r), BF_aux (ir, _)), None, value) ->
     (match ir with
     | BF_single n ->
         perform_action env (Write_reg (r, Some(n, n), value))
     | BF_range (n, m) ->
         perform_action env (Write_reg (r, Some(n, m), value))
     | BF_concat _ -> failwith "unimplemented: non-contiguous register write")
 (* memory *)
 | Read_mem (id, V_lit(L_aux((L_num n),_)), sub) ->
     slice (Mem.find n mem) sub, env
 | Write_mem (id, V_lit(L_aux(L_num n,_)), None, value) ->
     unit_lit, (reg, Mem.add n value mem)
 (* multi-byte accesses to memory *)
 | Read_mem (id, V_tuple [V_lit(L_aux(L_num n,_)); V_lit(L_aux(L_num size,_))], sub) ->
     let rec fetch k acc =
       if eq_big_int k size then slice acc sub else
         let slice = Mem.find (add_big_int n k) mem in
         fetch (succ_big_int k) (vconcat acc slice)
     in
     fetch zero_big_int (V_vector (zero_big_int, true, [])), env
 (* XXX no support for multi-byte slice write at the moment *)
 | Write_mem (id, V_tuple [V_lit(L_aux(L_num n,_)); V_lit(L_aux(L_num size,_))], None, V_vector (m, inc, vs)) ->
     (* normalize input vector so that it is indexed from 0 - for slices *)
     let value = V_vector (zero_big_int, inc, vs) in
     (* assumes smallest unit of memory is 8 bit *)
     let byte_size = 8 in
     let rec update k mem =
       if eq_big_int k size then mem else
         let n1 = mult_int_big_int byte_size k in
         let n2 = sub_big_int (mult_int_big_int byte_size (succ_big_int k)) (big_int_of_int 1) in
         let slice = slice_vector value n1 n2 in
         let mem' = Mem.add (add_big_int n k) slice mem in
         update (succ_big_int k) mem'
     in unit_lit, (reg, update zero_big_int mem)
 (* This case probably never happens in the POWER spec anyway *)
 | Write_mem (id, V_lit(L_aux(L_num n,_)), Some (start, stop), (V_vector _ as value)) ->
     let old_val = Mem.find n mem in
     let new_val = fupdate_vector_slice old_val value start stop in
     unit_lit, (reg, Mem.add n new_val mem)
 (* special case for slices of size 1: wrap value in a vector *)
 | Write_reg ((Reg (_, _) as r), (Some (start, stop) as slice), value) when eq_big_int start stop ->
     perform_action env (Write_reg (r, slice, V_vector(zero_big_int, true, [value])))
 | Write_mem (id, (V_lit(L_aux(L_num _,_)) as n), (Some (start, stop) as slice), value) when eq_big_int start stop ->
     perform_action env (Write_mem (id, n, slice, V_vector(zero_big_int, true, [value])))
 (* extern functions *)
 | Call_extern (name, arg) -> eval_external name arg, env
 | Interp.Step _ | Nondet _ | Exit _ -> unit_lit, env
 | _ -> assert false
;;

let debug = ref true
let debugf : ('a, out_channel, unit) format -> 'a = function f -> if !debug then eprintf f else ifprintf stderr f

type interactive_mode = Step | Run | Next

let mode_to_string = function
  | Step -> "step"
  | Run -> "run"
  | Next -> "next"

(* ANSI/VT100 colors *)
let disable_color = ref false
let color bright code s =
  if !disable_color then s
  else sprintf "\x1b[%s3%dm%s\x1b[m" (if bright then "1;" else "") code s
let red = color true 1
let green = color false 2
let yellow = color true 3
let blue = color true 4
let grey = color false 7

let run
  ?(entry=E_aux(E_app(Id_aux((Id "main"),Unknown), [E_aux(E_lit (L_aux(L_unit,Unknown)),(Unknown,None))]),(Unknown,None)))
  ?(reg=Reg.empty)
  ?(mem=Mem.empty)
  ?(eager_eval=true)
  ?mode
  (name, test) =
  let get_loc (E_aux(_, (l, _))) = loc_to_string l in
  let print_exp env e =
    debugf "%s: %s\n" (get_loc e) (Pretty_interp.pp_exp env Printing_functions.red e) in
  (* interactive loop for step-by-step execution *)
  let usage = "Usage:
    step    go to next action [default]
    next    go to next break point
    run     complete current execution,
    bt      print call stack
    cont    print continuation of the top stack frame
    env     print content of environment
    mem     print content of memory
    quit    exit interpreter" in
  let rec interact mode ((reg, mem) as env) stack =
    flush_all();
    let command = Pervasives.read_line () in
    let command' = if command = "" then mode_to_string mode else command in
    begin match command' with
    | "s" | "step" -> Step
    | "n" | "next" -> Next
    | "r" | "run" -> Run
    | "e" | "env" | "environment" ->
        Reg.iter (fun k v -> debugf "%s\n" (Reg.to_string k v)) reg;
        interact mode env stack
    | "m" | "mem" | "memory" ->
        Mem.iter (fun k v -> debugf "%s\n" (Mem.to_string k v)) mem;
        interact mode env stack
    | "bt" | "backtrace" | "stack" ->
        List.iter (fun (e,(env,mem)) -> print_exp env e) (compact_stack stack);
        interact mode env stack
    | "c" | "cont" | "continuation" ->
        (* print not-compacted continuation *)
      let (e,(lenv,lmem)) = top_frame_exp_state stack in
        print_exp lenv e;
        interact mode env stack
    | "show_casts" ->
        Pretty_interp.ignore_casts := false;
        interact mode env stack
    | "hide_casts" ->
        Pretty_interp.ignore_casts := true;
        interact mode env stack
    | "q" | "quit" | "exit" -> exit 0
    | _ -> debugf "%s\n" usage; interact mode env stack
    end
  in
  let rec loop mode env = function
  | Value v ->
      debugf "%s: %s %s\n" (grey name) (blue "return") (val_to_string_internal v);
      true, mode, env
  | Action (a, s) ->
      let (top_exp,(top_env,top_mem)) = top_frame_exp_state s in
      let loc = get_loc (compact_exp top_exp) in
      let return, env' = perform_action env a in
      let step ?(force=false) () =
        if mode = Step || force then begin
          debugf "%s\n" (Pretty_interp.pp_exp top_env Printing_functions.red top_exp);
          interact mode env s
       end else
         mode in
      let show act lhs arrow rhs = debugf "%s: %s: %s %s %s\n"
        (grey loc) (green act) lhs (blue arrow) rhs in
      let left = "<-" and right = "->" in
      let (mode',env',s) = begin match a with
      | Read_reg (reg, sub) ->
          show "read_reg" (reg_to_string reg ^ sub_to_string sub) right (val_to_string_internal return);
          step (),env',s
      | Write_reg (reg, sub, value) ->
          assert (return = unit_lit);
          show "write_reg" (reg_to_string reg ^ sub_to_string sub) left (val_to_string_internal value);
          step (),env',s
      | Read_mem (id, args, sub) ->
          show "read_mem" (id_to_string id ^ val_to_string_internal args ^ sub_to_string sub) right (val_to_string_internal return);
          step (),env',s
      | Write_mem (id, args, sub, value) ->
          assert (return = unit_lit);
          show "write_mem" (id_to_string id ^ val_to_string_internal args ^ sub_to_string sub) left (val_to_string_internal value);
          step (),env',s
      (* distinguish single argument for pretty-printing *)
      | Call_extern (f, (V_tuple _ as args)) ->
          show "call_lib" (f ^ val_to_string_internal args) right (val_to_string_internal return);
          step (),env',s
      | Call_extern (f, arg) ->
          show "call_lib" (sprintf "%s(%s)" f (val_to_string_internal arg)) right (val_to_string_internal return);
          step (),env',s
      | Interp.Step _ ->
          assert (return = unit_lit);
          show "breakpoint" "" "" "";
          step ~force:true (),env',s
      | Nondet exps ->
	let stacks = List.sort (fun (_,i1) (_,i2) -> compare i1 i2)
	                       (List.combine (List.map (set_in_context s) exps)
	                                     (List.map (fun _ -> Random.bits ()) exps)) in
	show "nondeterministic evaluation begun" "" "" "";
	let (_,_,env') = List.fold_right (fun (stack,_) (_,_,env') -> loop mode env' (resume {eager_eval = (mode = Run); track_values = false;} stack None)) stacks (false,mode,env'); in
	show "nondeterministic evaluation ended" "" "" "";
	step (),env',s
      | Exit e ->
	show "exiting current evaluation" "" "" "";
	step (),env', (set_in_context s e)
      | Barrier (_, _) | Write_next_IA _  ->
          failwith "unexpected action"
      end in
      loop mode' env' (resume {eager_eval = (mode' = Run);track_values = false} s (Some return))
  | Error(l, e) ->
      debugf "%s: %s: %s\n" (grey (loc_to_string l)) (red "error") e;
      false, mode, env in
  debugf "%s: %s %s\n" (grey name) (blue "evaluate") 
    (Pretty_interp.pp_exp Interp.eenv Printing_functions.red entry);
  let mode = match mode with
  | None -> if eager_eval then Run else Step
  | Some m -> m in
  try
    Printexc.record_backtrace true;
    loop mode (reg, mem) (interp {eager_eval = eager_eval; track_values = false} (fun id -> None) test entry)
  with e ->
    let trace = Printexc.get_backtrace () in
    debugf "%s: %s %s\n%s\n" (grey name) (red "interpretor error") (Printexc.to_string e) trace;
    false, mode, (reg, mem)
;;