(** 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)
      | [] -> []