import json
import ast
import inspect
import logging # Import logging
from typing import List, Dict, Any, Union, Optional, Tuple

from reward_kit.typed_interface import reward_function
from reward_kit.models import Message, EvaluateResult, MetricResult
from reward_kit.agent.resources.bfcl_sim_api_resource import BFCLSimAPIResource

# Get logger for bfcl_reward
logger = logging.getLogger("bfcl_reward")
logger.setLevel(logging.DEBUG) # Ensure debug logs are processed

# Helper functions adapted from BfclRubric (copied here)
def _parse_function_call(func_call_str: str):
    """
    Parses a function call string into a JSON-like dictionary.

    :param func_call_str: String representation of a function call.
    :return: JSON-like dictionary with function name and arguments.
    """
    try:
        tree = ast.parse(func_call_str, mode='eval')
        if not isinstance(tree.body, ast.Call):
            raise ValueError("Input is not a valid function call.")
        func_name = tree.body.func.id if isinstance(tree.body.func, ast.Name) else None
        if not func_name:
            raise ValueError("Could not determine function name.")
        args_dict = {}
        for kw in tree.body.keywords:
            args_dict[kw.arg] = ast.literal_eval(kw.value)
        for i, arg in enumerate(tree.body.args):
            args_dict[f"pos_arg_{i}"] = ast.literal_eval(arg) # Standardized positional arg key

        json_obj = {"name": func_name, "args": args_dict}
        return json_obj
    except Exception:
        raise ValueError(f"Error parsing function call string: {func_call_str}")

def _are_function_calls_equivalent(call1: Dict[str, Any], call2: Dict[str, Any]) -> bool:
    """
    Compares two parsed function call dictionaries for semantic equivalence.
    Special handling for 'sort' command arguments.
    """
    if not isinstance(call1, dict) or not isinstance(call2, dict):
        logger.warning(f"Invalid input to _are_function_calls_equivalent: call1={call1}, call2={call2}")
        return False 
    
    name1 = call1.get("name")
    name2 = call2.get("name")

    if name1 != name2:
        return False

    args1 = call1.get("args", {})
    args2 = call2.get("args", {})

    if not isinstance(args1, dict) or not isinstance(args2, dict):
        logger.warning(f"Invalid args in _are_function_calls_equivalent: args1={args1}, args2={args2}")
        return False # Should be dicts

    if name1 == "sort":
        val1 = args1.get("pos_arg_0", args1.get("file_name"))
        val2 = args2.get("pos_arg_0", args2.get("file_name"))
        
        # Check if both extracted values are not None and are equal
        # And that both arg dicts have only one relevant key for the sort value
        # (to avoid matching if one has extra unrecognized args beyond the file name)
        
        # Condition 1: Both args have exactly one key
        cond1_holds = len(args1) == 1 and len(args2) == 1
        # Condition 2: The values associated with the (potentially different) keys are the same
        cond2_holds = val1 is not None and val1 == val2

        if cond1_holds and cond2_holds:
            return True
        # Fallback to direct equality if the specific sort logic doesn't cleanly apply
        # (e.g. if one has pos_arg_0 and the other has file_name AND other args, or different number of args)
        return args1 == args2
    else:
        # For all other functions, use direct argument dictionary equality
        return args1 == args2

def _is_subsequence_unordered(list1: List[Dict[str, Any]], list2: List[Dict[str, Any]]) -> tuple[bool, list]:
    """
    Checks if all elements of list1 are present in list2, using _are_function_calls_equivalent for comparison.
    Also returns the elements of list1 that are not present in list2.
    """
    if not list1: # If list1 is empty, it's always a subsequence
        return True, []
    # If list1 is not empty but list2 is, list1 cannot be a subsequence.
    # This also handles the case where list2 becomes empty during processing.
    if not list2 and list1: 
        return False, list1[:] 
    
    list2_copy = list2[:] # Make a copy to modify
    missing_elements = []

    for item1 in list1:
        found_match_in_list2 = False
        for i, item2 in enumerate(list2_copy):
            if _are_function_calls_equivalent(item1, item2):
                list2_copy.pop(i) # Remove the matched element from list2_copy
                found_match_in_list2 = True
                break # Move to the next item in list1
        if not found_match_in_list2:
            missing_elements.append(item1)
            
    is_subsequence = len(missing_elements) == 0
    return is_subsequence, missing_elements

def compare_comparable_states(model_state: Dict[str, Any], gt_state: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
    """
    Compares two comparable state dictionaries.
    Returns True if they match, False otherwise, and a dictionary of differences.
    """
    if set(model_state.keys()) != set(gt_state.keys()):
        return False, {"keys_mismatch": {"model_keys": list(model_state.keys()), "gt_keys": list(gt_state.keys())}}

    all_match = True
    differences = {}
    for class_name in model_state.keys():
        model_instance_state = model_state[class_name]
        gt_instance_state = gt_state[class_name]

        if model_instance_state != gt_instance_state:
            all_match = False
            diffs = {k: (model_instance_state.get(k), gt_instance_state.get(k)) for k in set(model_instance_state) | set(gt_instance_state) if model_instance_state.get(k) != gt_instance_state.get(k)}
            differences[class_name] = diffs

    return all_match, differences

@reward_function
def bfcl_reward(
    messages: List[Message],      # Full conversation, assistant responses are at the end
    ground_truth: Dict[str, Any], # Contains 'function_calls' and 'comparable_state'
    state: Dict[str, Any],        # Runtime state (BFCLSimAPIResource, successful_func_calls)
    **kwargs: Any
) -> EvaluateResult:
    """
    Evaluates agent performance on BFCL tasks based on state, function calls, and format.
    """
    ground_truth_function_calls: Optional[List[List[str]]] = ground_truth.get("function_calls")
    ground_truth_comparable_state: Optional[Dict[str, Any]] = ground_truth.get("comparable_state")

    # Log ground truth data received
    logger.debug(f"Ground truth function calls from input: {ground_truth_function_calls}")
    logger.debug(f"Ground truth comparable state from input: {ground_truth_comparable_state}")

    if ground_truth_function_calls is None or ground_truth_comparable_state is None:
        return EvaluateResult(score=0.0, reason="Ground truth 'function_calls' or 'comparable_state' not found in ground_truth dict.", metrics={})

    # Access the BFCLSimAPIResource instance from the state
    bfcl_resource: Optional[BFCLSimAPIResource] = state.get("resource")

    if not isinstance(bfcl_resource, BFCLSimAPIResource):
         return EvaluateResult(score=0.0, reason="BFCLSimAPIResource instance not found in state.", metrics={})

    # --- State Matches Check ---
    model_comparable_state = bfcl_resource.get_comparable_state()
    state_match, state_diffs = compare_comparable_states(model_comparable_state, ground_truth_comparable_state)

    state_match_score = 0.5 if state_match else 0.0
    # num_state_matches = int(state_match) # Not used
    # num_state_total = 1 if ground_truth_comparable_state else 0 # Not used

    # --- Function Call Matches Check ---
    # model_successful_func_calls is List[List[Dict[str, Any]]], one inner list per user turn's accumulated calls
    model_successful_func_calls_per_turn = state.get("successful_func_calls", []) 
    
    num_func_matches_for_score = 0  # Number of user turns where model's calls matched GT's calls for that turn
    func_match_score = 0.0

    num_gt_turns_with_calls = len(ground_truth_function_calls) if ground_truth_function_calls else 0
    num_model_turns_with_actual_calls = len(model_successful_func_calls_per_turn)

    # Iterate over GT turns to see if the model matched them
    # This handles cases where model makes fewer turns with calls than GT expects.
    if num_gt_turns_with_calls > 0:
        for i in range(num_gt_turns_with_calls):
            gt_calls_str_for_this_turn = ground_truth_function_calls[i] # List[str]
            
            model_calls_for_this_turn = [] # List[Dict]
            if i < num_model_turns_with_actual_calls:
                model_calls_for_this_turn = model_successful_func_calls_per_turn[i]
            
            try:
                gt_calls_for_this_turn = [_parse_function_call(call_str) for call_str in gt_calls_str_for_this_turn]
                logger.debug(f"GT calls for turn {i}: {json.dumps(gt_calls_for_this_turn)}")
                logger.debug(f"Model calls for turn {i}: {json.dumps(model_calls_for_this_turn)}")
                
                is_match_for_turn, missing_gt_calls = _is_subsequence_unordered(gt_calls_for_this_turn, model_calls_for_this_turn)
                if is_match_for_turn:
                    num_func_matches_for_score += 1
                    logger.debug(f"Turn {i} matched.")
                else:
                    logger.debug(f"Turn {i} did NOT match. Missing GT calls in model's calls: {json.dumps(missing_gt_calls)}")
            except Exception as e:
                logger.error(f"Error comparing function calls for GT turn index {i}: {e}")
        
        if num_func_matches_for_score == num_gt_turns_with_calls and num_model_turns_with_actual_calls == num_gt_turns_with_calls:
            func_match_score = 0.5
        elif num_func_matches_for_score == num_gt_turns_with_calls and num_model_turns_with_actual_calls != num_gt_turns_with_calls:
             func_match_score = 0.0 
        else:
            func_match_score = 0.0 

    elif num_gt_turns_with_calls == 0: 
        if num_model_turns_with_actual_calls == 0: 
            func_match_score = 0.5 
        else: 
            func_match_score = 0.0
    
    reason_num_total_gt_turns_with_calls = num_gt_turns_with_calls if num_gt_turns_with_calls > 0 else "0 (no GT calls expected)"

    # --- Format Check (on model's response messages from the `messages` list) ---
    format_score = 0.2 
    valid_assistant_messages = 0 
    total_assistant_messages = 0
    assistant_message_found = False
    
    # Iterate over all messages to find assistant responses
    # The actual model response messages are part of the `messages` list.
    # Typically, these would be the last few messages if it's a multi-turn interaction,
    # or messages[-1] if it's a single assistant response.
    # For simplicity in format check, we scan all assistant messages in the provided `messages`.
    for msg in messages: 
        if isinstance(msg, Message) and msg.role == 'assistant':
            assistant_message_found = True
            total_assistant_messages +=1 
            # Check for any content or any tool_call
            if (msg.content and msg.content.strip()) or msg.tool_calls:
                valid_assistant_messages +=1
    
    if not assistant_message_found: 
        format_score = 0.0
    elif total_assistant_messages > 0 and valid_assistant_messages == 0:
        # Assistant messages were found, but none had content or tool_calls
        format_score = 0.0 
    # If valid_assistant_messages > 0, format_score remains 0.2 (or could be scaled)

    # --- Combine Scores ---
    base_score = state_match_score + func_match_score

    if base_score >= 1.0:
        final_score = base_score + format_score
        reason = "State and function calls matched ground truth."
        if format_score == 0.2:
            reason += " Format was also correct."
        else:
            reason += " Format was incorrect."
    else:
        final_score = 0.0
        reason = "State or function calls did not perfectly match ground truth."
        if state_match_score < 0.5:
            reason += f" State match failed."
            if state_diffs:
                 reason += f" Differences: {json.dumps(state_diffs)}"
        if func_match_score < 0.5: # Check against 0.5 as perfect score for this component
             reason += f" Function call match failed ({num_func_matches_for_score}/{reason_num_total_gt_turns_with_calls} GT turns with calls matched)."


    # Add metrics
    metrics = {} # Initialize metrics dictionary
    metrics["state_match"] = MetricResult(
        score=state_match_score,
        success=state_match_score == 0.5,
        reason=f"State match: {state_match}" + (f", Differences: {json.dumps(state_diffs)}" if state_diffs else "")
    )
    metrics["function_call_match"] = MetricResult(
        score=func_match_score,
        success=func_match_score == 0.5, # Success if it gets the full 0.5 for this part
        reason=f"{num_func_matches_for_score}/{reason_num_total_gt_turns_with_calls} GT turns with calls matched by model. Model made calls in {num_model_turns_with_actual_calls} turn(s)."
    )
    metrics["format_check"] = MetricResult(
        score=format_score,
        success=format_score == 0.2, # Success if it gets the full 0.2 for format
        reason=f"{valid_assistant_messages}/{total_assistant_messages} assistant messages had correct format."
    )

    return EvaluateResult(
        score=final_score,
        reason=reason,
        metrics=metrics
    )
