import json
from typing import Dict, Any
from datetime import datetime

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.types import interrupt

from pymongo import MongoClient
from src.core.config import MONGO_URI
from src.server.agent.agents.init_llm import persona_agent
from src.server.agent.memory.memory_manager import MemoryManager
from src.server.agent.state import AgentState

client = MongoClient(MONGO_URI)
db = client["furniture_agent_db"]
persona_collection = db["user_persona"]

EXTRACTION_PROMPT = ChatPromptTemplate.from_messages([
    ("system", """你是一个专为家具设计师服务的用户画像提取专家。

当前已知用户画像（JSON格式）：
{current_persona_json}

任务：
1. 从下面所有用户消息中，提取或更新用户对家具设计的偏好信息。
   只提取明确提到或强烈暗示的内容，不要臆想或添加默认值。
2. 输出更新后的完整 persona JSON（只包含有值字段）。
   推荐使用的键名（优先使用这些）：
   - "风格偏好"        （如 "北欧"、"极简"、"工业风"）
   - "家具类型"         （如 "沙发"、"餐桌"、"书柜"、"办公椅"）
   - "颜色偏好"         （如 "原木色"、"白色"、"深灰"、"大地色系"）
   - 其他可选键：预算范围、空间大小、材质偏好、使用场景等

3. 判断当前画像是否“足够完整”来支持针对家具设计师的市场趋势报告：
   - **必须条件**：至少包含 "风格偏好"、"家具类型"、"颜色偏好" 中的 2 项以上
   - 如果缺少 2 项或以上核心信息，返回 complete: false，并生成一个自然、礼貌、具体的追问（中文），优先询问缺失的核心项
   - 如果核心三项中至少有 2 项已明确，返回 complete: true，不需要追问

输出必须是严格的 JSON 对象，包含三个字段：
- "persona": 一个对象，键是画像属性，值是对应的字符串（或数组，如果有多个偏好）
- "complete": 布尔值 true 或 false
- "question": 字符串，如果 complete 为 true 则为空字符串，否则是具体的追问句子

示例输出结构（仅供参考，不要直接复制）：
{{
  "persona": {{
    "风格偏好": "北欧简约",
    "家具类型": "沙发",
    "颜色偏好": "原木色 + 浅灰"
  }},
  "complete": true,
  "question": ""
}}
"""),
    ("placeholder", "{messages}"),
])


def get_persona_from_mongo(thread_id: str) -> Dict[str, Any]:
    doc = persona_collection.find_one({"thread_id": thread_id}, sort=[("_id", -1)])
    if doc and "persona" in doc:
        return doc["persona"]
    return {}


def save_persona_to_mongo(thread_id: str, persona: Dict[str, Any], is_complete: bool):
    try:
        result = persona_collection.update_one(
            {"thread_id": thread_id},
            {
                "$set": {
                    "persona": persona,
                    "persona_complete": is_complete,
                    "updated_at": datetime.utcnow(),
                    "last_update_reason": "persona_node_update"
                }
            },
            upsert=True
        )
        print(f"[Persona Save] thread_id: {thread_id} | matched: {result.matched_count} | modified: {result.modified_count} | upserted: {result.upserted_id}")
    except Exception as e:
        print(f"[Persona Save Error] {e}")


def persona_node(state: AgentState, config: RunnableConfig):
    thread_id = config["configurable"]["thread_id"]

    # 读取已有画像（MongoDB 优先）
    persisted_persona = get_persona_from_mongo(thread_id)
    current_persona = state.get("persona", persisted_persona)

    # messages = state["messages"]
    messages = MemoryManager.build_llm_context(state)
    current_persona_json = json.dumps(current_persona, ensure_ascii=False, indent=None)

    chain = EXTRACTION_PROMPT | persona_agent
    result = chain.invoke({
        "current_persona_json": current_persona_json,
        "messages": messages,
    })

    updated_persona = result.persona
    is_complete = result.complete
    question = (result.question or "").strip()

    updates = {
        "persona": updated_persona,
        "persona_complete": is_complete,
        "persona_summary": json.dumps(updated_persona, ensure_ascii=False, indent=2),
    }

    # 持久化到 MongoDB
    save_persona_to_mongo(thread_id, updated_persona, is_complete)

    if is_complete:
        updates["messages"] = messages + [AIMessage(
            content=(
                "用户画像已足够完整（风格、家具类型、颜色偏好已明确），并已保存到项目记录。\n\n"
                "接下来是否需要我为您生成一份针对当前风格与家具类型的市场趋势报告？\n"
                "回复“是”或“需要”即可开始生成；回复“不需要”或“先不用”则直接进入家具设计阶段。"
            )
        )]
        return updates

    # 不完整 → 询问（优先问核心三项）
    if not question:
        missing = []
        core_keys = ["风格偏好", "家具类型", "颜色偏好"]
        for key in core_keys:
            if key not in updated_persona or not updated_persona[key]:
                missing.append(key)
        if missing:
            question = f"为了更好地为您生成趋势报告，能否补充一下您对{'、'.join(missing)}的偏好呢？"
        else:
            question = "您对家具的风格、类型或颜色有什么特别的偏好吗？可以多说一些～"

    updated_messages = messages + [AIMessage(content=question)]

    approved = interrupt({
        **updates,
        "messages": updated_messages,
        "persona_complete": False,
        "__interrupt__": {
            "type": "persona_question",
            "question": question,
            "node": "Persona",
            "wait_for": "human_response",
            # 可选：当前画像快照，便于前端显示或调试
            "current_persona_snapshot": updated_persona
        }
    })
    return approved