open Ast

module NameMap = Map.Make(struct
  type t = string
  let compare x y = Pervasives.compare x y
end)

exception ReturnException of curvet * var_decl NameMap.t


(* Main entry point: check a program *)

let rec check_layer cvs env count = 
  let locals, globals = env in
  match cvs with
  | [] -> true
  | hd :: tl -> 
      if count < 11 then
        if NameMap.mem hd locals && (NameMap.find hd locals).t = Curvet then
          check_layer tl env (count+1)
        else if NameMap.mem hd globals && (NameMap.find hd globals).t = Curvet then
          check_layer tl env (count+1)
        else false
      else 
        raise (Failure ("Too many curves in a layer!"))

let check_semantic (vars, funcs) =
  (* Put function declarations in a symbol table *)
  let func_decls = List.fold_left
      (fun funcs fdecl -> NameMap.add fdecl.fname fdecl funcs)
      NameMap.empty funcs
  in

  (* Invoke a function and return an updated global symbol table *)
  let rec call fdecl actuals globals checked =

    (* Evaluate an expression and return (value, updated environment) *)
    let rec eval env = function
        Literal(i) -> Literalt, env;
      | Dotop(var, op, e) -> 
    let (locals, globals) = env in
    if op = "getX" || op = "getY" then
      if NameMap.mem var locals then
        if (NameMap.find var locals).t = Pointt then
          if List.length e = 0 then
              Literalt, env
          else 
            raise (Failure ("Invalid parameter for "^op))
        else 
          raise (Failure (var^" is not a Point"))
      else if NameMap.mem var globals then
        if (NameMap.find var globals).t = Pointt then
          if List.length e = 0 then
              Literalt, env
          else 
            raise (Failure ("Invalid parameter for "^op))
        else 
          raise (Failure (var^" is not a Point"))
      else
        raise (Failure ("Undeclared variable "^var))
    else if op = "setX" || op = "setY" then 
      if NameMap.mem var locals then
        if (NameMap.find var locals).t = Pointt then
          if List.length e = 1 then
            let l, (locals, globals) = eval env (List.hd e) in
            if l = Literalt then
              Literalt, env
            else
              raise (Failure ("The parameter should be integer for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else 
          raise (Failure (var^" is not a Point"))
      else if NameMap.mem var globals then
        if (NameMap.find var globals).t = Pointt then
          if List.length e = 1 then
            let l, (locals, globals) = eval env (List.hd e) in
            if l = Literalt then
              Literalt, env
            else
              raise (Failure ("The parameter should be integer for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else 
          raise (Failure (var^" is not a Point"))
      else
        raise (Failure ("Undeclared variable "^var))
    else if op = "getPoint" then
      if NameMap.mem var locals then
        if (NameMap.find var locals).t = Curvet then
          if List.length e = 1 then
            let l, (locals, globals) = eval env (List.hd e) in
            if l = Literalt then
              Pointt, env
            else
              raise (Failure ("Parameter should be 0-3 for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else
          raise (Failure (var^" is not a Curve"))
      else if NameMap.mem var globals then
        if (NameMap.find var globals).t = Curvet then
          if List.length e = 1 then
            let l, (locals, globals) = eval env (List.hd e) in
            if l = Literalt then
              Pointt, env
            else
              raise (Failure ("Parameter should be 0-3 for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else
          raise (Failure (var^" is not a Curve"))
      else
        raise (Failure ("Undeclared variable "^var))
    else if op = "setPoint" then
      if NameMap.mem var locals then
        if (NameMap.find var locals).t = Curvet then
          if List.length e = 2 then
            let l1, (locals, globals) = eval env (List.hd e) in
            let l2, (locals, globals) = eval env (List.nth e 1) in 
            if l1 = Literalt then
              if l2 = Pointt then
                Literalt, env
              else 
                raise (Failure ("The 2nd parameter should be a Point for "^op))
            else
              raise (Failure ("The 1st parameter should be 0-3 for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else 
          raise (Failure (var^" is not a Curve"))
      else if NameMap.mem var globals then
        if (NameMap.find var globals).t = Curvet then
          if List.length e = 2 then
            let l1, (locals, globals) = eval env (List.hd e) in
            let l2, (locals, globals) = eval env (List.nth e 1) in 
            if l1 = Literalt then
              if l2 = Pointt then
                Literalt, env
              else 
                raise (Failure ("The 2nd parameter should be a Point for "^op))
            else
              raise (Failure ("The 1st parameter should be 0-3 for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else 
          raise (Failure (var^" is not a Curve"))
      else
        raise (Failure ("Undeclared variable "^var))
    else if op = "getCurve" then
      if NameMap.mem var locals then
        if (NameMap.find var locals).t = Layert then
          if List.length e = 1 then
            let l, (locals, globals) = eval env (List.hd e) in
            if l = Literalt then
              Curvet, env
            else
              raise (Failure ("Parameter should be integer for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else
          raise (Failure (var^" is not a Layer"))
      else if NameMap.mem var globals then
        if (NameMap.find var globals).t = Layert then
          if List.length e = 1 then
            let l, (locals, globals) = eval env (List.hd e) in
            if l = Literalt then
              Curvet, env
            else
              raise (Failure ("Parameter should be integer for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else
          raise (Failure (var^" is not a Layer"))
      else
        raise (Failure ("Undeclared variable "^var))
    else if op = "setCurve" then
      if NameMap.mem var locals then
        if (NameMap.find var locals).t = Layert then
          if List.length e = 2 then
            let l1, (locals, globals) = eval env (List.hd e) in
            let l2, (locals, globals) = eval env (List.nth e 1) in 
            if l1 = Literalt then
              if l2 = Curvet then
                Curvet, env
              else 
                raise (Failure ("The 2nd parameter should be a Curve for "^op))
            else
              raise (Failure ("The 1st parameter should be an integer for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else 
          raise (Failure (var^" is not a Layer"))
      else if NameMap.mem var globals then
        if (NameMap.find var globals).t = Layert then
          if List.length e = 2 then
            let l1, (locals, globals) = eval env (List.hd e) in
            let l2, (locals, globals) = eval env (List.nth e 1) in 
            if l1 = Literalt then
              if l2 = Curvet then
                Curvet, env
              else 
                raise (Failure ("The 2nd parameter should be a Curve for "^op))
            else
              raise (Failure ("The 1st parameter should be an integer for "^op))
          else 
            raise (Failure ("Invalid parameter for "^op))
        else 
          raise (Failure (var^" is not a Layer"))
      else
        raise (Failure ("Undeclared variable "^var))
    else if op = "getSize" then
      if NameMap.mem var locals then
        if (NameMap.find var locals).t = Layert then
          if List.length e = 0 then
              Literalt, env
          else 
            raise (Failure ("Invalid parameter for "^op))
        else
          raise (Failure (var^" is not a Layer"))
      else if NameMap.mem var globals then
        if (NameMap.find var globals).t = Layert then
          if List.length e = 0 then
              Literalt, env
          else 
            raise (Failure ("Invalid parameter for "^op))
        else
          raise (Failure (var^" is not a Layer"))
      else
        raise (Failure ("Undeclared variable "^var))
    else raise (Failure ("Undeclared dot operation " ^ op))
      | Curve(e1, e2, e3, e4, e5, e6, e7, e8) ->
    let l1, env = eval env e1 in
      let l2, env = eval env e2 in
        let l3, env = eval env e3 in
          let l4, env = eval env e4 in
            let l5, env = eval env e5 in
              let l6, env = eval env e6 in
                let l7, env = eval env e7 in
                  let l8, env = eval env e8 in
                    if List.map (fun(x) -> if x = Literalt then 1 else 0) [l1;l2;l3;l4;l5;l6;l7;l8] = [1;1;1;1;1;1;1;1] then
                      Curvet, env
                    else raise (Failure ("Invalid curve definition."))
      | Point(e1, e2) -> 
    let l1, env = eval env e1 in
      let l2, env = eval env e2 in
        if l1 = Literalt && l2 = Literalt then
          Pointt, env
        else raise (Failure ("Invalid point definition"))
      | Layer(cvs) ->
    if check_layer cvs env 0 then
      Layert, env
    else raise (Failure ("Invalid layer definition"))
      | Noexpr -> Literalt, env (* must be non-zero for the for loop predicate *)
      | Id(var) ->
    let locals, globals = env in
    if NameMap.mem var locals then
      (NameMap.find var locals).t, env
    else if NameMap.mem var globals then
      (NameMap.find var globals).t, env
    else raise (Failure ("Undeclared identifier " ^ var))
      | Binop(e1, op, e2) ->
    let l1, env = eval env e1 in
      let l2, env = eval env e2 in
        if l1 = Literalt && l1 = l2 then
          Literalt, env
        else raise (Failure ("Parameter should be integer for binary operation"))
      | Assign(var, e) ->
    let l, (locals, globals) = eval env e in
    if NameMap.mem var locals then
      if (NameMap.find var locals).t = l then
        l, (locals, globals)
      else raise (Failure ("Type mismatch for "^var))
    else if NameMap.mem var globals then
      if (NameMap.find var globals).t = l then
        l, (locals, globals)
      else raise (Failure ("Type mismatch for "^var))
    else raise (Failure ("Undeclared identifier " ^ var))
      | Call("draw", e) ->
    if List.length e = 1 then
    let l, env = eval env (List.hd e) in
      if l = Layert then
        (Literalt, env)
      else raise (Failure ("The parameter of draw() function should be Layer"))
    else raise (Failure ("Invalid parameter for draw()"))
      | Call("print", e) ->
    if List.length e = 1 then
    let l, env = eval env (List.hd e) in
      if l = Literalt then
        (Literalt, env)
      else raise (Failure ("The parameter of print() function should be integer"))
    else raise (Failure ("Invalid parameter for print()"))
      | Call("pause", e) ->
    if List.length e = 1 then
    let l, env = eval env (List.hd e) in
      if l = Literalt then
        (Literalt, env)
      else raise (Failure ("The parameter of pause() function should be integer"))
    else raise (Failure ("Invalid parameter for pause()"))
      | Call("random", e) ->
    if List.length e = 1 then
    let l, env = eval env (List.hd e) in
      if l = Literalt then
        (Literalt, env)
      else raise (Failure ("The parameter of random() function should be integer"))
    else raise (Failure ("Invalid parameter for random()"))
      | Call("clear", e) ->
    if List.length e = 0 then
        (Literalt, env)
    else raise (Failure ("Invalid parameter for clear()"))
      | Call(f, actuals) -> 
    let fdecl =
      try NameMap.find f func_decls;
      with Not_found -> raise (Failure ("Undefined function " ^ f))
    in
    let actuals, env = List.fold_left
        (fun (actuals, env) actual ->
    let v, env = eval env actual in v :: actuals, env)
          ([], env) (List.rev actuals)
    in
    let (locals, globals) = env in
    try
      if NameMap.mem f checked then
        (NameMap.find f func_decls).return, (locals, globals)
      else
        let globals = call fdecl actuals globals (NameMap.add f 1 checked)
        in fdecl.return, (locals, globals)
    with ReturnException(v, globals) -> 
      if v = fdecl.return then
        v, (locals, globals)
      else 
        raise (Failure ("Return type mismatch for function "^f))
    in

    (* Execute a statement and return an updated environment *)
    let rec exec env = function
  Block(stmts) -> List.fold_left exec env stmts
      | Expr(e) -> let _, env = eval env e in env
      | If(e, s1, s2) ->
    let l, env = eval env e in
      if l = Literalt then
        List.fold_left exec env [s1;s2]
      else raise (Failure ("Invalid conditional statement."))
      | While(e, s) ->
    let l, env = eval env e in
      if l = Literalt then
        exec env s
      else raise (Failure ("Invalid conditional statement."))
      | For(e1, e2, e3, s) ->
    let _, env = eval env e1 in
      let l, env = eval env e2 in
        if l = Literalt then
          let _, env = eval (exec env s) e3 in
          exec env s
        else raise (Failure ("Invalid conditional statement."))
      | Return(e) ->
    let v, (locals, globals) = eval env e in
    raise (ReturnException(v, globals))
    in
    
    (* Enter the function: bind actual values to formal arguments *)
    let locals =
      try List.fold_left2
    (fun locals formal actual -> if formal.t = actual then 
                                    NameMap.add formal.name { t = formal.t; name = formal.name; value = [0] } locals 
                                 else raise (Failure ("Type mismatch for parameter " ^ formal.name ^ " in function " ^ fdecl.fname)))
    NameMap.empty fdecl.formals actuals
      with Invalid_argument(_) ->
  raise (Failure ("wrong number of arguments passed to " ^ fdecl.fname))
    in
    let locals = List.fold_left
  (fun locals local -> NameMap.add local.name { t = local.t; name = local.name; value = local.value } locals) locals fdecl.locals
    in
    (* Execute each statement in sequence, return updated global symbol table *)
    let env = List.fold_left exec (locals, globals) fdecl.body
    in 
    if fdecl.fname = "main" then snd env
    else raise (Failure ("No return statement for function "^fdecl.fname))

  (* Run a program: initialize global variables to 0, find and run "main" *)
  in let globals = List.fold_left
      (fun globals vdecl -> NameMap.add vdecl.name { t = vdecl.t; name = vdecl.name; value = vdecl.value } globals) NameMap.empty vars
  in try
    ignore (call (NameMap.find "main" func_decls) [] globals NameMap.empty)
  with Not_found -> raise (Failure ("did not find the main() function"))
