(** Bombs-Must-Detonate: FSM Syntax Extension Compiler @author Brian Go *)



open Ast
exception Fsm_error of string


let raise_fsm_error strlist = raise (Fsm_error (String.concat "" strlist))
let concat s = String.concat "" s



(** A node in the state graph used for connectivity checking *)


type state_node = int * int list * bool ref 

(** id, target id list, visited flag *)





(** State Graph *)


let node_list = ref []



(** State Id -> State name association *)


let state_env = ref []



(** Returns the state name given an id *)


let get_state_name id = 
  if List.mem_assoc id !state_env then 
    List.assoc id !state_env
  else
    string_of_int id



(** Reverses key/value pairs of an association list *)


let rec reverse_assoc lst = match lst with
    (x,y)::rest -> (y,x)::(reverse_assoc rest)
  | [] -> []



(** Retrieves the state_node object with the given id from the given graph *)


let rec get_state_helper node_list search_id = match node_list with
    (id,targets,visited)::rest when id = search_id -> (id,targets,visited)
  | _::rest -> get_state_helper rest search_id
  | [] -> raise_fsm_error ["State machine error: state ";
                           get_state_name search_id;
                           " has an incoming transition but is not defined."]



(** Retrieves the state_node object with the given id from the globally defined graph *)


let get_state search_id = get_state_helper !node_list search_id



(** Visit all nodes stemming from nodes in node_list *)


let rec visit_graph node_list = 
  match node_list with
      (id,targets,visited)::rest ->
        if !visited then
          visit_graph rest
        else 
          let _ = visited := true in
          let _ = visit_graph (List.map get_state targets) in
            visit_graph rest
    | [] -> ()




(** Check if the graph node_list with root(s) start_nodes is connected *)


let check_connected start_nodes node_list =
  let _ = visit_graph start_nodes in
  let rec helper node_list = match node_list with
      (id,_,visited)::rest -> 
        if !visited then 
          helper rest 
        else 
          raise_fsm_error ["State machine error: state graph is disconnected.  State ";
                           get_state_name id;
                           " is not reachable."]
    | [] -> ()
  in 
    helper node_list




(** Used by make_id() *)


let cur_id = ref (-1);;


(** Returns a unique integer every call *)


let make_id () = let _ = cur_id := !cur_id + 1 in !cur_id



(** Returns the first element of any 2-tuple *)


let rec get_first lst = match lst with
    (x,y)::rest -> x::(get_first rest)
  | [] -> []



(** Gets a string list of state names from a list of state_nodes *)


let rec get_state_names statelist = match statelist with
    (name,_,_)::rest ->
      name::(get_state_names rest)
  | [] -> []



(** Appends the first string to the second *)


let append_string app_str str = 
  concat [str;app_str]



(** Checks if the id is associated with some name in the environment. If it is, returns the name, and then unchanged environment. If it isn't, creates a new ID, adds it to the environment, and returns them *)


let get_id name env = 
  if List.mem_assoc name env then
    List.assoc name env, env
  else
    let id = make_id() in
      id, (name,id)::env



(** Turns a string into a AST variable identifier *)


let varify str = SynVarIdentifier (SynVarName str)



(** Compiles the callback list of a state *)


let rec compile_callbacks callbacks state_name = 
  match callbacks with 
      cur::rest ->
        let field, foo = 
          (match cur with
               SynFsmOnInit s -> "onInit", s 

(** Not actually a callback. Called when the state is viisted. *)


             | SynFsmOnMoveRequest s -> "onMoveRequest", s
             | SynFsmOnTeammateDeath s -> "onTeammateDeath", s
             | SynFsmOnBombDetonate s -> "onBombDetonate", s
             | SynFsmOnDeath s -> "onDeath", s)
        in
          (if field = "onInit" then 

(** Process onInit differently *)


             [ SynVarAssign
                 (SynVarAssignment
                    (SynStructOrEnumValue
                       (SynVarName state_name,
                        "init",
                        ref VIdUnspecified),
                     SynVarIdentifier
                       (SynVarName foo))) ]
           else
             [ SynVarAssign
                 (SynVarAssignment
                    ((SynStructOrEnumValue
                        (SynVarName (concat [state_name;"_callbacks"]),
                         field, ref VIdUnspecified)),
                     (SynVarIdentifier
                        (SynVarName foo)))) ])
          @ compile_callbacks rest state_name
    | [] -> []



(** Compile a transition target of a state. A transition target is a state name with a weight. If there are multiple targets then the individual weight divided by the total weight is the probability that target state is visited next. *)


let compile_transition_target (targ_name,targ_weight) transition_name env = 
  let targ_id, new_env = 
    if List.mem_assoc targ_name env then
      List.assoc targ_name env, env
    else 
      let id = make_id() in
        id, (targ_name,id)::env
  in
  let tuple_name = concat [transition_name;"_target_";string_of_int (make_id())] in
    [ 

(** declare transition target tuple *)


      SynVarDeclare
        (SynVarDeclareNoInit
           (SynEnumOrStructType
              ("fsmTransitionTarget", ref []),
            tuple_name));
      

(** set target id *)


      SynVarAssign
        (SynVarAssignment
           (SynStructOrEnumValue
              (SynVarName tuple_name,
               "id",
               ref VIdUnspecified),
            SynValue
              (SynIntValue targ_id)));
      

(** set the target weight *)


      SynVarAssign
        (SynVarAssignment
           (SynStructOrEnumValue
              (SynVarName tuple_name,
               "weight",
               ref VIdUnspecified),
            SynValue
              (SynFloatValue targ_weight)));
      

(** add tuple to transition list *)


      SynVarAssign
        (SynVarModify
           (SynStructOrEnumValue
              (SynVarName transition_name,
               "targets",
               ref VIdUnspecified),
            SynBinopCons,
            SynVarIdentifier
              (SynVarName tuple_name))) ], new_env, targ_id
           



(** Compiles all transition targets of a given state. See compile_transition_target *)


let rec compile_transition_targets targs transition_name env = match targs with
    cur::rest -> 
      let cur_result, cur_env, targ_id = compile_transition_target cur transition_name env in 
      let rest_result, rest_env, targ_ids = compile_transition_targets rest transition_name cur_env in
        cur_result @ rest_result, rest_env, targ_id::targ_ids
  | [] -> [], env, []
   



(** Compiles a given transition in a state. A transition is a list of predicates which act on the fsmStateInformation object (see lib_fsm.bmd), all of which much be satisfied to visit (non-deterministaically) one of the transition targets. *)


let compile_transition (preds,targs) state_name env = 
  let name = concat [state_name;"_transition_";string_of_int (make_id())] in
  let predvars = List.map varify preds in
  let target_result, target_env, target_ids = compile_transition_targets targs name env in
    [ 

(** declare transition *)


      SynVarDeclare
        (SynVarDeclareNoInit
           (SynEnumOrStructType ("fsmTransition",ref []),
            name));
      

(** assign predicate list *)


      SynVarAssign
        (SynVarAssignment
           (SynStructOrEnumValue
              (SynVarName name,
               "predicates",
               ref VIdUnspecified),
            (SynListValueProducer
               (SynListList predvars))))]
      

(** compile target list *)


    @ target_result
      

(** add to transition list *)


    @ [ SynVarAssign
          (SynVarModify
             (SynStructOrEnumValue
                (SynVarName state_name,
                 "transitionList",
                 ref VIdUnspecified),
              SynBinopCons,
              (SynVarIdentifier
                 (SynVarName name)))) ], target_env, target_ids




(** Compiles all transitions of a given state. See compile_transition *)

      
let rec compile_transitions transitions state_name env = 
  match transitions with
      cur::rest -> 
        let cur_result, cur_env, targ_ids = compile_transition cur state_name env in
        let rest_result, rest_env, id_list = compile_transitions rest state_name cur_env in
          cur_result @ rest_result, rest_env, targ_ids @ id_list
    | [] -> [], env, []




(** Compiles all states. A state has an initialization function called every time it is visited, a set of game event callback functions, and transitions to other states. The mechanism of state transition is described in lib_fsm.bmd. *)


let rec compile_states_helper state_list env = match state_list with
    (name, callback_list, transition_list)::rest ->
      let id,new_env = get_id name env in
      let transition_code, transition_env, transition_ids = compile_transitions transition_list name new_env in
      let _ = node_list := (id,transition_ids,ref false)::!node_list in
        [ 

(** declare state *)


          SynVarDeclare
            (SynVarDeclareNoInit
               (SynEnumOrStructType ("fsmState",ref []),
                name));
          

(** assign id *)


          SynVarAssign 
            (SynVarAssignment
               (SynStructOrEnumValue
                  (SynVarName name,"id",ref VIdUnspecified),
                SynValue
                  (SynIntValue id)));
          

(** declare callback set *)


          SynVarDeclare
            (SynVarDeclareNoInit
               (SynEnumOrStructType ("callbackSet",ref []),
                concat [name;"_callbacks"]))]
        @ 

(** compile callbacks *)


          compile_callbacks callback_list name
        @ [ 

(** assign callback set *)


          SynVarAssign
            (SynVarAssignment
               (SynStructOrEnumValue
                  (SynVarName name,"callbacks",ref VIdUnspecified),
                SynVarIdentifier
                  (SynVarName (concat [name;"_callbacks"])))) ]
        @ transition_code
        @ compile_states_helper rest transition_env
  | [] -> let _ = state_env := reverse_assoc env in []



(** Produces code to set the state list of the finite state machine to the given state names *)

           
let compile_set_statelist state_names fsm_name =
  let state_name_vars = List.map varify state_names in
    [ SynVarAssign
        (SynVarAssignment
           (SynStructOrEnumValue
              (SynVarName "ret",
               "states",
               ref VIdUnspecified),
            SynListValueProducer
              (SynListList state_name_vars))) ]



(** Compiles the state list and sets the initial state of the FSM. See compile_states_helper. *)


let compile_states state_list fsm_name =
  

(** Compile all states *)

 
  let states = compile_states_helper state_list [] in
  match state_list with 
      (init_name,_,_)::rest ->
        states
          

(** Set the initial state *)


        @ [SynVarAssign
             (SynVarAssignment
                (SynStructOrEnumValue
                   (SynVarName "ret","currentState",ref VIdUnspecified),
                 SynVarIdentifier
                   (SynVarName init_name)))]
          

(** Make the state list and set the fsm field *)


        @ (compile_set_statelist (get_state_names state_list) fsm_name)
    | [] -> raise_fsm_error ["State machine error: No states specified for state machine."]




(** Compiles each state machine described in the AST, leaving other syntax elements untouched. *)


let rec compile_fsm src_prog = 
    match src_prog with
        (SynStateMachine (name,init,state_list))::rest ->
          let _ = node_list := [] in
          let machine_code = 
            

(** Declare fsm, set the init function *)


            [SynVarDeclare 
               (SynVarDeclareNoInit 
              (SynEnumOrStructType ("finiteStateMachine",ref []),
               "ret"));
             SynVarAssign
               (SynVarAssignment 
                  (SynStructOrEnumValue 
                     (SynVarName "ret","initialize",ref VIdUnspecified),
                   SynVarIdentifier
                     (SynVarName init)))]
              

(** Compile the states *)


            @ compile_states state_list name
              

(** Return the state machine *)


            @ [SynReturnStatement 
                 (SynValueReturn 
                    (SynVarIdentifier
                       (SynVarName "ret")))]
          in
          let machine_function = (SynFunctionDefine ((SynEnumOrStructType ("finiteStateMachine",ref []), name, []), machine_code)) in
          let start_node = match List.rev !node_list with
              start::rest -> start
            | [] -> raise_fsm_error ["State mahine error: no states specified."]
          in
          let _ = check_connected [start_node] !node_list in
            (SynInclude (ref (IncludeFileName "lib_fsm.bmd")))::(machine_function::(compile_fsm rest))
      | cur::rest -> cur::(compile_fsm rest)
      | [] -> []