import os
import json
import time
from pydantic import BaseModel, Field
from tqdm import tqdm
import google.generativeai as genai
import google.generativeai.types as genai_types

# API Keys 列表
API_KEYS = [
    "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
    "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
    "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
]

def get_configured_client(api_key):
    try:
        genai.configure(api_key=api_key)
        return genai.GenerativeModel(
            model_name="gemini-2.5-flash",
            # model_name="gemini-2.5-pro",
        )
    except Exception as e:
        print(f"Error configuring Gemini client with key {api_key[:10]}...: {e}")
        return None

# --- 2. 定义评估维度和输出格式 ---

class EvaluationResult(BaseModel):
    """
    定义了LLM-as-a-judge的输出结构，用于评估模型CoT的质量。
    """
    reasoning_correctness_score: int = Field(
        ..., ge=1, le=5, description="分数 1-5: 推理过程是否正确反映了视频内容，逻辑是否无误。"
    )
    reasoning_consistency_score: int = Field(
        ..., ge=1, le=5, description="分数 1-5: 推理过程是否与最终答案保持一致，无内部矛盾。"
    )
    reasoning_redundancy_score: int = Field(
        ..., ge=1, le=5, description="分数 1-5: 推理过程是否简洁，有无不必要的重复。5分代表最简洁。"
    )
    final_answer_correctness: bool = Field(
        ..., description="最终答案是否与'Ground Truth Answer'匹配 (True/False)。"
    )
    evaluation_rationale: str = Field(
        ..., description="对所有评分的详细解释，指出被评估模型回答的优点和缺点。"
    )

# --- 3.1 辅助函数 ---
def process_video_path(video_path):
    """
    处理视频路径，确保路径是绝对路径或者相对于当前工作目录的正确路径。
    这里可以根据实际情况添加更多处理逻辑，例如替换路径前缀等。
    """
    # 示例：如果视频路径是相对路径，可以尝试拼接一个基础路径
    # base_path = "/data/videos"
    # if not os.path.isabs(video_path):
    #     return os.path.join(base_path, video_path)
    return video_path

# --- 3. 设计评估Prompt模板 ---

EVALUATION_PROMPT_TEMPLATE = """
You are an impartial and meticulous AI evaluator. Your task is to critically assess the quality of a model's response to a video-based question. The model's response includes its step-by-step thinking process (Chain-of-Thought) and a final answer.

**You are provided with:**
1.  **The Video:** The primary source of information.
2.  **The Context and Question:** The user's query and any associated text or options.
3.  **The Ground Truth Answer:** The officially correct answer.
4.  **The Model's Full Response to be Evaluated:** This contains the model's entire reasoning process and its final conclusion.

**Your Evaluation Criteria:**

*   **1. Reasoning Correctness (Score 1-5):** Evaluate if the reasoning steps accurately interpret the video. Are the facts derived from the video correct? Is the logic sound?
    - 5: Perfectly accurate and logical.
    - 3: Contains minor errors or inaccuracies but the overall direction is correct.
    - 1: Completely hallucinatory or illogical.

*   **2. Reasoning Consistency (Score 1-5):** Evaluate if the final answer is a logical conclusion of the reasoning steps.
    - 5: The answer is perfectly supported by the reasoning.
    - 3: The answer is plausible but not strongly supported by the reasoning.
    - 1: The answer contradicts the reasoning.

*   **3. Reasoning Redundancy (Score 1-5):** Evaluate the conciseness of the reasoning.
    - 5: Highly concise and to the point. No redundant steps.
    - 3: Contains some repetition or slightly verbose steps.
    - 1: Extremely repetitive and convoluted.

*   **4. Final Answer Correctness (True/False):** Strictly compare the model's final answer with the ground truth answer.

**The Task to be Evaluated:**
---
**Context and Question:**
{question_text}

**Ground Truth Answer:**
{ground_truth_answer}

**Model's Full Response (CoT and Final Answer):**
{model_answer_cot}
---

Your output **MUST** be a single, valid JSON object that adheres to the following structure. Do not include any text, explanations, or markdown formatting outside of the JSON object.

Example JSON structure:
reasoning_correctness_score
"""

# --- 4. 主逻辑 ---

def main():
    # 输入：包含模型输出的JSON文件路径
    input_model_outputs_file = "xxxxxxxxxx"
    # 输出：保存评估结果的目录
    save_dir = "xxxxxxxxxxxxxxxxxxxxx"
    os.makedirs(save_dir, exist_ok=True)

    try:
        with open(input_model_outputs_file, "r") as f:
            dataset = json.load(f)
        if isinstance(dataset, dict): # 如果根是字典，则转换为项列表
             dataset = list(dataset.values())
    except (FileNotFoundError, json.JSONDecodeError) as e:
        print(f"Error loading dataset: {e}")
        return

    last_video_path = ""
    last_upload_api_key = None
    video_file_object = None

    for idx, item in enumerate(tqdm(dataset, desc="Evaluating Model Outputs")):
        question_id = item.get("question_id", f"item_{idx}")

        # 修复question_id包含非法字符（如/）的问题，因为它是文件名的一部分
        safe_question_id = question_id.replace("/", "_").replace("\\", "_")

        video_path = item.get("video_path")

        save_path = os.path.join(save_dir, f"{safe_question_id}_evaluation.json")
        # 检查是否已经评估过，如果文件存在且非空，则跳过
        if os.path.exists(save_path):
            try:
                with open(save_path, "r") as f:
                    # 尝试读取并解析JSON，如果成功且包含必要字段，则视为有效
                    existing_data = json.load(f)
                    if "llm_judge_evaluation" in existing_data:
                        print(f"Skipping {question_id}, already evaluated.")
                        continue
            except (json.JSONDecodeError, ValueError):
                print(f"Found corrupted or empty file for {question_id}, re-evaluating...")

        # 检查是否在之前的失败日志中（可选，但主要依靠文件是否存在来判断）
        # 即使在failed_items.log中，只要save_path不存在或无效，都会重新评估

        if not video_path:
            print(f"Skipping {question_id}, no video path provided.")
            continue

        # 处理视频路径
        processed_video_path = process_video_path(video_path)

        if not os.path.exists(processed_video_path):
            print(f"Video not found for {question_id}: {processed_video_path}")
            continue

        # 准备评估Prompt
        full_prompt = EVALUATION_PROMPT_TEMPLATE.format(
            question_text=item["question"],
            ground_truth_answer=item["answer"],
            model_answer_cot=item["model_answer"]
        )

        # 调用API进行评估 (整合上传逻辑以解决跨API文件访问问题)
        for attempt_idx, api_key in enumerate(API_KEYS):
            client = get_configured_client(api_key)
            if not client:
                continue

            # 检查是否需要上传/重新上传视频
            # 如果视频路径变了，或者当前API Key与上次上传的Key不一致，则必须重新上传
            if processed_video_path != last_video_path or api_key != last_upload_api_key or video_file_object is None:
                try:
                    print(f"\nUploading video: {processed_video_path} with API Key {attempt_idx + 1}")
                    # genai.configure 已在 get_configured_client 中调用，设置了全局 active client

                    video_file_object = genai.upload_file(path=processed_video_path)

                    # 等待视频处理完成
                    print(f"Waiting for video processing: {video_file_object.name}")
                    while video_file_object.state.name == "PROCESSING":
                        time.sleep(5)
                        video_file_object = genai.get_file(video_file_object.name)
                        print(f"Current state: {video_file_object.state.name}")

                    if video_file_object.state.name == "FAILED":
                        print(f"Video processing failed: {video_file_object.name}")
                        continue # 当前Key上传失败，尝试下一个Key

                    # 上传成功，更新状态
                    last_video_path = processed_video_path
                    last_upload_api_key = api_key
                except Exception as e:
                    print(f"Error uploading video with Key {attempt_idx + 1}: {e}")
                    continue # 当前Key上传出错，尝试下一个Key
            else:
                print(f"\nUsing previously uploaded video for {question_id} (Key {attempt_idx + 1}).")

            try:
                print(f"Generating evaluation for {question_id} with API Key {attempt_idx + 1}/{len(API_KEYS)}...")

                # 定义安全设置，禁用所有过滤以避免 PROHIBITED_CONTENT
                safety_settings = [
                    {
                        "category": "HARM_CATEGORY_HARASSMENT",
                        "threshold": "BLOCK_NONE"
                    },
                    {
                        "category": "HARM_CATEGORY_HATE_SPEECH",
                        "threshold": "BLOCK_NONE"
                    },
                    {
                        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                        "threshold": "BLOCK_NONE"
                    },
                    {
                        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                        "threshold": "BLOCK_NONE"
                    },
                ]

                response = client.generate_content(
                    contents=[full_prompt, video_file_object], # Prompt和视频文件
                    generation_config={
                        'response_mime_type': 'application/json',
                    },
                    safety_settings=safety_settings,
                    request_options={'timeout': 120} # 设置120秒超时
                )

                # 解析返回的JSON结果
                evaluation_data = json.loads(response.text)

                # 验证结果是否符合Pydantic模型 (可选但推荐)
                EvaluationResult(**evaluation_data) 

                # 准备最终保存文件
                final_output = {
                    "source_data": item,
                    "llm_judge_evaluation": evaluation_data
                }

                # 保存评估结果
                with open(save_path, "w", encoding="utf-8") as f:
                    json.dump(final_output, f, indent=2, ensure_ascii=False)

                print(f"Successfully saved evaluation for {question_id}")
                break # 成功后跳出重试循环

            except Exception as e:
                print(f"Error with API Key {attempt_idx + 1}: {e}")
                if attempt_idx == len(API_KEYS) - 1:
                    print(f"All API keys failed for {question_id}.")
                    # 记录失败的项，以便后续重试
                    with open(os.path.join(save_dir, "failed_items.log"), "a") as log_file:
                        log_file.write(f"{question_id}: {e}\n")
                else:
                    print("Switching to next API key...")
                    time.sleep(10) # 稍作等待再尝试下一个Key

        # API速率控制
        time.sleep(10) # 每次成功请求后稍作休息

if __name__ == "__main__":
    main()

            if processed_video_path != last_video_path or api_key != last_upload_api_key or video_file_object is None:
                try:
                    print(f"\nUploading video: {processed_video_path} with API Key {attempt_idx + 1}")
                    # genai.configure 已在 get_configured_client 中调用，设置了全局 active client

                    video_file_object = genai.upload_file(path=processed_video_path)

                    # 等待视频处理完成
                    print(f"Waiting for video processing: {video_file_object.name}")
                    while video_file_object.state.name == "PROCESSING":
                        time.sleep(5)
                        video_file_object = genai.get_file(video_file_object.name)
                        print(f"Current state: {video_file_object.state.name}")

                    if video_file_object.state.name == "FAILED":
                        print(f"Video processing failed: {video_file_object.name}")
                        continue # 当前Key上传失败，尝试下一个Key

                    # 上传成功，更新状态
                    last_video_path = processed_video_path
                    last_upload_api_key = api_key
                except Exception as e:
                    print(f"Error uploading video with Key {attempt_idx + 1}: {e}")
                    continue # 当前Key上传出错，尝试下一个Key
            else:
                print(f"\nUsing previously uploaded video for {question_id} (Key {attempt_idx + 1}).")