import json
from typing import TypedDict, Optional
from langgraph.graph import StateGraph, START, END
from langchain_core.messages import HumanMessage

from ..llm.client import get_llm
from ..prompts.questionnaire import PARENT_TEMPLATE, CHILD_TEMPLATE
from ..schemas.questionnaire import Questionnaire, GenerateRequest


class GraphState(TypedDict):
    request: GenerateRequest
    _prompt: Optional[str]
    raw_response: str
    questionnaire: Optional[dict]
    error: Optional[str]


def _build_context(req: GenerateRequest) -> str:
    """格式化家庭背景信息字符串"""
    if not req.family_context:
        return "信息暂缺，请基于基本关系作答"
    parts = []
    members = req.family_context.get("members", [])
    if members:
        parts.append(f"家庭共有 {len(members)} 位成员")
    cohabitation = req.family_context.get("cohabitation")
    if cohabitation:
        parts.append(f"是否同住：{'是' if cohabitation else '否'}")
    conflicts = req.family_context.get("recent_conflicts", [])
    if conflicts:
        parts.append(f"近期有 {len(conflicts)} 次主要冲突")
    return "；".join(parts) if parts else "信息暂缺，请基于基本关系作答"


def build_prompt(state: GraphState) -> GraphState:
    req = state["request"]
    # 优先使用 Java 后端传入的 prompt 模板（可配置），缺省回退内置模板
    template = req.prompt_template or (
        PARENT_TEMPLATE if req.relationship_type == "parent" else CHILD_TEMPLATE
    )

    # 构建双方画像数据
    resp_prof = req.respondent_profile or {}
    mem_prof = req.member_profile or {}

    def safe_get(d, key, default=""):
        return str(d.get(key)) if d and key in d and d.get(key) else default

    respondent_name = safe_get(resp_prof, "nickname", "家长")
    member_name = safe_get(mem_prof, "nickname", req.member_name)
    respondent_age = safe_get(resp_prof, "age", "未知")
    member_age = safe_get(mem_prof, "age", "未知")
    respondent_gender = safe_get(resp_prof, "gender", "未知")
    member_gender = safe_get(mem_prof, "gender", "未知")
    respondent_zodiac = safe_get(resp_prof, "zodiac", "未知")
    member_zodiac = safe_get(mem_prof, "zodiac", "未知")
    relationship_desc = "家长与孩子" if req.relationship_type == "parent" else "孩子与家长"
    family_ctx = _build_context(req)

    prompt = template.format(
        member_name=member_name,
        respondent_name=respondent_name,
        relationship_desc=relationship_desc,
        respondent_age=respondent_age,
        respondent_gender=respondent_gender,
        respondent_zodiac=respondent_zodiac,
        member_age=member_age,
        member_gender=member_gender,
        member_zodiac=member_zodiac,
        family_context=family_ctx,
    )
    return {**state, "_prompt": prompt}


def call_llm(state: GraphState) -> GraphState:
    llm = get_llm()
    messages = [HumanMessage(content=state["_prompt"])]
    response = llm.invoke(messages)
    return {**state, "raw_response": response.content}


def validate(state: GraphState) -> GraphState:
    raw = state["raw_response"]
    # 提取 JSON 块
    text = raw.strip()
    if "```json" in text:
        text = text.split("```json")[1].split("```")[0].strip()
    elif "```" in text:
        text = text.split("```")[1].split("```")[0].strip()

    try:
        data = json.loads(text)
        q = Questionnaire(**data)
        if not q.validate_structure():
            return {**state, "error": "问卷结构校验失败：至少需要1题且每题有options或scale"}
        return {**state, "questionnaire": data, "error": None}
    except Exception as e:
        return {**state, "error": f"JSON 解析失败: {str(e)}"}


def build_questionnaire_graph():
    graph = StateGraph(GraphState)
    graph.add_node("build_prompt", build_prompt)
    graph.add_node("call_llm", call_llm)
    graph.add_node("validate", validate)
    graph.add_edge(START, "build_prompt")
    graph.add_edge("build_prompt", "call_llm")
    graph.add_edge("call_llm", "validate")
    graph.add_edge("validate", END)
    return graph.compile()


_questionnaire_graph = None


def get_questionnaire_graph():
    global _questionnaire_graph
    if _questionnaire_graph is None:
        _questionnaire_graph = build_questionnaire_graph()
    return _questionnaire_graph