"""
健康方案生成 LangGraph - 分步骤结构化方案生成

工作流:
1. 数据收集: 家庭成员信息 + 健康指标 + 知识库检索
2. LLM 生成: 总览概述 + 营养/饮食/运动 三个 section
3. 返回结构化 JSON

每次调用 LLM 时遵循统一的数据组装方式:
  - 用户指标 (来自 Java context API)
  - 知识库参考 (RAG 检索)
  - 系统提示词 (带输出格式模板)
"""
from typing import TypedDict, Literal
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage
from app.rag.retriever import RagRetriever
from app.config import settings
from app.tools.java_client import JavaClient
from app.prompt_service import get_prompt
import logging

logger = logging.getLogger(__name__)

SECTION_KEYS = ["nutrition", "diet", "exercise"]

DEFAULT_PROMPT = """你是一个专业的家庭健康方案规划师。根据用户的健康数据和目标，生成结构化的改善方案。

## 输出格式（必须严格遵守 JSON）

```json
{
  "overview": "总体概述（100字以内，说明方案目标和核心策略）",
  "sections": [
    {
      "key": "nutrition",
      "title": "营养补充建议",
      "content": "Markdown 格式的详细内容，包含具体产品推荐、用量、服用时间",
      "items": [
        {"name": "产品名", "dosage": "用量", "timing": "服用时间", "reason": "推荐理由"}
      ]
    },
    {
      "key": "diet", 
      "title": "饮食建议",
      "content": "Markdown 格式的餐饮建议，包含早餐/午餐/晚餐建议",
      "items": [{"meal": "餐型", "food": "食物建议", "notes": "注意事项"}]
    },
    {
      "key": "exercise",
      "title": "运动计划",
      "content": "Markdown 格式的运动建议，包含频率、时长、类型",
      "items": [{"type": "运动类型", "duration": "时长", "frequency": "频率", "notes": "注意事项"}]
    }
  ]
}
```

## 原则
1. 基于实际数据给出建议，不编造
2. 引用知识库时标注来源
3. 建议要具体可执行，避免空泛
4. 营养补充部分要具体到品牌/产品类型和用量
5. 严重健康问题建议咨询医生
"""

REGENERATE_SECTION_PROMPT = """你是一个家庭健康方案规划师。根据用户反馈重新生成指定部分的内容。

## 当前方案内容
{existing_section_content}

## 用户反馈
{feedback}

## 相关背景数据
{context_summary}

请重新生成该部分内容，保持与原格式一致。只输出新的 content 字段值（Markdown 格式），不需要输出 JSON 结构。"""


class PlanState(TypedDict):
    member_ids: str
    dimensions: str
    goal: str
    family_id: int
    members_info: dict
    indicators: list
    abnormal_indicators: list
    kb_results: list
    overview: str
    nutrition_section: str
    diet_section: str
    exercise_section: str
    full_response: dict
    error: str


async def collect_data(state: PlanState) -> dict:
    """Step 1: 收集用户数据 + 知识库检索"""
    java = JavaClient()
    retriever = RagRetriever(collection_name="cfc_knowledge")

    member_ids_str = state.get("member_ids", "")
    member_ids = [m.strip() for m in member_ids_str.split(",") if m.strip()]
    goal = state.get("goal", "")
    dimensions = state.get("dimensions", "") or ""

    # 1a. 获取家庭成员信息
    members_info = []
    for uid_str in member_ids:
        ctx = await java.get_family_context(int(uid_str), "child_info")
        children = ctx.get("children", []) if isinstance(ctx, dict) else []
        for child in children:
            cid = str(child.get("用户ID", ""))
            if cid == uid_str:
                members_info.append({
                    "id": cid,
                    "name": child.get("姓名", f"成员{cid}"),
                    "age": child.get("年龄", "未知"),
                    "energy": child.get("能量", 0),
                })
                break
        if not any(m["id"] == uid_str for m in members_info):
            members_info.append({"id": uid_str, "name": f"成员{uid_str}", "age": "未知", "energy": 0})

    # 1b. 获取健康指标
    all_indicators = []
    abnormal_list = []
    for member in members_info:
        reports = await java.get_member_reports(int(member["id"]))
        if not reports:
            continue
        latest = max(reports, key=lambda r: r.get("reportDate", ""))
        report_id = latest.get("id")
        indicators = await java.get_report_indicators(report_id)
        for ind in indicators:
            ind["_member_id"] = member["id"]
            ind["_member_name"] = member["name"]
            all_indicators.append(ind)

    # 1c. 识别异常指标
    known_indicators = {}
    for ind in all_indicators:
        name = ind.get("indicatorName", "").strip()
        if not name or name in known_indicators:
            continue
        for itype in ["indicator", "bacteria", "nutrient"]:
            kb = await java.query_health_knowledge(itype, name)
            if kb:
                known_indicators[name] = kb
                break

    for ind in all_indicators:
        name = ind.get("indicatorName", "")
        status = ind.get("status", "")
        kb = known_indicators.get(name, {})
        entry = {
            "member": ind.get("_member_name", ""),
            "indicator": name,
            "value": ind.get("indicatorValue", ""),
            "unit": ind.get("unit", ""),
            "ref_range": kb.get("normalRange", ind.get("refRange", "")),
            "description": kb.get("description", ""),
            "suggestion": kb.get("suggestion", ""),
        }
        if status in ("abnormal", "high", "low", "偏高", "偏低"):
            abnormal_list.append(entry)

    # 1d. 知识库检索
    kb_results = []
    queries = set()
    for ind in abnormal_list:
        queries.add(ind["indicator"])
    queries.add(goal)
    if dimensions:
        queries.add(dimensions)
    for q in list(queries)[:8]:
        results = await retriever.retrieve(q, k=3)
        kb_results.extend(results)

    return {
        "members_info": members_info,
        "indicators": all_indicators,
        "abnormal_indicators": abnormal_list,
        "kb_results": kb_results,
    }


async def build_prompt(state: PlanState) -> str:
    """组装 LLM prompt"""
    base_prompt = await get_prompt("health_plan") or DEFAULT_PROMPT
    parts = [base_prompt]

    parts.append(f"\n## 用户目标\n{state['goal']}")
    if state.get("dimensions"):
        parts.append(f"\n## 重点关注维度\n{state['dimensions']}")

    parts.append("\n## 家庭成员")
    for m in state["members_info"]:
        parts.append(f"- {m['name']} (年龄: {m['age']})")

    if state["abnormal_indicators"]:
        parts.append("\n## 异常指标")
        for ind in state["abnormal_indicators"][:8]:
            parts.append(
                f"- {ind['member']} - {ind['indicator']}: {ind['value']}{ind.get('unit','')} "
                f"(参考: {ind['ref_range']})"
            )
            if ind.get("description"):
                parts.append(f"  说明: {ind['description']}")

    if state["kb_results"]:
        parts.append("\n## 知识库参考")
        for r in state["kb_results"][:6]:
            title = r.get("metadata", {}).get("title", "")
            content = r.get("content", "")[:200]
            parts.append(f"---\n{title}\n{content}")

    return "\n".join(parts)


async def generate_plan(state: PlanState) -> dict:
    """Step 2: 调用 LLM 生成结构化方案"""
    llm = ChatOpenAI(
        model=settings.llm_model,
        api_key=settings.llm_api_key,
        base_url=settings.llm_base_url,
        temperature=0.3,
    )
    prompt = await build_prompt(state)
    messages = [SystemMessage(content=prompt)]
    try:
        response = await llm.ainvoke(messages)
        answer = response.content
        # 解析 JSON
        import json
        try:
            # 提取 JSON 块
            start = answer.find("{")
            end = answer.rfind("}") + 1
            if start >= 0 and end > start:
                json_str = answer[start:end]
                parsed = json.loads(json_str)
                return {"full_response": parsed, "overview": parsed.get("overview", "")}
        except (json.JSONDecodeError, Exception) as e:
            logger.warning("解析方案 JSON 失败，使用原始文本: %s", e)
        return {"full_response": {"raw": answer}, "overview": answer[:200]}
    except Exception as e:
        logger.error("LLM 生成方案失败: %s", e)
        return {"error": str(e)}


async def regenerate_section(state: PlanState) -> dict:
    """重新生成指定 section"""
    section_key = state.get("section", "nutrition")
    feedback = state.get("feedback", "")
    existing_content = state.get("existing_section_content", "")
    llm = ChatOpenAI(
        model=settings.llm_model,
        api_key=settings.llm_api_key,
        base_url=settings.llm_base_url,
        temperature=0.3,
    )

    # 构建上下文摘要
    ctx_parts = []
    for m in state.get("members_info", []):
        ctx_parts.append(f"- {m['name']} (年龄: {m['age']})")
    if state.get("goal"):
        ctx_parts.append(f"目标: {state['goal']}")
    if state.get("abnormal_indicators"):
        for ind in state["abnormal_indicators"][:5]:
            ctx_parts.append(f"- {ind['member']}: {ind['indicator']}={ind['value']}")

    regenerate_template = await get_prompt("health_plan_regenerate") or REGENERATE_SECTION_PROMPT
    try:
        prompt = regenerate_template.format(
            existing_section_content=existing_content[:500],
            feedback=feedback,
            context_summary="\n".join(ctx_parts),
        )
    except (KeyError, IndexError, ValueError):
        logger.warning("Java 配置的 health_plan_regenerate 模板缺少占位符，回退本地模板")
        prompt = REGENERATE_SECTION_PROMPT.format(
            existing_section_content=existing_content[:500],
            feedback=feedback,
            context_summary="\n".join(ctx_parts),
        )
    try:
        response = await llm.ainvoke([SystemMessage(content=prompt)])
        return {"regenerated_content": response.content}
    except Exception as e:
        logger.error("重新生成方案 section 失败: %s", e)
        return {"error": str(e)}


def create_health_plan_graph():
    builder = StateGraph(PlanState)
    builder.add_node("collect_data", collect_data)
    builder.add_node("generate_plan", generate_plan)
    builder.add_node("regenerate_section", regenerate_section)
    builder.add_edge(START, "collect_data")
    builder.add_edge("collect_data", "generate_plan")
    builder.add_edge("generate_plan", END)
    # regenerate_section 从外部直接调用，不走图
    graph = builder.compile()
    return graph