"""多学生作业处理循环节点：完全并行处理，确保数据隔离"""
import logging
import threading
from typing import List, Dict, Any, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed, TimeoutError as FuturesTimeoutError
from langchain_core.runnables import RunnableConfig
from langgraph.runtime import Runtime
from coze_coding_utils.runtime_ctx.context import Context

from graphs.state import (
    ProcessImagesInput,
    ProcessImagesOutput,
    StudentResult,
    StudentHomework,
    SingleImageResult,
    SubgraphInput,
    ImageInfo,
    CorrectAnswer
)
from graphs.loop_graph import single_image_subgraph
from utils.file.file import File

logger = logging.getLogger(__name__)


# 默认等级标准（与GraphInput中的默认值保持一致）
DEFAULT_GRADE_STANDARDS = {
    "A+": {"min_percentage": 95, "description": "优秀"},
    "A": {"min_percentage": 90, "description": "良好"},
    "B": {"min_percentage": 80, "description": "合格"},
    "C": {"min_percentage": 70, "description": "及格"},
    "D": {"min_percentage": 0, "description": "需努力"}
}

# 超时配置（秒）
SINGLE_IMAGE_TIMEOUT = 120  # 单张图片处理超时120秒


def calculate_grade(score_rate: float, incorrect_count: int, total_questions: int, grade_standards: Dict[str, Any]) -> tuple:
    """
    根据得分率、错误数量和等级标准计算等级
    
    核心规则（按优先级）：
    1. A+ 和 A：首要条件是"全对"（incorrect_count == 0），与得分率无关
       - A+：全对
       - A：全对
       （A+和A的区别由其他因素决定，这里都返回全对的最高等级）
    2. B/C/D：有错误时，按得分率判断
    
    Args:
        score_rate: 得分率（百分比，如95.5）
        incorrect_count: 错误题目数量
        total_questions: 总题目数量
        grade_standards: 等级标准字典
        
    Returns:
        (等级, 等级描述)
    """
    # 使用传入的等级标准，如果没有则使用默认值
    standards = grade_standards if grade_standards else DEFAULT_GRADE_STANDARDS
    
    # 规则1：全对 → A+ 或 A（这里统一返回A+，因为全对就是最高等级）
    if incorrect_count == 0 and total_questions > 0:
        return "A+", standards.get("A+", {}).get("description", "优秀")
    
    # 规则2：有错误 → 按得分率判断 B/C/D
    # 按min_percentage降序排序（只排B/C/D）
    bcd_grades = [(k, v) for k, v in standards.items() if k not in ("A+", "A")]
    sorted_grades = sorted(
        bcd_grades,
        key=lambda x: x[1].get("min_percentage", 0),
        reverse=True
    )
    
    # 遍历找到匹配的等级
    for grade, config in sorted_grades:
        min_pct = config.get("min_percentage", 0)
        if score_rate >= min_pct:
            return grade, config.get("description", "")
    
    # 默认返回D
    return "D", "需努力"


def process_single_image(
    student_id: int,
    idx: int,
    image_url: str,
    correct_answers: List[CorrectAnswer],
    comment_max_length: int,
    config: RunnableConfig
) -> tuple:
    """
    处理单个学生的单张图片（线程安全）
    返回: (student_id, image_index, SingleImageResult) 元组，确保学生ID关联
    
    注意：correct_answers 参数是只读的，不会被修改
    """
    logger.info(f"Processing student {student_id}, image {idx + 1}")
    try:
        # 将URL字符串转换为File对象
        homework_image = File(url=image_url, file_type="image")
        
        # 构建子图输入（创建新的输入对象，确保数据隔离）
        # 注意：correct_answers 是只读的，不需要复制
        subgraph_input = SubgraphInput(
            homework_image=homework_image,
            correct_answers=correct_answers,  # 只读，不需要复制
            image_index=idx,
            comment_max_length=comment_max_length
        )
        
        # 调用子图（config 是只读的，LangGraph 保证线程安全）
        subgraph_output = single_image_subgraph.invoke(subgraph_input, config)
        
        # 提取结果（兼容字典和Pydantic对象两种返回类型）
        image_result = None
        if subgraph_output:
            if isinstance(subgraph_output, dict):
                result_data = subgraph_output.get("image_result")
                if result_data:
                    if isinstance(result_data, dict):
                        image_result = SingleImageResult(**result_data)
                    else:
                        image_result = result_data
            elif hasattr(subgraph_output, 'image_result'):
                image_result = subgraph_output.image_result
        
        if image_result:
            logger.info(f"Student {student_id}, image {idx + 1} processed: {len(image_result.annotations)} annotations")
            return (student_id, idx, image_result)
        else:
            logger.warning(f"Student {student_id}, image {idx + 1} returned invalid output")
            return (student_id, idx, SingleImageResult(
                image_index=idx,
                image_info=ImageInfo(width=0, height=0, dpi=72),
                annotations=[]
            ))
    except Exception as e:
        logger.error(f"Failed to process student {student_id}, image {idx + 1}: {e}", exc_info=True)
        return (student_id, idx, SingleImageResult(
            image_index=idx,
            image_info=ImageInfo(width=0, height=0, dpi=72),
            annotations=[]
        ))


def process_images_node(
    state: ProcessImagesInput,
    config: RunnableConfig,
    runtime: Runtime[Context]
) -> ProcessImagesOutput:
    """
    title: 多学生作业批改处理
    desc: 完全并行处理所有学生的所有图片，确保每个学生的数据完全隔离
    integrations: 
    """
    ctx = runtime.context
    
    # === 输入参数校验 ===
    if not state.student_homework:
        logger.warning("No student homework provided, returning empty result")
        return ProcessImagesOutput(student_results=[])
    
    # 获取并发数限制（从参数获取，默认10，限制在合理范围1-50）
    max_concurrent = max(1, min(getattr(state, 'max_concurrent', 10), 50))
    
    # 过滤有效的图片URL，统计总任务数
    valid_tasks = []  # [(student_id, idx, image_url), ...]
    for student in state.student_homework:
        if not student.homework_images:
            continue
        for idx, image_url in enumerate(student.homework_images):
            # 验证图片URL有效性
            if image_url and isinstance(image_url, str) and image_url.strip():
                valid_tasks.append((student.student_id, idx, image_url.strip()))
            else:
                logger.warning(f"Invalid image URL for student {student.student_id}, image {idx}")
    
    if not valid_tasks:
        logger.warning("No valid images to process, returning empty result")
        return ProcessImagesOutput(student_results=[])
    
    logger.info(f"Starting to process {len(state.student_homework)} students, {len(valid_tasks)} valid images (concurrent={max_concurrent})")
    
    # 第一步：提交所有任务（所有学生的所有图片）
    # 注意：使用字典存储结果，as_completed 在主线程顺序处理，无需加锁
    student_image_results: Dict[int, Dict[int, SingleImageResult]] = {}
    
    with ThreadPoolExecutor(max_workers=max_concurrent) as executor:
        # 提交所有有效任务
        future_to_task = {}  # future -> (student_id, idx, image_url)
        for student_id, idx, image_url in valid_tasks:
            future = executor.submit(
                process_single_image,
                student_id,
                idx,
                image_url,
                state.correct_answers,
                state.comment_max_length,
                config
            )
            future_to_task[future] = (student_id, idx, image_url)
        
        logger.info(f"Submitted {len(future_to_task)} tasks for all students")
        
        # 第二步：收集所有结果，按student_id分组
        # as_completed 在主线程中迭代，不需要加锁
        # 双重保险：使用 task_info 中的 student_id（而非 result 中的），确保数据归属正确
        try:
            for future in as_completed(future_to_task, timeout=SINGLE_IMAGE_TIMEOUT * len(future_to_task)):
                task_info = future_to_task[future]
                student_id, idx, image_url = task_info  # 从任务信息获取，不依赖返回值
                
                try:
                    # 返回值中的 student_id 应该与 task_info 一致，但我们优先使用 task_info
                    # 添加单任务超时保护
                    result_student_id, image_index, image_result = future.result(timeout=SINGLE_IMAGE_TIMEOUT)
                    
                    # 双重校验：确保返回的学生ID与任务一致
                    if result_student_id != student_id:
                        logger.warning(f"Student ID mismatch: task={student_id}, result={result_student_id}, using task_id")
                    
                    # 使用 task_info 中的 student_id（更可靠）
                    if student_id not in student_image_results:
                        student_image_results[student_id] = {}
                    
                    # 存储该学生的图片结果
                    student_image_results[student_id][image_index] = image_result
                    
                except FuturesTimeoutError:
                    logger.error(f"Task timeout for student {student_id}, image {idx} (timeout={SINGLE_IMAGE_TIMEOUT}s)")
                    # 存储一个空结果，确保该学生的结果完整
                    if student_id not in student_image_results:
                        student_image_results[student_id] = {}
                    student_image_results[student_id][idx] = SingleImageResult(
                        image_index=idx,
                        image_info=ImageInfo(width=0, height=0, dpi=72),
                        annotations=[]
                    )
                except Exception as e:
                    logger.error(f"Task failed for student {student_id}, image {idx}: {e}", exc_info=True)
                    # 存储一个空结果，确保该学生的结果完整
                    if student_id not in student_image_results:
                        student_image_results[student_id] = {}
                    student_image_results[student_id][idx] = SingleImageResult(
                        image_index=idx,
                        image_info=ImageInfo(width=0, height=0, dpi=72),
                        annotations=[]
                    )
        
        except FuturesTimeoutError:
            # as_completed 总超时，处理未完成的任务
            logger.error(f"Total timeout exceeded, processing remaining {len(future_to_task)} tasks")
            for future, task_info in future_to_task.items():
                if not future.done():
                    student_id, idx, _ = task_info
                    if student_id not in student_image_results:
                        student_image_results[student_id] = {}
                    student_image_results[student_id][idx] = SingleImageResult(
                        image_index=idx,
                        image_info=ImageInfo(width=0, height=0, dpi=72),
                        annotations=[]
                    )
        
        logger.info(f"Collected results for {len(student_image_results)} students")
    
    # 第三步：为每个学生计算独立的结果（数据隔离）
    student_results: List[StudentResult] = []
    
    for student in state.student_homework:
        student_id = student.student_id
        student_name = student.student_name
        homework_image_urls = student.homework_images
        
        # 获取该学生的图片结果（确保只使用该学生的数据）
        image_results_dict = student_image_results.get(student_id, {})
        
        # 按顺序组装该学生的图片结果
        image_results: List[SingleImageResult] = []
        for idx in range(len(homework_image_urls)):
            if idx in image_results_dict:
                image_results.append(image_results_dict[idx])
            else:
                # 补充缺失的结果
                image_results.append(SingleImageResult(
                    image_index=idx,
                    image_info=ImageInfo(width=0, height=0, dpi=72),
                    annotations=[]
                ))
        
        # 计算该学生的总分（只基于该学生的图片结果）
        total_score = 0
        full_score = 0
        total_questions = 0
        correct_count = 0
        incorrect_count = 0
        
        for result in image_results:
            for annotation in result.annotations:
                total_score += annotation.score
                full_score += annotation.full_score
                total_questions += 1
                if annotation.status == "correct":
                    correct_count += 1
                elif annotation.status == "incorrect":
                    incorrect_count += 1
        
        # 计算得分率
        score_rate = (total_score / full_score * 100) if full_score > 0 else 0
        
        # 获取等级标准（从state中获取，如果没有则使用默认值）
        grade_standards = getattr(state, 'grade_standards', {})
        
        # 使用等级标准计算等级（传入错误数量和总题数）
        grade, grade_description = calculate_grade(score_rate, incorrect_count, total_questions, grade_standards)
        
        # 生成该学生的整体评价（与等级匹配）
        if total_questions == 0:
            overall_comment = "未识别到题目内容"
            grade = "D"
        elif grade == "A+":
            # A+：全部正确
            overall_comment = f"优秀！{total_questions}题全部正确，掌握扎实，继续保持。"
        elif grade == "A":
            # A：全部正确（已确保无错误）
            overall_comment = f"良好！{total_questions}题全部正确，步骤规范。"
        elif grade == "B":
            # B：有少量错误
            if incorrect_count > 0:
                overall_comment = f"合格。得分率{score_rate:.0f}%，错{incorrect_count}题，需加强练习。"
            else:
                overall_comment = f"合格。得分率{score_rate:.0f}%，继续努力。"
        elif grade == "C":
            overall_comment = f"及格。错{incorrect_count}题，部分知识点掌握不牢，建议复习。"
        else:
            overall_comment = f"需努力。得分率{score_rate:.0f}%，错{incorrect_count}题，建议认真复习，多做练习。"
        
        # 创建该学生的完整结果（数据完全隔离）
        student_result = StudentResult(
            student_id=student_id,
            student_name=student_name,
            total_images=len(image_results),
            image_results=image_results,
            overall_comment=overall_comment,
            total_score=total_score,
            full_score=full_score,
            grade=grade
        )
        
        student_results.append(student_result)
        
        logger.info(f"Student {student_id}: {total_score}/{full_score}, grade={grade}")
    
    logger.info(f"Completed processing {len(student_results)} students")
    
    return ProcessImagesOutput(student_results=student_results)
