from langgraph.graph import StateGraph, START, END
from langgraph.runtime import Runtime

from alayaflow.component.model import ModelManager

from .schemas import WorkflowInitArgs, WorkflowState, WorkflowContext, Input, Output

def mk_chat_node():
    model_manager = ModelManager()

    def chat_node(state: WorkflowState, runtime: Runtime[WorkflowContext]):
        model_id = runtime.context.chat_model_id
        chat_model = model_manager.get_model(model_id)
        if not chat_model:
            raise ValueError(f"无法找到模型ID为 '{model_id}' 的模型配置")

        messages = state["messages"].copy()
        updated_state = state.copy()

        retrieved_docs = state.get("retrieved_docs", [])
        if retrieved_docs:
            context_text = "\n\n".join([str(doc) for doc in retrieved_docs])
            from langchain_core.messages import SystemMessage
            context_message = SystemMessage(
                content=f"以下是相关的参考资料，请基于这些资料回答用户的问题：\n\n{context_text}"
            )
            messages.insert(0, context_message)

        response = chat_model.invoke(messages)
        updated_state['chat_response'] = response
        return updated_state
    return chat_node


def create_graph(init_args: WorkflowInitArgs | dict):
    if isinstance(init_args, dict):
        init_args = WorkflowInitArgs(**init_args)
    alayamem_url = init_args.alayamem_url

    graph = StateGraph(WorkflowState, WorkflowContext, input_type=Input, output_type=Output)


    graph.add_node("chat_node", mk_chat_node())


    graph.add_edge(START, "chat_node")
    graph.add_edge("chat_node",  END)

    return graph.compile()
