open Ast
module StringMap = Map.Make(String)

(* Function Type *)
type functionType =
  Constructor of string
| TableConstructor of string * int (* tableName, columnSize *)
| FuncCall

(* Var Info *)
type varInfo =
   RegVar
 | ArrayVar

(* Translation Environment *)
type environment = {
  function_index : functionType StringMap.t; 
  var_index : (progType * varInfo * int) StringMap.t;
  method_index : (string * string) list;
  type_list: (string * int) list;
  csv_class: string option;
  }

(* Handle Binary Operation *)
let string_op = function
    Add -> "+" 
  | Sub -> "-" 
  | Mult -> "*" 
  | Div -> "/" 
  | Mod -> "%" 
  | Equal -> "==" 
  | Neq -> "!=" 
  | Less -> "<" 
  | Leq -> "<=" 
  | Greater -> ">" 
  | Geq -> ">=" 

let rec get_type env = function
    Literal(v) -> (Literal(v), Int )
  | LitFloat(v) -> (LitFloat(v), Float )
  | LitBool(v) -> (LitBool(v), Boolean )
  | LitString(v) -> (LitString(v), String )
  | Id(name) -> (
    let (t,v,num) = 
     (try 
          StringMap.find name env.var_index
       with Not_found -> raise( Failure "Variable Not Found")) in
    (Id(name),t))
  | Binop(e1,op,e2) -> get_type env e1 
  | Assign(eVar,e) -> (
     let (_,t) = get_type env e in
     Assign(eVar,e),t)
  | Call( name, expr_li ) -> (* TODO Check function table for type *)
       Call(name,expr_li), Int
  | IfThenElse( e1, e2, e3, sl1, sl2 ) ->
     (* Get Types *)
     let (_,t1) = get_type env e1 in
     IfThenElse( e1, e2, e3, sl1, sl2), t1
   | MemberAccess(name,field) -> MemberAccess(name,field),Int (* Lookup type *)
   | ArrayAccess(name,e) -> ArrayAccess(name,e),Int (* Lookup type *)
   | MethodCall(name,field,args) -> MethodCall(name,field,args),Int (* Lookup type *)
   | CSVInput -> CSVInput,Int (* No Type here *)
   | Noexpr -> raise( Failure "No expression here")

(* Handle Expressions *)
let rec expr env = function
    Literal(l) -> string_of_int l
  | LitFloat(l) -> string_of_float l
  | LitBool(l) -> string_of_bool l
  | LitString(l) -> l
  | Id(s) -> s
  | Binop(e1,op,e2) -> "(" ^ expr env e1 ^ ")" ^ string_op op ^ "(" ^ expr env e2 ^ ")"
  | Assign(var,exp) -> "(" ^ expr env var ^ "=" ^ expr env exp ^ ")"
  | Call( name, args ) -> if name = "Print" then (
          match args with
            [] -> "printf(\"\")"
          | hd::tl -> (
              let (_,t) = get_type env hd in
              match t with
                Int -> "printf(\"%d\", " ^ expr env hd ^ " )"
              | Float -> "printf(\"%f\", " ^ expr env hd ^ " )"
              | Boolean -> "printf(\"%s\", (" ^ expr env hd ^ "==true ? \"true\" : \"false\"))"
              | String -> "printf(\"%s\", " ^ expr env hd ^ " )"
              | _ -> "printf( \"Cannot print this type\"\\n)"))
          else (name ^ "(" ^ (String.concat "," (List.map (expr env) args)) ^ ")")
  | MemberAccess(name,field) -> name ^ "->" ^ field
  | MethodCall(name,field,args) -> if field = "print" 
                                   then "__TablePrint( " ^ name ^ ")"
                                   else "Unimplemented()"
  | CSVInput -> ""
  | ArrayAccess(name,e) -> (
    let (t,varT,num) =
    (try 
        StringMap.find name env.var_index
     with Not_found -> (Int,RegVar,100)) in
    match t with
      Record(oName) -> "&" ^ name ^ "[" ^ expr env e ^ "]"
    | _ -> name ^ "[" ^ expr env e ^ "]" )
  | IfThenElse(pred,e2,e3,sl1,sl2) -> "(" ^ expr env pred ^ ") ? (" ^ expr env e2 ^ ") : (" ^ expr env e3 ^ ")"
  | Noexpr -> ""

(* Handle list of statements *)
let rec stmt env = function
    Block sl -> String.concat " " (List.map (stmt env) sl)
  | Expr e -> expr env e ^ ";"  
  | For (e1,e2,e3,sl) -> "for(" ^ expr env e1 ^ " ; " ^ expr env e2 ^ " ; " ^ expr env e3 ^ ")" ^
    "{ " ^ (String.concat " " (List.map (stmt env) sl)) ^ "}"
  | Foreach(s1,sl) -> "Unimplemented"
  | Return(e) -> "return " ^ expr env e ^ ";"

let translate_c_type t =
  match t with
    Ast.Int -> "int"
  | Ast.String -> "char *"
  | Ast.Boolean -> "Boolean"
  | Ast.Float -> "double"
  | Ast.Record(name) -> name ^ "* "
  | Ast.Table(name) -> "__Table_ * "

let translate_c_type_array t =
  match t with
    Ast.Int -> "int"
  | Ast.String -> "char *"
  | Ast.Boolean -> "Boolean"
  | Ast.Float -> "double"
  | Ast.Record(name) -> name 
  | Ast.Table(name) -> "__Table_ * "

let translate_c_type_id = function
    Ast.Int -> "__ProgramType_Int"
  | Ast.Float -> "__ProgramType_Float"
  | Ast.Boolean-> "__ProgramType_Boolean"
  | Ast.String -> "__ProgramType_String"
  | Ast.Record(str) -> "__ProgramType_" ^ str
  | Ast.Table(str) -> "Unimplemented"

let translate_c_type_id_for name = "__ProgramType_" ^ name;;

let translate_object_constructor oName = 
  "(" ^ oName ^ " *) ObjectNew( __ProgramType_" ^ oName ^ " )"

let translate_table_constructor name tblColSize dataClass dataArrName dataArrSize =
  "__TableNew( \"" ^ name ^ "\", TABLE_GETITEM( " ^ name ^ " ),
    TABLE_COLUMNS( " ^ name ^ "), " ^ string_of_int tblColSize ^
    ", " ^ translate_c_type_id_for dataClass ^ ", " ^ dataArrSize ^ ", " ^
    dataArrName ^ " )"

let translate_table_constructor_single name tName colSize =
  "__TableNew( \"" ^ name ^ "\", TABLE_GETITEM( " ^ name ^ " ),
    TABLE_COLUMNS( " ^ name ^ " ), " ^ string_of_int colSize ^
    ", 0 , 0, NULL )"

let translate_varDecl env varDecl = 
  let csvClass = function
    Some(nm) -> nm
    |  _ -> " " in
  let name = translate_c_type varDecl.varType ^ " " ^ varDecl.vname in
  let name2 = translate_c_type_array varDecl.varType ^ " " ^ varDecl.vname in
  let funcMap = env.function_index in 
  let nonArrayAssign env = function
      Some(hd::tl) -> (match hd with
         Call( name, headExpr::tailExpr ) -> (let callType = 
           try StringMap.find name funcMap
           with Not_found -> FuncCall in
           match callType with
             Constructor(oName) -> translate_object_constructor name
           | TableConstructor(tName,colLen) -> (
              match headExpr with 
                CSVInput -> translate_table_constructor 
                  name colLen (csvClass env.csv_class) "dataArray" "__calls"
              | Id(str) -> translate_table_constructor 
                  name colLen "Int" str "10" (* TODO GET ARRAY ITEM COUNT *)
              | _ -> raise( Failure "Bad Table Constructor Call" ))
           | FuncCall -> expr env hd)
         | _ -> "=" ^ expr env hd )
    | _ -> "" in
  let assignable = 
    match varDecl.varArraySize with
      None -> name ^ if varDecl.varAssign = None then ""
        else nonArrayAssign env varDecl.varAssign
    | Some(size) -> name2 ^
       match varDecl.varAssign with
         None -> ""
       | Some(li) -> "[" ^ string_of_int size ^ "] = {" ^ 
         String.concat "," (List.map (expr env) li) ^ "}"
  in
  assignable ^ ";"

let translate_funcDecl env funcDecl =
 
  (* Handle formals *)
  let string_of_formal formal =
    translate_c_type formal.varType ^ " " ^ formal.vname
  in

  (* Print Function *)
  "static " ^ translate_c_type funcDecl.retType ^ " " ^
    funcDecl.name ^ "(" ^ (String.concat "," 
    (List.map string_of_formal funcDecl.formals)) ^ " ) {\n" 
    ^ (String.concat ";\n" (List.map (translate_varDecl env) funcDecl.locals))
    ^ (String.concat ";\n" (List.map (stmt env) funcDecl.body)) ^ "}\n"
  (* Function End *)

let translate_tableDecl env tableDecl = 
  
  (* C Table Function Return Type *)
  let c_table_ret_type = "__TableReturn_ " in

  (* C Column Function Name *)
  let c_column_decl_name name = "__" ^ tableDecl.tableName ^ 
    "__" ^ name in
  
  (* C Column Function Decl *)
  let c_method_decl name = c_table_ret_type ^ c_column_decl_name name ^  
    "( void * input )\n{\n" in

  (* C Get Item Function *)
  let c_getItem_decl = "void * _" ^ tableDecl.tableName ^ "_getItem_" ^
    "( void * data, int index )\n{\n" ^
    tableDecl.iterClass ^ "* d;\n" ^
    tableDecl.iterClass ^ "* item;\n\n" ^
    "d = (" ^ tableDecl.iterClass ^ "*) data;\n" ^
    "item = &d[index];\n" ^
    "return item;\n}\n" in

  (* Column Function translate *)
  let translate_table_column_func env col retType = 
    let retVal = (expr env) col.lastExpr in
    let func_ret_code = 
      match retType with
        Int -> "__ret_._int = " ^ retVal ^ ";\n"
      | Float -> "__ret_._float = " ^ retVal ^ ";\n"
      | String -> "__ret_._str = " ^ retVal ^ ";\n"
      | Boolean -> "__ret_._boolean = " ^ retVal ^ ";\n"
      | Table(s) -> "UNIMPLEMENTED" 
      | Record(s) -> "UNIMPLEMENTED" in
    (String.concat ";\n" (List.map (translate_varDecl env) col.locals))
    ^ (String.concat ";\n" (List.map (stmt env) col.statements)) ^ 
    func_ret_code ^ "return __ret_;" ^ ";\n}\n" in

  (* C Table Function Preface *)
  let c_table_func_preface = c_table_ret_type ^ "__ret_;\n" ^
    tableDecl.iterClass ^ " * item = (" ^ tableDecl.iterClass ^ " *) input;\n" in

  (* Translate Column *)
  let translate_column env col = c_method_decl col.tblColName ^
    c_table_func_preface ^
    translate_table_column_func env col.tblColFunc col.tblColType in

  (* C print Column Functions *)
  let c_column_functions = String.concat " " 
    (List.map (translate_column env) tableDecl.columns) in

  (* C Column item for list *)
  let c_column_Item col = "TABLE_COLUMN( \"" ^ col.tblColName ^ "\", " ^
    translate_c_type_id col.tblColType ^ ", " ^ c_column_decl_name col.tblColName ^
    " )" in

  (* Print Column Info *)
  let c_column_list = "TABLE_START( " ^ tableDecl.tableName ^ ")\n" ^
    String.concat "\n" (List.map c_column_Item tableDecl.columns) ^ "\n" ^
    "TABLE_END( " ^ tableDecl.tableName ^ ")\n" in

  (* Print *)
  c_getItem_decl ^ c_column_functions ^ c_column_list

let translate_objectDecl env obj =
  let c_prelude = "typedef struct " ^ obj.objName ^ " {\n OBJECT_HEADER \n" in
  let field_string (name,t) = translate_c_type t ^ " " ^ name in
  let fields = String.concat ";\n" (List.map field_string obj.objFields) in
  let c_end = ";\n} " ^ obj.objName ^ ";\n" in
  c_prelude ^ fields ^ c_end

let translate_main env varDecl_list stmt_list = 
  let start = "int main( int argc, char ** argv ) {\n" in
  let filter_only_assignables li = 
    List.filter (fun varDecl -> 
      if varDecl.varArraySize != None || varDecl.varAssign != None 
      then true else false) li in
  let varDecl_string = String.concat "\n" (List.map (translate_varDecl env)
    (filter_only_assignables varDecl_list)) in
  let stmts_string = String.concat "\n" (List.map (stmt env) stmt_list ) in
  start ^ varDecl_string ^ stmts_string ^ "\n}"

let load_constructors p _map = 
  let _map2 = List.fold_left 
    (fun m objD -> StringMap.add objD.objName (Constructor(objD.objName)) m ) _map p.objDecls in
  List.fold_left 
   (fun m tableD -> StringMap.add tableD.tableName (TableConstructor(tableD.tableName,List.length tableD.columns)) m )
    _map2 p.tableDecls

let filter_unassigned varDecls_li = 
  List.filter (fun varDecl -> if varDecl.varArraySize = None && varDecl.varAssign = None
  then true else false ) varDecls_li

(* List of counters *)
let rec inc_list = function
    0 -> []
  | num -> num :: inc_list (num-1)

(* Zip two lists *)
let rec zip l1 l2 = 
  match l1,l2 with
  | [],_ -> []
  | _, []-> []
  | (x::xs),(y::ys) -> (x,y) :: (zip xs ys)

let load_type_list objDecls = 
  let basic_types = [ ("Boolean",4); ("String",3); ("Float",2); ("Int",1) ] in
  let obj_len = List.length objDecls in
  let nums = inc_list (obj_len + 4) in
  let new_types = List.map (fun obj -> obj.objName) objDecls in
  let items = zip new_types nums in
  items @ basic_types

let print_program_typeInfo = "
typedef struct __ProgramTypeInfo_ {
  char * name;
  __ProgramType_ type;
  int size;
} __ProgramTypeInfo_;\n\n"

let translate_types_enum t_list =
  let preface = "typedef enum __ProgramType_ {\n" in
  let translate_type (t,num) = "__ProgramType_" ^ t ^ " = " ^ 
    string_of_int num in
  let epilogue = "} __ProgramType_;\n\n" in 
  preface ^ String.concat ",\n" (List.map translate_type t_list) ^ epilogue

let translate_type_table t_list =
  let preface = "__ProgramTypeInfo_ __programTypeList_[] = { \n" in
  let first_item = " { \"None\", 0, 0 },\n" in
  let translate_type_item (t,num) = " { \"" ^ t ^ "\" , " ^ "__ProgramType_" ^ 
    t ^ ", sizeof(" ^ t ^ ") }" in
  let epilogue = "\n};\n\n" in
  preface ^ first_item ^ String.concat ",\n" (List.map translate_type_item t_list) ^
  epilogue

let translate_varDecl2_globals env varDecl_list =  
  let name varDecl = translate_c_type varDecl.varType ^ " " ^ varDecl.vname in
  let name2 varDecl = translate_c_type_array varDecl.varType ^ " " ^ varDecl.vname in
(*  let assignable = List.filter (fun varDecl -> varDecl.varArraySize != None ||
    varDecl.varAssign != None) varDecl_list in*)
  let nonArrayAssign env = function
      Some(hd::tl) -> (match hd with
         Call( name, args ) -> " " (*(let callType = 
           try StringMap.find name funcMap
           with Not_found -> FuncCall in
           match callType with
             Constructor -> translate_object_constructor name
           | TableConstructor -> translate_table_constructor_single name
           | FuncCall -> expr hd)*)
         | _ -> "=" ^ expr env hd )
    | _ -> "" in
  let translate2 env varDecl = 
    match varDecl.varArraySize with
      None -> name varDecl ^ (if varDecl.varAssign = None then ""
        else nonArrayAssign env varDecl.varAssign)
    | Some(size) -> name2 varDecl ^
       match varDecl.varAssign with
         None -> "[" ^ string_of_int size ^ "]"
       | Some(li) -> "[" ^ string_of_int size ^ "] = {" ^ 
         String.concat "," (List.map (expr env) li) ^ "}"
  in String.concat ";\n" (List.map (translate2 env) varDecl_list) ^ ";\n"


let translate_varDecl2_assignments env varDecl_list =
  let csvClass = function
    Some(nm) -> nm
    |  _ -> " " in
  let assignable = List.filter (fun varDecl -> varDecl.varArraySize == None &&
    varDecl.varAssign != None) varDecl_list in
  let funcMap = env.function_index in 
  let nonArrayAssign env vname = function
      Some(hd::tl) -> (match hd with
         Call( name, args ) -> (let callType = 
           try StringMap.find name funcMap
           with Not_found -> FuncCall in
           match callType with
             Constructor(oName) -> vname ^ " = " ^ translate_object_constructor name
           | TableConstructor(tName,colSize) -> 
               (match args with
                  [] -> " "
                | car::cadr -> 
                   (match car with 
                     CSVInput ->  vname ^ " = " ^ translate_table_constructor name colSize (csvClass env.csv_class) "dataArray" "__calls"
                    | Id(nm) -> (
                      let (t,varT,num) =
                      (try 
                          StringMap.find nm env.var_index
                       with Not_found -> (Int,RegVar,100)) in
                      let cls = function
                        Record(n) -> n
                        | _ -> raise (Failure "Bad Array Access" ) in
                        vname ^ " = " ^ translate_table_constructor name colSize (cls t) nm (string_of_int num) )
                    | _ -> ""))
           | FuncCall -> "")
         | _ -> "" )
    | _ -> "" in
  let translate2 env varDecl = 
    match varDecl.varArraySize with
      None -> (if varDecl.varAssign = None then ""
        else nonArrayAssign env varDecl.vname varDecl.varAssign)
    | _ -> " "
  in String.concat ";\n" (List.map (translate2 env) assignable) ^ ";\n"

(*************************************************************************)
(* L O A D   I N F O                                                     *)
(*************************************************************************)

let load_var_decls varDecls _map = 
  let vtype = function
      None -> RegVar
    | _ -> ArrayVar in
  let vNum = function
      Some(n) -> n
    | _ -> 0 in
  List.fold_left (fun m v -> StringMap.add v.vname (v.varType, vtype v.varArraySize, vNum v.varArraySize) m) _map varDecls
  
(*************************************************************************)
(* C S V    P A R S I N G                                                *)
(*************************************************************************)

let rec zero_counter = function
    0 -> []
  | num -> (num-1)::(zero_counter (num-1))

let string_to_field = function
    Int -> "atoi( (char *) s )"
  | Float -> "atof( (char *) s )"
  | String -> "newString( (char *) s )"
  | Boolean -> "strcmp( (char *), \"true\" ) == 0 ? true : false"
  | Record(name) -> raise( Failure "Can't parse CSV object if it has an Object field")
  | Table(name) -> raise( Failure "Can't parse CSV object if it has an Table field")
  
let translate_csv_field_loader objDecl =
  let preface = "void __" ^ objDecl.objName ^ 
    "_cb1( void * s, size_t i, void * p )\n{\n" in
  let locals = objDecl.objName ^ " * arr;\n" ^ objDecl.objName ^ " * item;\n\n" in
  let ptrs = "arr = (" ^ objDecl.objName ^ " *) p;\nitem = &arr[ __calls ];\n" in
  let fieldLength = List.length objDecl.objFields in
  let counts = zero_counter fieldLength in
  let fields_counts = zip objDecl.objFields (List.rev counts) in
  let print_field_transform ((name,t),num) = 
    if num = 0 then "if( item->fields == 0 )\n  item->" ^ name ^ " = " 
      ^ string_to_field t ^ ";\n" 
    else "else if( item->fields == " ^ string_of_int num ^ " )\n item->" ^ name ^ " = "
      ^ string_to_field t ^ ";\n" in
  let epilogue = "item->fields++;\n}\n" in
  preface ^ locals ^ ptrs ^ String.concat "" 
    (List.map print_field_transform fields_counts) ^ epilogue

let translate_csv_loader env objDecls =
  let objDeclFind li className = 
    try
      List.find (fun obj -> obj.objName = className) li
    with Not_found -> raise( Failure "No Object definition for CSV class" ) in
  match env.csv_class with
    Some(cl) -> let objDecl = objDeclFind objDecls cl in
                translate_csv_field_loader objDecl
  | _ -> "";;

let print_csv_input = "if( argc > 2 ) {\n  if( strcmp( argv[1], \"-input\" ) == 0)\n" ^
  "   fileName = argv[2];\n }\n"

let load_csv_data csvClass = 
  let loadDataArray cl = "dataArray = ArrayNew( __ProgramType_" ^ cl ^ ", 1000 );\n" in
   match csvClass with
   Some(cl) -> loadDataArray cl ^ print_csv_input ^ "LoadCSVFile( fileName, __" ^ cl ^ "_cb1,  cb2, dataArray );\n"
  | _ -> "";;

let print_csv_preface = "char * fileName = NULL;\n";;

let get_csv_class = function
    [] -> None
  | hd::tl -> Some(hd.csvInputClass)         

(*************************************************************************)
(* P R I N T   M E T H O D S                                             *)
(*************************************************************************)

let load_method_list tblDecls = 
  List.fold_left (fun li tbl -> (tbl.tableName,"print") :: li) [] tblDecls
                                  
(*************************************************************************)
(* M A I N                                                               *)
(*************************************************************************)

let translate_main2 env varDecl_list stmt_list = 
  let start = "int main( int argc, char ** argv ) {\n" in  
  let varDecl_string = translate_varDecl2_assignments env varDecl_list in
  let stmts_string = String.concat "\n" (List.map (stmt env) stmt_list ) in
  start ^ print_csv_preface ^ load_csv_data env.csv_class ^ varDecl_string ^ stmts_string ^ "\n}"

let translate_program p = 
  let env = { function_index = load_constructors p StringMap.empty;
              var_index = load_var_decls p.varDecls StringMap.empty;
              method_index = load_method_list p.tableDecls;
              type_list = load_type_list p.objDecls;
              csv_class = get_csv_class p.csvDecls } in

  (* Includes *)
  let c_prelude = "#include \"base.h\"\n" in
  let c_table_include = "#include \"table.c\"\n\n" in

  (* Build *)
  c_prelude ^ translate_types_enum (List.rev env.type_list) ^ 
  String.concat "\n" (List.map (translate_objectDecl env) p.objDecls) ^ "\n" ^
  print_program_typeInfo ^ translate_type_table (List.rev env.type_list) ^
  c_table_include ^
  translate_varDecl2_globals env p.varDecls ^ "\n" ^ 
  String.concat " " (List.map (translate_tableDecl env) p.tableDecls) ^ "\n" ^

  (* Print functions for CSV data *)    
  translate_csv_loader env p.objDecls ^ 

  String.concat " " (List.map (translate_funcDecl env) p.funcDecls) ^ "\n" ^
(*  String.concat "\n" (List.map (translate_varDecl env) (filter_unassigned p.varDecls)) ^ "\n" ^ *)
  translate_main2 env p.varDecls (List.rev p.stmts)
  
            
