import json

def calculate_total_latency(json_file_path, how_many=1000, start_idx=500):
    """
    读取指定的 JSON 文件，计算所有 prompt 条目中 "latency" 值的总和。

    Args:
        json_file_path (str): 包含 prompt 数据的 JSON 文件的路径。

    Returns:
        float: 所有 "latency" 值的累加和。如果文件不存在或格式错误，
               或者没有找到任何 latency 值，则返回 0.0。
    """
    total_latency = 0.0
    total_step_compression = 0.0
    total_target_calling = 0
    cnt = 0

    try:
        # 使用 'with' 语句确保文件正确打开和关闭
        with open(json_file_path, 'r', encoding='utf-8') as f:
            # 使用 json.load() 从文件对象解析 JSON 数据为 Python 字典
            data = json.load(f)

        # 检查 data 是否是一个字典
        if not isinstance(data, dict):
            print(f"错误：JSON 文件的顶层结构不是一个字典。")
            return 0.0

        if start_idx + how_many > len(data):
            print(f"错误：请求的范围超出数据长度。")
            start_idx = len(data) - how_many

        # 遍历字典中所有的值 (每个值对应一个 prompt_i 的字典)
        for idx, (prompt_key, prompt_data) in enumerate(data.items()):
            if idx + 1 < start_idx:
                continue
            else:
                # 检查 prompt_data 是否是字典并且包含 'latency'键
                if isinstance(prompt_data, dict) and 'latency' in prompt_data:
                    latency_value = prompt_data['latency']
                    # 确保 latency 值是数字类型 (int 或 float)
                    if isinstance(latency_value, (int, float)):
                        total_latency += latency_value
                    else:
                        print(f"警告：在键 '{prompt_key}' 中找到非数字类型的 latency 值：{latency_value}")
                else:
                    print(f"警告：键 '{prompt_key}' 的值不是字典或缺少 'latency' 键。")

                if isinstance(prompt_data, dict) and 'step_compression' in prompt_data:
                    step_compression = prompt_data['step_compression']
                    # 确保 step_compression 值是数字类型 (int 或 float)
                    if isinstance(step_compression, (int, float)):
                        total_step_compression += step_compression
                    else:
                        print(f"警告：在键 '{prompt_key}' 中找到非数字类型的 latency 值：{step_compression}")
                else:
                    print(f"警告：键 '{prompt_key}' 的值不是字典或缺少 'latency' 键。")

                if isinstance(prompt_data, dict) and 'num_target_calling' in prompt_data:
                    target_calling = prompt_data['num_target_calling']
                    # 确保 step_compression 值是数字类型 (int 或 float)
                    if isinstance(target_calling, (int, float)):
                        total_target_calling += target_calling
                    else:
                        print(f"警告：在键 '{prompt_key}' 中找到非数字类型的 latency 值：{target_calling}")
                else:
                    print(f"警告：键 '{prompt_key}' 的值不是字典或缺少 'latency' 键。")

                cnt += 1

            if cnt >= how_many:
                break

    except FileNotFoundError:
        print(f"错误：找不到文件 '{json_file_path}'")
        return 0.0 # 或者可以抛出异常 raise
    except json.JSONDecodeError:
        print(f"错误：无法解码文件 '{json_file_path}' 中的 JSON 数据。请检查文件格式。")
        return 0.0 # 或者可以抛出异常 raise
    except Exception as e:
        print(f"发生意外错误：{e}")
        return 0.0 # 或者可以抛出异常 raise

    return total_latency, total_step_compression, total_target_calling

if __name__ == '__main__':

    # --- 使用示例 ---
    # 1. 将你的 JSON 文件路径替换下面的 'your_data.json'
    #    例如: json_file_path = 'data/results.json' 或 json_file_path = 'C:/Users/YourUser/Documents/data.json'
    json_file_path = ''

    # 2. 调用函数计算总延迟
    total_latency, total_step_compression, total_target_calling = calculate_total_latency(json_file_path, how_many=50, start_idx=0)


    # 3. 打印结果
    if total_latency is not None and total_step_compression is not None:
        print(f"\n所有 prompt 的 'latency' 总和为: {total_latency}")
        print(f"\n所有 prompt 的 'step_compression' 总和为: {total_step_compression}")
        print(f"\n所有 prompt 的 'num_target_calling' 总和为: {total_target_calling}")