from typing import TypedDict, List, Optional, ClassVar

from langchain_core.messages import BaseMessage, AIMessageChunk
from langgraph.graph import StateGraph, START, END

from alayaflow.component.memory import init_memory, query_message, add_message, query_vdb_message
from alayaflow.component.chat_model import mk_chat_model_deepseek

class WorkflowState(TypedDict):
    user_id: str
    session_id: str
    messages: List[BaseMessage]
    memory_initialized: Optional[bool]
    stream_chunks: List[AIMessageChunk]
    chat_response: Optional[dict]
    retrieved_docs: Optional[List[str]]
    context: Optional[str]

def mk_init_memory_node(alayamem_url: str):
    def init_memory_node(state: WorkflowState):
        user_id = state["user_id"]
        session_id = state["session_id"]
        original_result = init_memory(alayamem_url, user_id, session_id)
        updated_state = state.copy()
        updated_state["memory_initialized"] = original_result.get("status", "") == "success"
        return updated_state
    return init_memory_node

# Keep for integration
# def mk_query_message_node(alayamem_url: str):
#     def query_message_node(state: WorkflowState):
#         user_id = state["user_id"]
#         session_id = state["session_id"]
#         messages = state.get("messages", [])
#         original_result = query_message(alayamem_url, user_id, session_id, messages)
#         updated_state = state.copy()
#         updated_state["context"] = original_result.get("context", "")
#         return updated_state
#     return query_message_node

def mk_query_vdb_message_node(alayamem_url: str):
    def query_vdb_message_node(state: WorkflowState):
        messages = state.get("messages", [])
        limit = state.get("limit", 5)
        original_result = query_vdb_message(alayamem_url, messages, limit)
        updated_state = state.copy()
        updated_state["retrieved_docs"] = original_result.get("vdb_results", [])
        return updated_state
    return query_vdb_message_node

def mk_chat_node():
    chat_model = mk_chat_model_deepseek()
    def chat_node(state: WorkflowState):
        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 mk_add_message_node(alayamem_url: str):
    def add_message_node(state: WorkflowState):
        user_id = state["user_id"]
        session_id = state["session_id"]
        messages = state.get("messages", [])
        add_message(alayamem_url, user_id, session_id, messages)
        return state.copy()
    return add_message_node

def get_state_schema() -> ClassVar[WorkflowState]:
    return WorkflowState

def create_graph(init_args: dict):
    alayamem_url = init_args.get("alayamem_url")
    if not alayamem_url:
        raise ValueError("alayamem_url is required")

    graph = StateGraph(WorkflowState)

    graph.add_node("init_memory_node", mk_init_memory_node(alayamem_url))
    graph.add_node("query_vdb_message_node", mk_query_vdb_message_node(alayamem_url))
    graph.add_node("chat_node", mk_chat_node())
    graph.add_node("add_message_node", mk_add_message_node(alayamem_url))

    graph.add_edge(START, "init_memory_node")
    graph.add_edge("init_memory_node", "query_vdb_message_node")
    graph.add_edge("query_vdb_message_node", "chat_node")
    graph.add_edge("chat_node", "add_message_node")
    graph.add_edge("add_message_node", END)

    return graph.compile()
