"""AI 营养助手 LangGraph — 基于健康报告的个性化营养建议"""
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.tools.java_client import JavaClient
from app.tools.product_tools import (
    search_product_by_keyword,
    search_article_by_keyword,
)
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
import logging

logger = logging.getLogger(__name__)


class NutritionState(TypedDict):
    query: str
    user_id: int
    conversation_id: str
    child_id: int | None
    report_id: int | None
    context: dict | None
    answer: str | None
    tasks: list[dict]
    sources: list[dict]
    messages: list | None


DEFAULT_NUTRITION_PROMPT = """你是一个儿童营养健康顾问，专注于基于检测报告提供精准的营养改善建议。

你的能力:
1. 解读儿童健康报告（菌群、营养指标、体检数据）
2. 基于报告数据推荐合适的营养产品和膳食方案
3. 结合孩子体质给出个性化的饮食建议
4. 推荐相关的健康活动和科普文章

回答原则:
1. 用中文，语气温暖专业，像营养师一样亲切
2. 必须基于报告数据给出建议，不凭空猜测
3. 引用知识库中的营养素/食物信息作为证据
4. 需要产品推荐时使用搜索工具，说明推荐理由
5. 生成 [TASK: {"title": "任务名", "dimension": "身", "points": 10, "frequency": "每天"}] 标记创建营养行动任务
6. 严重健康问题建议咨询专业医生或营养师
7. 不要编造医学建议，不夸大效果
"""


def create_nutrition_graph():
    """创建营养助手 StateGraph"""
    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()
    memory_mgr = MemoryManager(embeddings)
    java = JavaClient()

    llm_with_tools = llm.bind_tools([
        search_product_by_keyword,
        search_article_by_keyword,
    ])

    builder = StateGraph(NutritionState)

    async def generate_answer(state: NutritionState) -> dict:
        """核心 LLM 调用 + 画像注入"""
        # 加载营养助手 prompt
        nutrition_prompt = await get_prompt("nutrition_assistant") or DEFAULT_NUTRITION_PROMPT
        messages = [SystemMessage(content=nutrition_prompt)]

        # 画像注入（在人格 prompt 之后）
        member_id = state.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:
            ctx_parts = []
            if ctx.get("child_id"):
                ctx_parts.append(f"孩子ID: {ctx['child_id']}")
            if ctx.get("family_id"):
                ctx_parts.append(f"家庭ID: {ctx['family_id']}")
            if ctx.get("report_id"):
                ctx_parts.append(f"报告ID: {ctx['report_id']}")
            if ctx_parts:
                messages.append(SystemMessage(content="用户背景信息: " + ", ".join(ctx_parts)))

        # 长期记忆
        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_with_tools.ainvoke(messages)
        answer = response.content

        # 提取任务和来源
        tasks = []
        sources = []
        if response.response_metadata.get("tool_calls"):
            for tc in response.response_metadata["tool_calls"]:
                sources.append({
                    "type": "tool",
                    "name": tc.get("name", ""),
                    "input": tc.get("args", {}),
                })

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

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

    # 构建图
    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