from typing import Dict, Any, Tuple


def process_action_agent(event: Dict[str, Any]) -> Tuple[str, bool]:
    if event["ActionAgent"].get("additional_kwargs"):
        return str(event["ActionAgent"].get("additional_kwargs")), False

    pred = event["ActionAgent"].get("messages", [{}])[-1]
    content = pred.content if pred and pred.content else "No prediction available"

    has_tool_call = bool(pred.tool_calls) if pred else False

    if has_tool_call:
        tool_calls = clean_tool_calls(pred.tool_calls, event)
        content = f"Tool calls:\n{tool_calls}"

    return content, has_tool_call


def process_ranker_agent(event: Dict[str, Any]) -> str:
    return str(event["RankerAgent"].get("elements", {}))


def process_call_tool(event: Dict[str, Any]) -> str:
    messages = event["call_tool"].get("messages", [])
    return "\n".join(msg.content for msg in messages) if messages else "No tool call information available"


def process_memory_agent(event: Dict[str, Any]) -> str:
    memory = event["MemoryAgent"].get("observation", "")[0]
    return f"MemoryAgent updated memory:\n\n{memory.content}"


def process_planner_agent(event: Dict[str, Any]) -> str:
    return str(event["PlannerAgent"].get("policy", "No policy available"))


def process_event(event: Dict[str, Any]) -> Tuple[str, bool]:
    event_processors = {
        "ActionAgent": process_action_agent,
        "RankerAgent": process_ranker_agent,
        "call_tool": process_call_tool,
        "MemoryAgent": process_memory_agent,
        "PlannerAgent": process_planner_agent,
    }

    for agent, processor in event_processors.items():
        if agent in event:
            if agent == "ActionAgent":
                prediction, has_tool_call = processor(event)
                return f"\n {agent} prediction:\n\n{prediction}", has_tool_call
            else:
                return f"\n {agent} prediction:\n\n{processor(event)}", False

    return "Unrecognized event type", False


def clean_tool_calls(tool_calls, event=None):
    if not isinstance(tool_calls, list):
        return tool_calls
    # Its every element beside the 'cur_state' key
    for call in tool_calls:
        if 'args' in call and 'state' in call['args']:
            if 'cur_state' in call['args']['state']:
                cur_state = call['args']['state'].pop('cur_state', None)
            try:
                call['args']['state']['env'] = event['ActionAgent']['env']
            except:
                print(f"Error: \n\n tool_calls = \n {tool_calls} \n\n call['args']['state'] = \n {call['args']['state']}")
            # call['args']['state']['env'] = env

    return tool_calls
