import json
from typing import Any, Dict
from langgraph.runtime import Runtime
from alayaflow.clients.alayamem.http_client import HttpAlayaMemClient
from alayaflow.component.llm_node import LLMComponent, ResponseFormat
from alayaflow.component.retrieve_node import MemoryUpdateComponent
from alayaflow.utils.logger import AlayaFlowLogger

from ..schemas import OverallState, TaskState, WorkflowContext

logger = AlayaFlowLogger()

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"},
            }

        form_patch = state.get("memory_patch", {}) or {}
        if not form_patch:
            # 如果没有需要更新的内容，视为成功
            return {
                "memory_update_success": True,
                "memory_update_resp": "{}",
            }

        # 1) Parse old memory (string -> dict), fallback to {}
        parse_err = None
        old_memory: Dict[str, Any] = {}
        try:
            mem_str = (state.get("memory") or "").strip()
            if mem_str:
                obj = json.loads(mem_str)
                if isinstance(obj, dict):
                    old_memory = obj
                else:
                    parse_err = "memory is not a JSON object"
        except Exception as e:
            parse_err = f"invalid memory json: {type(e).__name__}: {e}"
            old_memory = {}

        # 2) LLM prompt：直接输出最终 memory JSON
        MAX_CHARS = 180_000  # 按你们后端限制调整

        system_prompt = f"""
        你是一个“记忆管理员”，你全权负责把本次表单信息融合进历史记忆。

        你必须【只输出一个 JSON 对象】，该 JSON 对象本身就是最终的 memory。
        不要输出任何解释文字、不要输出 Markdown、不要包裹在 new_memory 或其它字段中。

        你可以做的事情：
        - 字段改名、同义键归一（如 name / 姓名 / 名字）
        - 结构层级调整（扁平 -> 嵌套，或反之）
        - 合并重复或等价信息
        - 删除明显冗余、低价值、重复字段

        必须遵守的规则：
        1) 不要凭空编造未提供的事实。
        2) 若历史记忆与表单信息冲突，优先保留更“具体 / 更新 / 可信”的值；不确定则保留旧值。
        3) 尽量保持结构稳定、简洁、可长期演进。
        4) 避免重复语义字段（同义字段必须归一）。
        5) 输出的整个 JSON 文本字符数【必须 ≤ {MAX_CHARS}】。

        如果内容可能超限，你必须主动裁剪：
        - 优先删除低价值信息、长文本、冗余字段
        - 对长文本进行摘要或截断
        - 对长列表只保留最少必要项
        - 永远优先保留结构化、确定性的关键信息
        """.strip()

        user_prompt = json.dumps(
            {
                "old_memory": old_memory,
                "form_patch": form_patch,
            },
            ensure_ascii=False,
        )

        # 3) Call LLM
        try:
            llm = LLMComponent(
                model_id=runtime.context.extract_model_id,
                system_prompt=system_prompt,
                prompt=user_prompt,
                response_format=ResponseFormat.JSON,
                temperature=0.6,
            )
            msg = llm()
            new_memory = json.loads(msg.content) if msg and msg.content else None
        except Exception as e:
            err = {"memory_update": f"llm exception: {type(e).__name__}: {e}"}
            if parse_err:
                err["memory_parse"] = parse_err
            return {**base_update, "errors": err}

        # 4) Validate LLM output
        if not isinstance(new_memory, dict):
            err = {
                "memory_update": "llm output is invalid (must be a JSON object)",
                "llm_raw": (msg.content[:2000] if msg and msg.content else ""),
            }
            if parse_err:
                err["memory_parse"] = parse_err
            return {**base_update, "errors": err}

        try:
            new_memory_str = json.dumps(new_memory, ensure_ascii=False)
        except Exception as e:
            err = {"memory_update": f"new memory not serializable: {type(e).__name__}: {e}"}
            if parse_err:
                err["memory_parse"] = parse_err
            return {**base_update, "errors": err}

        if len(new_memory_str) > MAX_CHARS:
            err = {
                "memory_update": f"llm output exceeds size limit: {len(new_memory_str)} chars"
            }
            if parse_err:
                err["memory_parse"] = parse_err
            return {**base_update, "errors": err}

        # 5) Update backend with final memory
        comp = MemoryUpdateComponent(client)
        try:
            resp_json, success = comp(
                user_id=user_id,
                patch=new_memory,
                profile_id=profile_id,
            )
        except Exception as e:
            err = {"memory_update": f"exception: {type(e).__name__}: {e}"}
            if parse_err:
                err["memory_parse"] = parse_err
            return {**base_update, "errors": err}

        update = {
            "memory_update_success": bool(success),
            "memory_update_resp": resp_json or "{}",
        }

        # 只有更新成功时才更新 state.memory，保持与后端一致
        if success:
            update["memory"] = new_memory_str

        if parse_err:
            update.setdefault("errors", {})
            update["errors"]["memory_parse"] = parse_err

        if not success:
            update.setdefault("errors", {})
            update["errors"]["memory_update"] = "update failed"

        return update

    return memory_update_node
