(* Implementation of exceptions in the CPS interpreter. *)

open Absyn
open Format 
open Env 
open EoplPrint 

exception TODO
exception RuntimeError of string

type value = 
    Num of int 
  | ProcVal of var option * env * var list * exp 
  | BoolVal of bool 
  | ConsVal of value * value 
  | NullVal
  | RefVal  of reference 
  | PairVal of value * value
and env = value Env.env
and reference = Reference of int * value array
         
(* --------- values ---------------------------- *)
  
let rec print_value (v:value) : unit = 
  match v with 
      Num i -> print_int i;
    | BoolVal b -> 
        if b then print_string "true" else 
          print_string "false"
    | ProcVal (None,env,v,str) -> print_string "<closure>"
    | ProcVal (Some s, env,v,str) -> print_string ("<" ^ s ^ ">")
    | RefVal r -> print_string "<reference>"
    | ConsVal (v1, v2) -> 
        print_string "cons(";
        print_value v1; 
        print_string ","; 
        print_value v2; 
        print_string ")"
    | NullVal -> print_string "[]"
    | PairVal (v1, v2) -> 
        print_string "pair(";
        print_value v1; 
        print_string ","; 
        print_value v2; 
        print_string ")"

let closure (ids:var list) (body:exp) (env:env) : value = 
  ProcVal (None, env, ids, body)
                
  (* ------------ References and environments ------------- *)

  let deref (Reference (i,vec)) : value = vec.(i) 
  let setref (Reference (i,vec)) (v:value) : unit =  vec.(i) <- v

  let extend (vars:var list) (vals:value list) (env:env) : env = 
    Extend(vars,Array.of_list vals, env)

  (* Unlike the old version of apply_env, when the location of the
   * value is found, that is what is returned as a reference.  *)
  let rec apply_env_ref (env:env) (sym:var) : reference = 
    match env with 
        Empty -> raise (RuntimeError ("Unbound variable " ^ sym))
      | Extend (vars, vals, old_env) -> 
          (match list_find_position sym vars with 
               Some pos -> Reference(pos,vals)
             | None -> apply_env_ref old_env sym)
                                                       
  let extend_env_rec (proc_names:var list) (idss:var list list) (bodies:exp list) (old_env:env) : env = 
    let num_procs = List.length proc_names             in 
    let array     = Array.make num_procs (Num 0)       in
    let new_env   = Extend (proc_names, array,old_env) in 
    let rec loop proc_names idss bodies n : unit = 
      match proc_names, idss, bodies with 
          (pn::pns), (ids::is), (body::bs)  -> 
            (array.(n) <- ProcVal (Some pn, new_env, ids, body);
             loop pns is bs (n+1))
        | [], [], [] -> () 
        | _, _,_ -> failwith "BUG: Cannot happen" in
      begin
        loop proc_names idss bodies 0;
        new_env
      end 
                

  (* --------- Booleans --------------- *)
  let is_true (v:value) :bool = 
    match v with 
        BoolVal b -> b
      | _ -> raise (RuntimeError "Invalid value used as bool")
                         
  let true_value :value  = BoolVal true
  let false_value :value = BoolVal false
 
(* ----------- Primitives ----------- *)

  let rec apply_prim (p:prim) (args:value list) = 
    match p,args with 
        Plus,  [Num v1;Num v2] -> Num (v1 + v2)
      | Times, [Num v1;Num v2] -> Num (v1 * v2)
      | Minus, [Num v1;Num v2] -> Num (v1 - v2)
      | And,   [BoolVal b1; BoolVal b2] -> BoolVal ( b1 && b2 )
      | Or,    [BoolVal b1; BoolVal b2] -> BoolVal (b1 || b2)
      | Not,   [BoolVal b1] -> BoolVal (not b1)
      | Add1, [Num v1] -> Num (v1 + 1)
      | Sub1, [Num v1] -> Num (v1 - 1)
      | Zero, [Num v1] -> if v1 = 0 then true_value else false_value
      | LT, [Num v1; Num v2] -> if v1 < v2 then true_value else false_value
      | GT, [Num v1; Num v2] -> if v1 > v2 then true_value else false_value
      | Equal, [Num v1;Num v2] -> 
          if (v1 = v2) then true_value  else false_value

      | Hd, [ConsVal(v1,v2)] -> v1
      | Tl, [ConsVal(v1,v2)] -> v2
      | Cons, [v1; v2] -> ConsVal (v1,v2)
      | IsNull, [v] -> 
          (match v with 
               ConsVal (_,_) -> BoolVal false
             | NullVal -> BoolVal true
             | _ -> raise (RuntimeError "Incorrect argument for null?"))
          
      | Pair, [v1; v2] -> PairVal (v1,v2)
      | Fst, [PairVal (v1,v2)] -> v1
      | Snd, [PairVal (v1,v2)] -> v2

      | _, _ -> raise (RuntimeError "Unimplemented  primitive or invalid arguments")

(* --------- Continuations --------------------- *)

type cont = 
    Halt 
  | Test of exp * exp * env * cont
  | EvalFirst of exp list * env * list_cont
  | EvalRator of exp list * env * cont
  | VarassignCont of var * env * cont
  | HandlerCont of var * exp * env * cont 
  | RaiseCont of cont

and list_cont = 
    PrimArgs of prim * cont 
  | EvalRest of value * list_cont 
  | LetBody of var list * exp * env * cont 
  | EvalRands of value * cont 

let rec apply_list_cont (cont:list_cont) (args:value list) : unit = 
  match cont with 
      PrimArgs (prim, cont) ->  
        apply_cont cont (apply_prim prim args)
    | EvalRest (first,cont) -> 
        apply_list_cont cont (first::args)
    | LetBody (vars, body, env, cont) -> 
        eval body (extend vars args env) cont
    | EvalRands (proc,cont) -> 
        apply_procval proc args cont

and apply_cont (cont:cont) (v:value) : unit = 
  match cont with 
      Halt -> print_value v
    | Test (true_exp,false_exp,env,cont) ->
        if (is_true v)
        then (eval true_exp env cont)
        else (eval false_exp env cont)
    | EvalFirst (rest,env,cont) -> 
        eval_rands rest env (EvalRest (v,cont))
    | EvalRator (rands, env, cont) -> 
        eval_rands rands env (EvalRands (v,cont))
    | VarassignCont (id, env, cont) -> 
        begin
          setref (apply_env_ref env id) v;
          apply_cont cont (Num 1)
        end 
	 | HandlerCont (id,exp,env,cont) -> 
		  apply_cont cont v
	 | RaiseCont (cont) ->
		  find_handler cont v

and find_handler (cont:cont) (v:value) = 
  match cont with 
		Halt -> print_string "UNCAUGHT EXCEPTION"
	 | Test (_,_,_,c) -> find_handler c v
	 | EvalFirst (_,_,c) -> find_handler_list c v
	 | EvalRator (_,_,c) -> find_handler c v
	 | VarassignCont (_,_,c) -> find_handler c v
	 | HandlerCont (id,body,env,cont) -> 
		  eval body (extend_env [id] [v] env) cont
	 | RaiseCont (c) -> find_handler c v 
and find_handler_list (cont:list_cont) (v:value) = 
  match cont with 
		PrimArgs (_,c) -> find_handler c v
	 | EvalRest (_,c) -> find_handler_list c v
	 | LetBody  (_,_,_,c) -> find_handler c v
	 | EvalRands (_,c) -> find_handler c v

(* ------------ Evaluator ---------------------- *)
                  
and apply_procval (v:value) (args:value list) (cont:cont) : unit = 
    match v with 
        ProcVal (_, env, ids, body) -> 
          if (List.length ids = List.length args) then 
            eval body (extend_env ids args env) cont
          else 
            raise (RuntimeError "Wrong number of arguments to procedure")
      | _ -> raise (RuntimeError "Attempt to apply non-procedure")

  and eval (exp:exp) (env:env) (cont:cont) : unit = 
    match exp with 
        LitInt i -> apply_cont cont (Num i)
      | LitBool b ->  apply_cont cont (BoolVal b) 
      | LitNull -> apply_cont cont NullVal
      | Var v -> apply_cont cont (apply_env env v)
      | Prim (prim, rands) -> 
          eval_rands rands env (PrimArgs (prim, cont))
      | If (test_exp, true_exp, false_exp) -> 
          eval test_exp env (Test (true_exp, false_exp, env, cont))
      | Let (vars,rands,body) -> 
          eval_rands rands env (LetBody (vars, body, env, cont))
      | Lambda (tys,ids,body) ->            
          apply_cont cont (closure ids body env)
      | App (rator,rands) -> 
          eval rator env 
            (EvalRator (rands, env, cont))
      | Letrec (tys, proc_names, arg_tys, idss, bodies, letrec_body) -> 
          eval letrec_body (extend_env_rec proc_names idss bodies env) cont
      | Varassign (id, rhs) -> 
          eval rhs env
            (VarassignCont (id, env, cont))
		| Failwith string -> 
          print_string string 
		| Try(exp1, id,exp2) -> 
			 eval exp1 env (HandlerCont (id,exp2,env,cont))
		| Raise (exp) -> 
			 eval exp env (RaiseCont cont)
            
  and eval_rands (rands:exp list) (env:env) (cont:list_cont) : unit = 
    match rands with 
        [] -> apply_list_cont cont []
      | (hd::tl) -> 
          eval hd env (EvalFirst (tl, env, cont))

  let eval_program (exp:exp) : unit= 
       eval exp Empty Halt