from typing import Optional
import json


def format_output_entry(
    original_data: dict,
    a1_result: dict,
    a3_result: dict,
    a4_result: Optional[dict],
    a5_result: dict,
    ablation_type: str,
    token_usage: Optional[dict] = None
) -> dict:
    """
    构造一个标准的样本结果条目，用于输出到统一 JSON 文件。
    """

    # ✅ Step 1: 构造基本字段（先插入a1和a3）
    output = {
        "id": original_data.get("id", ""),
        "question": original_data.get("question", ""),
        "context": original_data.get("context", ""),
        "a1": a1_result,
        "a3": a3_result,
    }

    # ✅ Step 2: 插入 a4（应在 a3 后、a5 前）
    if ablation_type != "woFOL" and a4_result is not None:
        output["a4"] = a4_result

    # ✅ Step 3: 插入 a5 和预测结果
    output["a5"] = a5_result
    output["predicted_answer"] = a5_result.get("s1_truth", "Uncertain")
    output["answer"] = original_data.get("answer", "")

    # ✅ Step 4: 插入 token 使用记录
    if token_usage:
        output["token_usage"] = {
            "prompt_tokens": token_usage.get("prompt_tokens", 0),
            "completion_tokens": token_usage.get("completion_tokens", 0),
            "total_tokens": token_usage.get("total_tokens", 0)
        }

    return output
