import json
import os
import re
from glob import glob
from typing import List, Tuple

def get_agent_idx_from_action(action: str) -> int:
    """Extract agent index from action string (e.g., 'goto_agent1_storage0' -> 1)."""
    match = re.search(r'agent(\d+)', action)
    if match:
        return int(match.group(1))
    return -1

def calculate_success_rate(action_history: List[List[str]], success_history: List[List[bool]], num_agents: int) -> Tuple[List[float], List[int], List[int]]:
    """Calculate success rate for each agent by parsing agent indices from action strings,
    counting all actions (including noops)."""
    # Initialize counters for each agent
    total_actions = [0] * num_agents
    successful_actions = [0] * num_agents
    
    # Iterate over each step in the flattened histories
    for actions, successes in zip(action_history, success_history):
        if not actions or not successes:  # Skip empty steps
            continue
            
        # Process each action in the step
        for action, success in zip(actions, successes):
            agent_idx = get_agent_idx_from_action(action)
            if 0 <= agent_idx < num_agents:
                total_actions[agent_idx] += 1
                if success:
                    successful_actions[agent_idx] += 1
    
    # Compute success rates (rounded to three decimals)
    success_rates = []
    for total, successful in zip(total_actions, successful_actions):
        rate = successful / total if total > 0 else 0
        success_rates.append(round(rate, 3))
    
    return success_rates, total_actions, successful_actions

def process_result_file(file_path: str, num_agents: int) -> Tuple[str, List[float]]:
    """Process a single result file and return level name and success rates only."""
    with open(file_path, 'r') as f:
        data = json.load(f)
    
    # Get level name from filename (assuming format: result_level_X_*.json)
    level_name = os.path.basename(file_path).split('_')[2]
    
    # Retrieve the alpha level data (assumes key "1.0" exists and holds a list)
    alpha_data = data.get("1.0", [{}])[0]
    
    # Get the raw (possibly nested) action and success histories
    raw_action_history = alpha_data.get("action_history", [])
    raw_success_history = alpha_data.get("action_success_history", [])
    
    # Flatten the nested lists
    action_history = [step for episode in raw_action_history for step in episode]
    success_history = [step for episode in raw_success_history for step in episode]
    
    success_rates, _, _ = calculate_success_rate(action_history, success_history, num_agents)
    return level_name, success_rates

def main():
    # Process for different numbers of agents
    for num_agents in [1, 2, 3]:
        result_files = glob(f"results/gpt-4o-v2/planner/{num_agents}/result_level_*.json")
        success_rates_by_level = {}
        
        for file_path in result_files:
            level_name, rates = process_result_file(file_path, num_agents)
            # Only include the success rate per agent
            level_data = {}
            for i in range(num_agents):
                level_data[f"agent{i}_success_rate"] = rates[i] if i < len(rates) else 0
            success_rates_by_level[level_name] = level_data
        
        output_file = f"success_rates_{num_agents}agent.json"
        with open(output_file, 'w') as f:
            json.dump(success_rates_by_level, f, indent=2)
        
        print(f"Created {output_file} with success rates for {num_agents} agent(s)")

if __name__ == "__main__":
    main()
