"""单图片处理子图：封装完整的单张图片批改流程（优化版：合并识别和批改）"""
from langgraph.graph import StateGraph, END
from langchain_core.runnables import RunnableConfig
from langgraph.runtime import Runtime
from coze_coding_utils.runtime_ctx.context import Context

from graphs.state import (
    SubgraphState,
    SubgraphInput,
    SubgraphOutput,
    SingleImageResult,
    ImageInfo
)

# 导入子图需要的节点
from graphs.nodes.image_preprocess_node import image_preprocess_node
from graphs.nodes.recognize_and_correct_node import recognize_and_correct_node
from graphs.nodes.result_merge_node import result_merge_node


def wrap_result_node(
    state: SubgraphState,
    config: RunnableConfig,
    runtime: Runtime[Context]
) -> SubgraphOutput:
    """
    title: 包装子图结果
    desc: 将子图的中间状态包装为SingleImageResult输出
    integrations: 
    """
    from graphs.state import SingleImageResult
    
    image_result = SingleImageResult(
        image_index=state.image_index,
        image_info=state.image_info,
        image_url=state.image_url,
        annotations=state.annotations
    )
    
    return SubgraphOutput(image_result=image_result)


# 创建单图片处理子图
subgraph_builder = StateGraph(
    SubgraphState,
    input_schema=SubgraphInput,
    output_schema=SubgraphOutput
)

# 添加节点（优化后只需4个节点）
subgraph_builder.add_node("image_preprocess", image_preprocess_node)
subgraph_builder.add_node("recognize_and_correct", recognize_and_correct_node, metadata={"type": "agent", "llm_cfg": "config/homework_recognize_llm_cfg.json"})
subgraph_builder.add_node("result_merge", result_merge_node)
subgraph_builder.add_node("wrap_result", wrap_result_node)

# 设置入口点
subgraph_builder.set_entry_point("image_preprocess")

# 添加边（优化后流程：预处理 -> 识别+批改 -> 整合 -> 包装）
subgraph_builder.add_edge("image_preprocess", "recognize_and_correct")
subgraph_builder.add_edge("recognize_and_correct", "result_merge")
subgraph_builder.add_edge("result_merge", "wrap_result")
subgraph_builder.add_edge("wrap_result", END)

# 编译子图
single_image_subgraph = subgraph_builder.compile()
