from langgraph.runtime import Runtime

from alayaflow.clients.alayamem.http_client import HttpAlayaMemClient
from alayaflow.component.retrieve_node import RetrieveComponent, MemoryRetrieveComponent, MemoryUpdateComponent
from alayaflow.utils.logger import AlayaFlowLogger

from ..schemas import OverallState, TaskState, WorkflowContext
from ..utils import slim_docs

logger = AlayaFlowLogger()


def create_retrieve_node(client: HttpAlayaMemClient):
    TOP_K: int = 10
    # MAX_DOC_CHARS: int = 2000
    
    def retrieve_node(state: TaskState, runtime: Runtime[WorkflowContext]):
        user_id = runtime.context.user_id
        query = state.get("query")
        if query is None or query == "":
            task = state.get("task")
            if task:
                query_parts = list(task.path) + list(task.keys)
                query = "；".join([p for p in query_parts if p])
            else:
                query = ""
        
        try:
            retrieve_component = RetrieveComponent(client=client)
            docs = retrieve_component(query=query, limit=TOP_K, user_id=user_id)
            # docs = slim_docs(docs, MAX_DOC_CHARS)
            return {"docs": docs}
        except Exception as e:
            logger.error(f"Retrieve error {type(e).__name__}: {e}")
            return {"docs": []}

    return retrieve_node


def create_memory_retrieve_node(client: HttpAlayaMemClient):
    def memory_retrieve_node(state: OverallState, runtime: Runtime[WorkflowContext]):
        user_id = runtime.context.user_id
        profile_id = runtime.context.profile_id
        if not user_id:
            return {"memory": "{}"}

        component = MemoryRetrieveComponent(client)
        memory_json = component(user_id=user_id, profile_id=profile_id)
        return {"memory": memory_json}

    return memory_retrieve_node



# def create_memory_update_node(client: HttpAlayaMemClient):
#     def memory_update_node(state: OverallState, runtime: Runtime[WorkflowContext]):
#         user_id = runtime.context.user_id
#         profile_id = runtime.context.profile_id

#         base_update = {
#             "memory_update_success": False,
#             "memory_update_resp": "{}",
#         }
        
#         if not user_id:
#             return {
#                 **base_update,
#                 "errors": {"memory_update": "missing user_id"},
#             }

#         patch = state.get("memory_patch", {}) or {}
#         if not patch:
#             return base_update

#         comp = MemoryUpdateComponent(client)
#         try:
#             success, resp_json = comp(user_id=user_id, patch=patch, profile_id=profile_id)
#         except Exception as e:
#             return {
#                 **base_update,
#                 "errors": {"memory_update": f"exception: {type(e).__name__}: {e}"},
#             }
#         update = {
#             "memory_update_success": bool(success),
#             "memory_update_resp": resp_json or "{}",
#         }

#         if not success:
#             update["errors"] = {"memory_update": "update failed"}

#         return update
#     return memory_update_node
