"""4. 结果整合节点：将识别结果和批改结果合并为最终批注"""
from typing import List
from langchain_core.runnables import RunnableConfig
from langgraph.runtime import Runtime
from coze_coding_utils.runtime_ctx.context import Context

from graphs.state import (
    ResultMergeInput,
    ResultMergeOutput,
    Annotation,
    QuestionItem,
    CorrectionResult
)


def result_merge_node(
    state: ResultMergeInput,
    config: RunnableConfig,
    runtime: Runtime[Context]
) -> ResultMergeOutput:
    """
    title: 结果整合
    desc: 将识别结果和批改结果合并为最终批注列表
    integrations: 
    """
    ctx = runtime.context
    
    annotations: List[Annotation] = []
    
    # 创建批改结果字典，方便查找
    correction_dict = {c.question_id: c for c in state.correction_results}
    
    for q in state.question_items:
        # 查找对应的批改结果
        correction = correction_dict.get(q.question_id)
        
        if correction is None:
            # 如果没有批改结果，默认为错误
            annotations.append(Annotation(
                question_id=q.question_id,
                parent_id=q.parent_id,
                status="incorrect",
                question_text=q.question_text,
                student_answer=q.student_answer,
                answer_bbox=q.answer_bbox,
                comment="未找到批改结果",
                mark_position=q.mark_position,
                score=0,
                full_score=q.full_score
            ))
        else:
            annotations.append(Annotation(
                question_id=q.question_id,
                parent_id=q.parent_id,
                status=correction.status,
                question_text=q.question_text,
                student_answer=q.student_answer,
                answer_bbox=q.answer_bbox,
                comment=correction.comment,
                mark_position=q.mark_position,
                score=correction.score,
                full_score=correction.full_score
            ))
    
    return ResultMergeOutput(annotations=annotations)
