"""
AI 健康教练 LangGraph - 专为儿童健康咨询设计的对话图
Uses RAG retrieval from health knowledge base.
"""
from typing import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage
from app.rag.retriever import RagRetriever
from app.memory.store import MemoryManager
from app.config import settings
from app.prompt_service import get_prompt
from app.portrait_service import get_user_portrait_text
from app.tools.java_client import JavaClient
import logging

logger = logging.getLogger(__name__)


class HealthCoachState(TypedDict):
    query: str
    user_id: int
    child_id: int | None
    conversation_id: str | None
    context: dict | None
    answer: str | None
    sources: list[dict]
    memory_messages: list | None


DEFAULT_COACH_PROMPT = """你是一个儿童健康教练，喜欢用提问的方式引导家长发现问题和解决方案。

核心原则：
1. 先问后答：不要直接给答案，用问题引导家长思考
   - "孩子这种情况多久了？"、"平时饮食习惯是怎样的？"、"体重身高发育曲线有没有记录？"
2. 用开放式问题了解情况，再给针对性建议
3. 语气温暖，像一个耐心的教练在陪伴家长成长

你的职责：
1. 儿童体质调理（如挑食、瘦弱、易疲劳、反复生病）—— 用提问了解情况
2. 科学喂养和营养建议 —— 先问饮食现状再给建议
3. 日常保健（睡眠、运动、季节防护）—— 引导家长发现生活中的问题
4. 体质辨识基础咨询 —— 通过提问帮助家长理解孩子体质特点
5. 健康问题判断 —— 先收集症状信息再建议是否就医

回答结构（教练风格）：
1. 提问：先问1-2个关键问题了解情况
2. 分析：简短总结你观察到的问题
3. 建议：给出2-3个可操作的建议
4. 跟进问句：留一个后续跟进的问题，体现教练陪伴感

原则：
- 涉及严重症状（急性腹痛、高烧不退、严重外伤等）立刻建议就医，不提问
- 不确定时明确说"建议咨询专业医生"
- 始终用中文回答
"""


def create_health_coach_graph():
    """创建AI健康教练 StateGraph"""
    # LLM
    llm = ChatOpenAI(
        model=settings.llm_model,
        api_key=settings.llm_api_key,
        base_url=settings.llm_base_url,
        temperature=settings.llm_temperature,
    )

    # 嵌入和检索
    from app.rag.embeddings import get_embeddings
    embeddings = get_embeddings()
    retriever = RagRetriever(collection_name="health_knowledge")
    memory_mgr = MemoryManager(embeddings)
    java = JavaClient()

    # 带工具的 LLM（健康知识库检索）
    async def health_retrieve(query: str, user_id: int) -> list[dict]:
        try:
            results = await retriever.retrieve(
                query,
                filters={"user_id": user_id},
                k=3,
            )
            return results
        except Exception as e:
            logger.warning("健康知识检索失败: %s", e)
            return []

    async def generate_answer(state: HealthCoachState) -> dict:
        """带健康知识检索的 LLM 生成"""
        # 人格路由：context.coach_id (xibao/fubao) -> 专属prompt，三级回退保证可用性
        _ctx = state.get("context") or {}
        _coach_id = _ctx.get("coach_id")
        if _coach_id not in ("xibao", "fubao"):
            _coach_id = None
        _persona_prompt = await get_prompt(f"health_coach_{_coach_id}") if _coach_id else None
        _base_prompt = await get_prompt("health_coach")
        messages = [SystemMessage(content=_persona_prompt or _base_prompt or DEFAULT_COACH_PROMPT)]

        # 画像注入
        member_id = _ctx.get("child_id")
        portrait_text = await get_user_portrait_text(java, state["user_id"], member_id)
        if portrait_text:
            messages.insert(1, SystemMessage(content=portrait_text))

        # 家庭上下文
        ctx = state.get("context") or {}
        if ctx:
            child_info = []
            if ctx.get("child_id"):
                child_info.append(f"孩子ID: {ctx['child_id']}")
            if ctx.get("family_id"):
                child_info.append(f"家庭ID: {ctx['family_id']}")
            if child_info:
                messages.append(SystemMessage(
                    content="用户背景信息: " + ", ".join(child_info)
                ))

        # 健康知识检索
        retrieved = await health_retrieve(state["query"], state["user_id"])
        sources = []
        if retrieved:
            context_texts = []
            for i, r in enumerate(retrieved[:3]):
                src = r.get("metadata", {})
                title = src.get("title", "健康知识")
                content = r["content"][:300]
                context_texts.append(f"[{i+1}] {title}:\n{content}")
                sources.append({
                    "type": "knowledge",
                    "title": title,
                    "content": content,
                })
            knowledge_context = "\n\n---\n".join(context_texts)
            messages.append(SystemMessage(
                content=f"以下是与问题相关的健康知识参考：\n{knowledge_context}"
            ))

        # 长期记忆
        try:
            memories = await memory_mgr.recall(state["user_id"], state["query"])
            if memories:
                mem_text = "\n".join([f"- {m}" for m in memories[:3]])
                messages.append(SystemMessage(
                    content=f"该用户的历史健康咨询记录:\n{mem_text}"
                ))
        except Exception as e:
            logger.warning("召回健康记忆失败: %s", e)

        # 用户问题
        messages.append(HumanMessage(content=state["query"]))

        response = await llm.ainvoke(messages)
        answer = response.content

        return {
            "answer": answer,
            "sources": sources,
            "memory_messages": [
                {"role": "user", "content": state["query"]},
                {"role": "assistant", "content": answer},
            ],
        }

    async def save_memory(state: HealthCoachState) -> dict:
        """保存对话到长期记忆"""
        try:
            if state.get("memory_messages"):
                await memory_mgr.save_conversation(
                    state["user_id"],
                    state.get("conversation_id") or f"health_{state['user_id']}",
                    state["memory_messages"],
                )
        except Exception as e:
            logger.warning("保存健康教练记忆失败: %s", e)
        return {}

    # ── 构建图 ──

    builder = StateGraph(HealthCoachState)

    builder.add_node("generate_answer", generate_answer)
    builder.add_node("save_memory", save_memory)

    builder.add_edge(START, "generate_answer")
    builder.add_edge("generate_answer", "save_memory")
    builder.add_edge("save_memory", END)

    checkpointer = MemorySaver()
    graph = builder.compile(checkpointer=checkpointer)

    return graph