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

from alayaflow.component.memory import query_file, commit_turn, turns, summary, profile
from alayaflow.component.model import ModelManager

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

def mk_turns_node(alayamem_url: str):
    def turns_node(state: WorkflowState, runtime: Runtime[WorkflowContext]):
        session_id = runtime.context.session_id
        user_id = runtime.context.user_id
        original_result = turns(alayamem_url, session_id, user_id)

        # 打印返回字段
        print(f"\n[1. turns 接口返回]")
        print(f"  返回类型: {type(original_result)}")
        print(f"  返回值: {original_result}")

        updated_state = state.copy()
        # turns 接口直接返回列表
        updated_state["history_turns"] = original_result if isinstance(original_result, list) else []
        return updated_state
    return turns_node

def mk_query_file_node(alayamem_url: str):
    def query_file_node(state: WorkflowState, runtime: Runtime[WorkflowContext]):
        user_id = runtime.context.user_id
        messages = state.get("messages", [])
        message = messages[-1] if messages else None
        original_result = query_file(alayamem_url, user_id, message)

        # 打印返回字段
        print(f"\n[2. query_file 接口返回]")
        print(f"  返回类型: {type(original_result)}")
        if isinstance(original_result, dict):
            print(f"  字段列表: {list(original_result.keys())}")
        print(f"  返回值: {original_result}")

        updated_state = state.copy()
        # query_file 返回 dict，results 是 dict，包含 documents 字段
        if isinstance(original_result, dict):
            results = original_result.get("results", {})
            updated_state["retrieved_docs"] = results.get("documents", []) if isinstance(results, dict) else []
        else:
            updated_state["retrieved_docs"] = []
        return updated_state
    return query_file_node

def mk_summary_node(alayamem_url: str):
    def summary_node(state: WorkflowState, runtime: Runtime[WorkflowContext]):
        session_id = runtime.context.session_id
        user_id = runtime.context.user_id
        original_result = summary(alayamem_url, session_id, user_id)

        # 打印返回字段
        print(f"\n[2. summary 接口返回]")
        print(f"  返回类型: {type(original_result)}")
        if isinstance(original_result, dict):
            print(f"  字段列表: {list(original_result.keys())}")
        print(f"  返回值: {original_result}")

        updated_state = state.copy()
        # summary 可能是 None，需要处理
        if isinstance(original_result, dict):
            summary_value = original_result.get("summary", "")
            updated_state["session_summary"] = summary_value if summary_value is not None else ""
        else:
            updated_state["session_summary"] = ""
        return updated_state
    return summary_node

def mk_profile_node(alayamem_url: str):
    def profile_node(state: WorkflowState, runtime: Runtime[WorkflowContext]):
        user_id = runtime.context.user_id
        original_result = profile(alayamem_url, user_id)

        # 打印返回字段
        print(f"\n[3. profile 接口返回]")
        print(f"  返回类型: {type(original_result)}")
        if isinstance(original_result, dict):
            print(f"  字段列表: {list(original_result.keys())}")
        print(f"  返回值: {original_result}")

        updated_state = state.copy()
        # profile 可能是 None，需要处理
        if isinstance(original_result, dict):
            profile_value = original_result.get("profile", {})
            updated_state["user_profile"] = profile_value if profile_value is not None else {}
        else:
            updated_state["user_profile"] = {}
        return updated_state
    return profile_node

def mk_commit_turn_node(alayamem_url: str):
    def commit_turn_node(state: WorkflowState, runtime: Runtime[WorkflowContext]):
        session_id = runtime.context.session_id
        user_id = runtime.context.user_id
        messages = state.get("messages", [])
        chat_response = state.get("chat_response")

        # 提取用户消息和助手回复
        user_text = ""
        if messages:
            last_msg = messages[-1]
            if hasattr(last_msg, 'content'):
                user_text = last_msg.content
            elif isinstance(last_msg, dict):
                user_text = last_msg.get('content', '')

        assistant_text = ""
        if chat_response:
            if hasattr(chat_response, 'content'):
                assistant_text = chat_response.content
            elif isinstance(chat_response, dict):
                assistant_text = chat_response.get('content', '')

        # 调用 commit_turn
        original_result = commit_turn(alayamem_url, session_id, user_text,
                                     assistant_text, user_id, window_size=3)

        # 打印返回字段
        print(f"\n[4. commit_turn 接口返回]")
        print(f"  返回类型: {type(original_result)}")
        if isinstance(original_result, dict):
            print(f"  字段列表: {list(original_result.keys())}")
        print(f"  返回值: {original_result}")

        return state.copy()
    return commit_turn_node

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}' 的模型配置")

        from langchain_core.messages import SystemMessage, HumanMessage, AIMessage

        # 获取 prompt 模板配置
        ctx = runtime.context
        system_prompt = ctx.system_prompt
        user_profile_prompt = ctx.user_profile_prompt
        session_summary_prompt = ctx.session_summary_prompt
        retrieved_docs_prompt = ctx.retrieved_docs_prompt
        history_turns_prompt = ctx.history_turns_prompt
        context_wrapper_prompt = ctx.context_wrapper_prompt

        # 构建完整的消息列表
        full_messages = []

        # 添加系统基础提示词
        if system_prompt:
            full_messages.append(SystemMessage(content=system_prompt))

        # 构建上下文信息部分
        context_parts = []

        # 添加用户画像
        # user_profile = state.get("user_profile", {})
        # if user_profile:
        #     profile_text = user_profile_prompt.format(user_profile=str(user_profile))
        #     context_parts.append(profile_text)

        # 添加会话摘要
        session_summary = state.get("session_summary", "")
        if session_summary:
            summary_text = session_summary_prompt.format(session_summary=session_summary)
            context_parts.append(summary_text)

        # 添加检索文档
        retrieved_docs = state.get("retrieved_docs", [])
        if retrieved_docs:
            docs_content = "\n\n".join([str(doc) for doc in retrieved_docs])
            docs_text = retrieved_docs_prompt.format(retrieved_docs=docs_content)
            context_parts.append(docs_text)

        # 添加历史对话
        history_turns = state.get("history_turns", [])
        if history_turns:
            history_content = "\n".join([
                f"用户: {turn.get('user_text', '')}\n助手: {turn.get('assistant_text', '')}"
                for turn in history_turns
            ])
            history_text = history_turns_prompt.format(history_turns=history_content)
            context_parts.append(history_text)

        # 如果有上下文信息，添加上下文系统消息
        if context_parts:
            context_content = context_wrapper_prompt.format(
                context_parts="\n\n".join(context_parts)
            )
            full_messages.append(SystemMessage(content=context_content))

        # 添加当前用户消息
        current_messages = state["messages"]
        full_messages.extend(current_messages)

        # 打印最终输入到模型的消息
        print(f"\n[5. 模型输入消息]")
        print(f"  消息数量: {len(full_messages)}")
        print("-" * 50)
        full_content = ""
        for msg in full_messages:
            msg_type = type(msg).__name__
            content = msg.content if hasattr(msg, 'content') else str(msg)
            full_content += f"[{msg_type}]\n{content}\n\n"
        print(full_content)
        print("-" * 50)

        # 调用模型
        response = chat_model.invoke(full_messages)

        updated_state = state.copy()
        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("turns_node", mk_turns_node(alayamem_url))
    graph.add_node("query_file_node", mk_query_file_node(alayamem_url))
    # 
    graph.add_node("summary_node", mk_summary_node(alayamem_url))
    graph.add_node("profile_node", mk_profile_node(alayamem_url))
    graph.add_node("chat_node", mk_chat_node())
    graph.add_node("commit_turn_node", mk_commit_turn_node(alayamem_url))

    # 添加边：获取历史 -> 获取摘要 -> 获取画像 -> 聊天 -> 提交记录
    graph.add_edge(START, "turns_node")
    graph.add_edge("turns_node",  "summary_node")
    graph.add_edge("summary_node", "profile_node")
    graph.add_edge("profile_node", "chat_node")
    graph.add_edge("chat_node", "commit_turn_node")
    graph.add_edge("commit_turn_node", END)

    return graph.compile()
