open Elang
open Batteries
open BatList
open Prog
open Utils
open Builtins
open Utils

let binop_bool_to_int f x y = if f x y then 1 else 0

(* [eval_binop b x y] évalue l'opération binaire [b] sur les arguments [x]
   et [y]. *)
let eval_binop (b: binop) : int -> int -> int =
  match b with
   | Eadd -> fun x y -> x + y
   | Esub -> fun x y -> x - y
   | Emul -> fun x y -> x * y
   | Ediv -> fun x y -> x / y
   | Emod -> fun x y -> (x mod y)
   | Exor -> fun x y -> x lxor y
   | Eclt -> fun x y -> binop_bool_to_int (<) x y
   | Ecle -> fun x y -> binop_bool_to_int (<=) x y
   | Ecgt -> fun x y -> binop_bool_to_int (>) x y
   | Ecge -> fun x y -> binop_bool_to_int (>=) x y
   | Eceq -> fun x y -> binop_bool_to_int (=) x y
   | Ecne -> fun x y -> binop_bool_to_int (!=) x y

(* [eval_unop u x] évalue l'opération unaire [u] sur l'argument [x]. *)
let eval_unop (u: unop) : int -> int =
  match u with
   | Eneg -> fun x -> -x

(* [eval_eexpr oc ep st e] évalue l'expression [e] dans l'état [st]. Renvoie une
   erreur si besoin. *)
let rec eval_eexpr oc (ep: eprog) st (e : expr) : (int * int Prog.state) res =
   match e with
   | Ebinop(op,expr1,expr2) ->
      let* (value1, st) = eval_eexpr oc ep st expr1 in
      let* (value2, st) = eval_eexpr oc ep st expr2 in
          OK((eval_binop op value1 value2), st)

   | Eunop(op,expr) ->
      let* (value,st) = eval_eexpr oc ep st expr in
      OK((eval_unop op value), st)

   | Evar(str) ->
      (match Hashtbl.find_option st.env str with 
      | None -> Error (Format.sprintf "Unknown variable %s\n" str)
      | Some(name) -> OK(name, st))
   | Eint(n) -> OK(n, st)
   | Ecall(name,expr_list) ->
      let* (args, st) = List.fold_left (fun acc expr ->
         let* (expr_l, st) = acc in
         let* (arg_expr, st) = eval_eexpr oc ep st expr in
         OK(expr_l@[arg_expr], st)
         ) (OK([], st)) expr_list in 
      let* f = find_function ep name in
      let* (res,st) = eval_efun oc ep st f name args in
      match res with
      | Some(n) -> OK(n, st)
      | None -> Error (Format.sprintf "Error with function %s" name)

(* [eval_einstr oc st ins] évalue l'instrution [ins] en partant de l'état [st].

   Le paramètre [oc] est un "output channel", dans lequel la fonction "print"
   écrit sa sortie, au moyen de l'instruction [Format.fprintf].

   Cette fonction renvoie [(ret, st')] :

   - [ret] est de type [int option]. [Some v] doit être renvoyé lorsqu'une
   instruction [return] est évaluée. [None] signifie qu'aucun [return] n'a eu
   lieu et que l'exécution doit continuer.

   - [st'] est l'état mis à jour. *)
and eval_einstr oc (ep: eprog) (st: int state) (ins: instr) :
  (int option * int state) res =
  let rec eval_list_instrs st li : (int option * int state) res =
   match li with 
   | [] -> OK(None,st)
   | instr::next_instrs -> eval_einstr oc ep st instr >>= fun result ->
       match result with
       | (None, st) -> eval_list_instrs st next_instrs
       | (Some(n), st) -> OK(result)
 in
 let rec eval_while st e i: (int option * int state) res =
   let* (condition,st) = eval_eexpr oc ep st e in
      if condition = 1 then 
         eval_einstr oc ep st i >>= fun result_loop ->
         match result_loop with
         | (None, st)-> (eval_while st e i)
         | (Some(n), st) -> OK(result_loop)
      else OK(None, st)
  in 
   match ins with
   | Iassign(str, expr) ->
      let* (value, st) = eval_eexpr oc ep st expr in
      Hashtbl.replace st.env str value;
      OK(None,st)

   | Iif(expr,instr1,instr2) ->
      let* (condition,st) = eval_eexpr oc ep st expr in
        if condition=1 then eval_einstr oc ep st instr1 else eval_einstr oc ep st instr2
   
   | Iwhile(expr,instr) -> eval_while st expr instr

   | Iblock(instrs) -> eval_list_instrs st instrs

   | Ireturn(expr) ->
      let* (result,st) = eval_eexpr oc ep st expr
      in OK(Some(result), st)

   | Iprint(expr) ->
      let* (result, st) = eval_eexpr oc ep st expr in
      Format.fprintf oc "%d\n" result; OK(None, st)
   
   | Icall(name,expr_list) ->
      let* (args, st) = List.fold_left (fun acc expr ->
         let* (expr_l, st) = acc in
         let* (arg_expr, st) = eval_eexpr oc ep st expr in
         OK(expr_l@[arg_expr], st)
         ) (OK([], st)) expr_list in
      let* f = find_function ep name in
      eval_efun oc ep st f name args


(* [eval_efun oc st f fname vargs] évalue la fonction [f] (dont le nom est
   [fname]) en partant de l'état [st], avec les arguments [vargs].

   Cette fonction renvoie un couple (ret, st') avec la même signification que
   pour [eval_einstr]. *)
and eval_efun oc (ep: eprog) (st: int state) ({ funargs; funbody}: efun)
    (fname: string) (vargs: int list)
  : (int option * int state) res =
  (* L'environnement d'une fonction (mapping des variables locales vers leurs
     valeurs) est local et un appel de fonction ne devrait pas modifier les
     variables de l'appelant. Donc, on sauvegarde l'environnement de l'appelant
     dans [env_save], on appelle la fonction dans un environnement propre (Avec
     seulement ses arguments), puis on restore l'environnement de l'appelant. *)
  let env_save = Hashtbl.copy st.env in
  let env = Hashtbl.create 17 in
  match List.iter2 (fun a v -> Hashtbl.replace env a v) funargs vargs with
  | () ->
    eval_einstr oc ep { st with env } funbody >>= fun (v, st') ->
    OK (v, { st' with env = env_save })
  | exception Invalid_argument _ ->
    Error (Format.sprintf
             "E: Called function %s with %d arguments, expected %d.\n"
             fname (List.length vargs) (List.length funargs)
          )

(* [eval_eprog oc ep memsize params] évalue un programme complet [ep], avec les
   arguments [params].

   Le paramètre [memsize] donne la taille de la mémoire dont ce programme va
   disposer. Ce n'est pas utile tout de suite (nos programmes n'utilisent pas de
   mémoire), mais ça le sera lorsqu'on ajoutera de l'allocation dynamique dans
   nos programmes.

   Renvoie:

   - [OK (Some v)] lorsque l'évaluation de la fonction a lieu sans problèmes et renvoie une valeur [v].

   - [OK None] lorsque l'évaluation de la fonction termine sans renvoyer de valeur.

   - [Error msg] lorsqu'une erreur survient.
   *)
let eval_eprog oc (ep: eprog) (memsize: int) (params: int list)
  : int option res =
  let st = init_state memsize in
  find_function ep "main" >>= fun f ->
  (* ne garde que le nombre nécessaire de paramètres pour la fonction "main". *)
  let n = List.length f.funargs in
  let params = take n params in
  let { env } = st in
  let { funargs } = f in
  List.iter2 (fun argname argval -> Hashtbl.replace env argname argval) funargs params;
  eval_efun oc ep st f "main" params >>= fun (v, st) ->
  OK v