module StringMap = Map.Make(String)

(* Symbol Table *)
type symbol_table = {
  parent : symbol_table option;
  vars : Ast.varDecl list;
  }

(* Translation Environment  *)
type env = {
  function_index : Ast.funcDecl StringMap.t;  
  object_index : Ast.objDecl StringMap.t;
  var_decls : Ast.varDecl StringMap.t;
  locals : symbol_table;
  }

let symbol_table_add_var scope (decl : Ast.varDecl) = 
  decl :: scope.vars 

let rec symbol_table_find_var scope name =
  try
    List.find (fun (decl : Ast.varDecl) -> decl.vname = name) scope.vars
  with Not_found ->
    match scope.parent with
      Some(parent) -> symbol_table_find_var parent name
    | _ -> raise Not_found

(* Expression Sum Type *)
type expression_detail =
    Literal of int
  | LitFloat of float
  | LitBool of bool
  | LitString of string
  | Id of string
  | Binop of Ast.expr * Ast.op * Ast.expr
  | Assign of Ast.expr * Ast.expr
  | Call of string * Ast.expr list 
  | IfThenElse of Ast.expr * Ast.expr * Ast.expr * Ast.stmt list * Ast.stmt list
  | ArrayAccess of string * Ast.expr  
  | MemberAccess of string * string
  | MethodCall of string * string * Ast.expr list
  | CSVInput
(*  | MemberAccess of expr * expr
  | ArrayAccess of expr * expr  )
  | If of expr * expr * expr *)
  | Noexpr

type expression = expression_detail * Ast.progType

type statement = 
    Block of Ast.stmt list (* Maybe a scope *)
  | Expression of expression * env * Ast.expr
  | For of Ast.expr * Ast.expr * Ast.expr * statement list * env
  | Foreach of string * statement list * env
  | Return of Ast.expr

(* Zip two lists *)
let rec zip l1 l2 = 
  match l1,l2 with
  | [],_ -> []
  | _, []-> []
  | (x::xs),(y::ys) -> (x,y) :: (zip xs ys)
               
(* Check expression -> answer a Sast.expression *)
let rec check_expr env = function
    Ast.Literal(v) -> (Literal(v), Ast.Int )
  | Ast.LitFloat(v) -> (LitFloat(v), Ast.Float )
  | Ast.LitBool(v) -> (LitBool(v), Ast.Boolean )
  | Ast.LitString(v) -> (LitString(v), Ast.String )
  | Ast.Id(name) -> 
     
     (* Check local variable is in scope *)
     let varDecl = 
       try
         StringMap.find name env.var_decls
       with Not_found ->
          raise( Failure "Could not find local variable")
     in
     Id(name), varDecl.varType

  | Ast.Binop(e1,op,e2) -> 
     let (_,t1) = check_expr env e1 
     and (_,t2) = check_expr env e2 in
         if op <> Ast.Equal && op <> Ast.Neq && op <> Ast.Leq && op <> Ast.Less && op <> Ast.Greater && op <> Geq then
           (* Make sure both left and right side are the same type TODO: Other types *) 
           if t1 <> t2 then raise( Failure "Mismatched type") 
           else Binop( e1, op, e2), t1
         else
           Binop( e1, op, e2), Ast.Boolean 
  | Ast.Assign(eVar,e) -> (
     let (_,t1) = check_expr env eVar 
     and (_,t2) = check_expr env e in
     Assign(eVar,e),t1)
     
          (*     if (t1 != t2) then (Assign(eVar,e),t1)
       else raise( Failure "Mismatch type in assignment" ))*)
     
  | Ast.Call( name, expr_li ) -> (
     
     match name with
       "Print" -> Call( name, expr_li), Ast.Int
     | _ -> (

     (* Check name exists in function table *)
      let func_decl name = 
        try
          StringMap.find name env.function_index
        with Not_found ->
          raise( Failure "A: No function name found" )
      in
     
     (* Check count of arguments and compare to func definition *)
      let compareArguments exprLi formals = 
        if List.length exprLi <> List.length formals 
        then raise( Failure "Arguments and parameters don't match")
        else
          let zipList = zip exprLi formals in
          let result = checkParameters env zipList in
          match result with
            false -> raise (Failure "Mismatch in arguments and parameters" )
          | true -> true
       in

     (* TODO: Check return type *)
       let fd = func_decl name in
       if ((compareArguments expr_li fd.formals) = false)
       then raise( Failure "Mismatch in arguments and parameters" )
       else Call(name,expr_li), fd.retType))

  | Ast.IfThenElse( e1, e2, e3, sl1, sl2 ) ->

     (* Get Types *)
     let (_,t1) = check_expr env e1 
     and (_,t2) = check_expr env e2 
     and (_,t3) = check_expr env e3 in

     (* Check type of predicate *)
     let predicate = t1 <> Ast.Boolean in

     (* Check that t2 and t3 match *)
     let thenAndElse = t2 <> t3 in

     (* Check error and answer *)
     if predicate || thenAndElse then raise( Failure "Mismatch type in if-then-else")
     else
       IfThenElse( e1, e2, e3, sl1, sl2), Ast.Int (* Use Int Type here *)
  | Ast.ArrayAccess(name,expr) -> let arrVar =
     try
          StringMap.find name env.var_decls
     with Not_found ->
       raise( Failure "No variable found for array access" ) in
     let (_,t) = check_expr env expr in ArrayAccess(name,expr),t
(*     if (t = arrVar.varType) then ArrayAccess(name,expr),t
       else raise( Failure "Type mismatch")*)


  | Ast.MemberAccess(name,field) -> 
     let varD = 
       try
         StringMap.find name env.var_decls
       with Not_found ->
          raise( Failure "Could not find local variable for member access") in
     let clName = function
       Ast.Record(name) -> name
      | _ -> raise( Failure "Not an object for member access" ) in
     let obj =
     try
          StringMap.find (clName varD.varType) env.object_index
     with Not_found ->
       raise( Failure "No object found for array access" ) in
     let (name,t) = 
       List.find (fun (a,b) -> a = field) obj.objFields in
     MemberAccess(name,field),t
  | Ast.MethodCall(s1,s2,e_li) -> MethodCall(s1,s2,e_li),Ast.Int
  | Ast.CSVInput -> CSVInput,Ast.Int
   | Ast.Noexpr -> raise( Failure "No expression here")

and checkParameters env zipList =
  let cmp ( e, (decl : Ast.varDecl) ) = 
    let _,t = check_expr env e in
    if t <> decl.varType then None
    else Some(true)
  in
  let results = List.map cmp zipList in
  let fail = List.filter (fun a -> if a = None then true else false) results in
  match fail with
    [] -> true
  | _ -> false


(* Check individual statements -> answer Sast.statement *)
let rec check_stmt environment = function
    Ast.Block(stmts) -> Block(stmts)   (* TODO *)
  | Ast.Expr(e) -> 
     let result = check_expr environment e in
     Expression( result, environment, e )
  | Ast.For(e1,e2,e3,li) -> 
     (* Check that e2 is a Boolean *)
     let (_,t) = check_expr environment e2 in
     
     (* Check individual statements *)
     let statementList = (List.map (check_stmt environment) li) in

     (* Answer checked for loop *)
     if t != Ast.Boolean then raise( Failure "For loop doesn't have a Boolean")
     else For(e1,e2,e3,statementList,environment)
  | Ast.Return(e) -> let (_,t) = check_expr environment e in
     Return(e)
  | Ast.Foreach( name, stmt_li ) ->(
     (* Lookup name *)

     (* Check statements *)
     let statementList = (List.map (check_stmt environment) stmt_li) in
     
     (* Answer checked foreach loop *)
     Foreach( name, statementList, environment ));;

let load_var_decls _map (varDecls : Ast.varDecl list) =
  List.fold_left (fun m (varD : Ast.varDecl) -> StringMap.add varD.vname varD m) _map varDecls

(* Check function definitions and answer bool option *)
let check_func_decl environment (funcDecl : Ast.funcDecl) =(

  (* Check if name exists in environment *)
  let checkFuncNameExists = 
    try
      StringMap.find funcDecl.name environment.function_index
    with Not_found ->
      raise( Failure "B: No function name found" )
  in

  (* Build new environment with empty locals *)
  let env = {
    function_index = environment.function_index;
    object_index = environment.object_index;
    var_decls = load_var_decls (load_var_decls (environment.var_decls) funcDecl.locals) funcDecl.formals;
    locals = { parent = None; vars = [] } } in

  (* check statements -> get back a statement list *)
  let stmts = (List.map (check_stmt env) funcDecl.body) in

  (* Answer OKAY *)
  checkFuncNameExists;
  stmts)


let load_function_index _map funcDecls = (
  List.fold_left (fun m (func : Ast.funcDecl) -> StringMap.add func.name func m) _map funcDecls)


let load_object_index _map objDecls = (
  List.fold_left (fun m (obj : Ast.objDecl) -> StringMap.add obj.objName obj m) _map objDecls)



let check ( program : Ast.program ) =
  
  (* Create Environment for scope and tracking *)
  let environment = {
    function_index = load_function_index StringMap.empty program.funcDecls;
    object_index = load_object_index StringMap.empty program.objDecls;
    var_decls = load_var_decls StringMap.empty program.varDecls;
    locals = { parent = None; vars = [] } } in

  let funcChecks = List.map (check_func_decl environment) program.funcDecls in
  let stmtChecks = List.map (check_stmt environment) program.stmts in
  funcChecks;
  stmtChecks;
  Some(true)
  

  (* var Decl names *)
(*  let var_decl_names varDeclList = 
    List.map (fun (varDecl : Ast.varDecl) -> varDecl.vname) varDeclList
  in*)

  (* Load global variables into the environment *)
(*  let load_global_vars = List.fold_left (fun m obj -> StringMap.add obj true m ) 
    environment.global_index (var_decl_names program.varDecls)
  in*)

  (* Load objects into environment *)

  (* get function names *)
(*  let function_names funcDecl = List.map 
    (fun (funcDecl : Ast.funcDecl) -> funcDecl.name) funcDecl
  in*)
  
  (* Also load declarations for all functions *)
(*  let load_function_names = List.fold_left (fun m obj -> StringMap.add obj true m )
    environment.function_index (function_names program.funcDecls )
  in*)

  (* Load function declarations *)
(*  let load_function_declarations (funcDecls : Ast.funcDecl list) = 
    List.fold_left (fun li obj -> obj :: li) [] funcDecls
  in*)

  (* Need to check that objects and variable declarations are okay *)

  (* Check the functions *)
(*  let check_func_decls env funcDecls = 
    List.fold_left (check_func_decl env) (Some true) funcDecls
  in*)

  (* Answer check of all functions *)
(*  load_global_vars;
  load_function_names;
  load_function_declarations program.funcDecls;
  check_func_decls environment program.funcDecls*)
  
                  
