(* Implementation of the multi-threaded interpreter discussed in class. *)

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
  (* Is the lock occupied *)
  | LockVal of bool ref 
  (* ML-style references *)
  | CellVal of value ref 

and env = value Env.env
and reference = Reference of int * value array
         
(* --------- values ---------------------------- *)
  
let rec is_list v = 
  match v with 
		NullVal -> true
	 | ConsVal(_,v2) -> is_list v2
	 | _ -> false

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) when is_list v2 -> 
		  print_string "["; print_value v1; print_list v2; print_string "]"
    | 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 ")"
	 | LockVal (_) -> print_string "<lock>"
	 | CellVal _ -> print_string "<cell>"
and print_list (v:value) : unit = 
  match v with 
		NullVal -> ()
	 | ConsVal(v1,v2) -> 
		  print_string ";";
		  print_value v1;
		  print_list v2
	 | _ -> failwith "Can't print that kind of list"

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

		| Cell, [v] -> CellVal (ref v)
		| Setcell, [CellVal r; v] -> 
			 (r := v;
			  Num 1)
		| Contents, [CellVal r] -> !r 

		| Print, [v] -> print_value v; Num 1

		| Lock, [] -> LockVal (ref false)
		| Release, [LockVal (occupied)] -> 
			 if !occupied then 
				(occupied:= false;
				 Num 1)
			 else 
				raise (RuntimeError "Must acquire lock before releasing")
      | _, _ -> 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
  | AcquireCont of cont
  | DieCont 

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

type thread = Done of value 
			   | Running of (unit -> thread)

let step_thread (t:thread) : thread = 
  match t with  
		Done v -> Done v
	 | Running proc -> proc ()

let make_thread (p:unit -> thread) : thread = Running p

let ready_queue : thread Queue.t = Queue.create ()

let the_final_answer : thread ref = ref (Done (Num 1))

let get_next_from_ready_queue _ : thread = 
  try 
	 Queue.pop ready_queue
  with
		Queue.Empty -> !the_final_answer

let rec schedule (quantum:int) (t1:thread) : value =
  Queue.add t1 ready_queue;
  let rec timer_loop (t2:thread)(ticks:int) = 
	 match t2 with 
		  Done v -> v
		| Running p -> if ticks > 0 then 
			 timer_loop (step_thread t2) (ticks - 1)
		  else schedule quantum t2
  in 
	 timer_loop (get_next_from_ready_queue ()) quantum
		
let rec apply_list_cont (cont:list_cont) (args:value list) : thread = 
  make_thread 
	 (fun _ -> 
		 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
			| BeginCont (cont) -> 
				 let rec last l = 
		         match l with 
						 [] -> failwith "BUG: Cannot happen"
					  | [v] -> apply_cont cont v
					  | (hd::tl) -> last tl
				 in last args
	 ) 

and apply_cont (cont:cont) (v:value) : thread = 
  make_thread 
	 (fun _ -> 
		 match cont with 
			  Halt -> 
				 begin
					print_string "The final answer is: ";
					print_value v;
					print_newline ();
					the_final_answer := Done v;
					get_next_from_ready_queue ()
				 end
			| 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
			| DieCont -> 
				 get_next_from_ready_queue ()
			| AcquireCont (cont1) -> 
				 match v with 
					  LockVal (occupied) -> 
						 if !occupied then
							(Queue.add (apply_cont cont v) ready_queue;
							 get_next_from_ready_queue ())
						 else 
							(occupied := true;
							 apply_cont cont1 (Num 1))
					| _ -> raise (RuntimeError "Non-lock acquired")						 
	 )
		  

and find_handler (cont:cont) (v:value) : thread = 
  make_thread 
	 (fun _ -> 
		 match cont with 
			  Halt -> 
				 (print_string "UNCAUGHT EXCEPTION"; 
				  get_next_from_ready_queue ())
			| DieCont -> 
				 (print_string "UNCAUGHT EXCEPTION";
				  get_next_from_ready_queue ())
			| 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 
			| AcquireCont c -> find_handler c v
	 )

and find_handler_list (cont:list_cont) (v:value) :thread = 
  make_thread
	 (fun _ -> 
		 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
			| BeginCont (c) -> find_handler c v
	 )
(* ------------ Evaluator ---------------------- *)
                  
and apply_procval (v:value) (args:value list) (cont:cont) : thread = 
    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")
      | _ -> (print_string "Nonproc:"; print_value v;
			 raise (RuntimeError ("Attempt to apply non-procedure ")))

  and eval (exp:exp) (env:env) (cont:cont) : thread = 
    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))
      | Begin(exp, exps) -> 
			 eval_rands (exp::exps) env (BeginCont(cont))
		| Try(exp1, id,exp2) -> 
			 eval exp1 env (HandlerCont (id,exp2,env,cont))
		| Raise (exp) -> 
			 eval exp env (RaiseCont cont)
	   (* Start a new process by adding a new thread to the 
			ready_queue. Return the value "1" to the current continuation. *)
		| Spawn exp -> 
			 begin
				Queue.add
				  (make_thread (fun _ -> eval exp env DieCont))
				  ready_queue;
				(apply_cont cont (Num 1))
			 end
	   (* Acquire a lock *)
		| Acquire exp ->
			 eval exp env (AcquireCont (cont))
            
  and eval_rands (rands:exp list) (env:env) (cont:list_cont) : thread = 
    match rands with 
        [] -> apply_list_cont cont []
      | (hd::tl) -> 
          eval hd env (EvalFirst (tl, env, cont))

  let eval_program (quantum:int) (exp:exp) : value = 
	 Queue.clear ready_queue;
	 schedule quantum
		(make_thread 
			(fun () -> eval exp Empty Halt))