from typing import TypedDict, Optional
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.report_tools import get_report_detail, get_survey_data, get_dimension_scores
from app.config import settings
from app.prompt_service import get_prompt
import logging

logger = logging.getLogger(__name__)


class AnalysisState(TypedDict):
    report_id: int
    user_id: int
    focus: Optional[str]          # nutrition / gut / chronic / overall
    report_data: Optional[dict]
    survey_data: Optional[dict]
    dimension_scores: Optional[dict]
    analysis: Optional[str]
    recommendations: list[str]


ANALYSIS_SYSTEM_PROMPT = """你是一个儿童健康报告解读专家。根据健康报告数据和调研问卷, 提供专业、易懂的分析。

分析原则:
1. 用通俗语言解释各项指标含义
2. 关注异常指标, 给出改善建议
3. 结合问卷数据提供个性化分析
4. 五维能量(身/心/智/行/富)角度解读整体状况
5. 输出格式: 总体评估 → 分项分析 → 改善建议
"""


def create_analysis_graph():
    """创建报告解读 StateGraph"""
    llm = ChatOpenAI(
        model=settings.llm_model,
        api_key=settings.llm_api_key,
        base_url=settings.llm_base_url,
        temperature=settings.llm_temperature,
    )
    llm_with_tools = llm.bind_tools([
        get_report_detail,
        get_survey_data,
        get_dimension_scores,
    ])

    builder = StateGraph(AnalysisState)

    async def gather_data(state: AnalysisState) -> dict:
        """收集报告 + 问卷 + 维度数据"""
        result = {}
        if state.get("report_id"):
            try:
                detail = await get_report_detail.ainvoke({"report_id": state["report_id"]})
                result["report_data"] = detail
            except Exception as e:
                logger.warning("获取报告详情失败: %s", e)
            try:
                survey = await get_survey_data.ainvoke({"report_id": state["report_id"]})
                result["survey_data"] = survey
            except Exception as e:
                logger.warning("获取问卷数据失败: %s", e)
        return result

    async def analyze(state: AnalysisState) -> dict:
        """LLM 分析"""
        context_parts = []
        if state.get("report_data"):
            context_parts.append(f"报告数据: {state['report_data']}")
        if state.get("survey_data"):
            context_parts.append(f"问卷数据: {state['survey_data']}")
        if state.get("focus"):
            context_parts.append(f"重点关注: {state['focus']}")

        context_text = "\n".join(context_parts) if context_parts else "暂无数据"

        messages = [
            SystemMessage(content=await get_prompt("analysis") or ANALYSIS_SYSTEM_PROMPT),
            SystemMessage(content=f"待分析数据:\n{context_text}"),
            HumanMessage(content="请分析以上健康数据, 给出评估和建议。"),
        ]

        response = await llm_with_tools.ainvoke(messages)
        return {"analysis": response.content}

    builder.add_node("gather_data", gather_data)
    builder.add_node("analyze", analyze)

    builder.add_edge(START, "gather_data")
    builder.add_edge("gather_data", "analyze")
    builder.add_edge("analyze", END)

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