"""
舌诊 LangGraph — 视觉 LLM 舌象分析

流程：
  START → load_image → analyze_tongue → END

输出：整体评估 + 7 类结构化舌诊指标
"""
import base64
from typing import TypedDict, Optional

from langgraph.graph import StateGraph, START, END
from langchain_core.messages import HumanMessage

from ..llm.client import get_vision_llm

TONGUE_INDICATOR_CODES = [
    "tongue_color", "coating_color", "coating_texture",
    "fissure", "teeth_mark", "sublingual_vein", "constitution",
]

# LLM 可能自创的 code 别名 → 规范 code（容错归一化）
INDICATOR_ALIASES = {
    "tongue_shape": "tongue_color",
    "tongue_coating_color": "coating_color",
    "tongue_coating_texture": "coating_texture",
    "coating": "coating_color",
    "sublingual_veins": "sublingual_vein",
    "tooth_mark": "teeth_mark",
    "teeth_marks": "teeth_mark",
    "body_constitution": "constitution",
}

SYSTEM_PROMPT = (
    "你是资深中医舌诊专家。根据用户上传的舌象图片，输出结构化 JSON，"
    "不要输出任何 JSON 之外的文字。"
    "indicators 数组的 code 字段必须严格从以下 7 个值中选择，禁止自创或改写："
    "tongue_color（舌色）、coating_color（苔色）、coating_texture（苔质）、"
    "fissure（裂纹）、teeth_mark（齿痕）、sublingual_vein（舌下络脉）、constitution（体质）。"
    "JSON 格式："
    '{"overall_assessment": "整体舌象评估结论", "indicators": ['
    '{"code": "tongue_color", "value": "淡红"}, '
    '{"code": "coating_color", "value": "薄白"}]}'
    "指标值用简短中文描述，如舌色「淡红」、苔色「薄白」、裂纹「无」、齿痕「轻」。"
)


class GraphState(TypedDict):
    request: dict
    image_base64: Optional[str]
    raw_response: str
    overall_assessment: str
    indicators: list
    error: Optional[str]


def load_image(state: GraphState) -> GraphState:
    req = state["request"]
    if req.get("image_base64"):
        return {**state, "image_base64": req["image_base64"]}
    # 支持 data URL 格式 (data:image/jpeg;base64,xxx) — TongueDiagnosisAgent 调用场景
    url = req.get("image_url", "")
    if url and url.startswith("data:image"):
        payload = url.split(",", 1)[1]
        return {**state, "image_base64": payload}
    if url:
        return {**state, "error": "舌诊暂不支持外部 URL，请传 image_base64 或 data URL"}
    return {**state, "error": "image_url 或 image_base64 至少提供一项"}


def analyze_tongue(state: GraphState) -> GraphState:
    if state.get("error"):
        return state
    llm = get_vision_llm()
    # 优先使用 Java 后端传入的 prompt 模板（可配置），缺省回退内置 SYSTEM_PROMPT
    prompt_text = state["request"].get("prompt_template") or SYSTEM_PROMPT
    content = [
        {"type": "text", "text": prompt_text},
        {
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{state['image_base64']}"},
        },
    ]
    response = llm.invoke([HumanMessage(content=content)])
    return {**state, "raw_response": response.content}


def parse_result(state: GraphState) -> GraphState:
    raw = state.get("raw_response", "").strip()
    text = raw
    if "```json" in text:
        text = text.split("```json")[1].split("```")[0].strip()
    elif "```" in text:
        text = text.split("```")[1].split("```")[0].strip()

    import json
    try:
        data = json.loads(text)
        assessment = data.get("overall_assessment", "")
        indicators = []
        for it in data.get("indicators", []):
            code = it.get("code")
            if code in INDICATOR_ALIASES:
                code = INDICATOR_ALIASES[code]
            if code not in TONGUE_INDICATOR_CODES:
                continue
            indicators.append({"code": code, "value": it.get("value")})
        if not assessment or not indicators:
            return {**state, "error": "舌诊结果缺少评估或指标"}
        return {
            **state,
            "overall_assessment": assessment,
            "indicators": indicators,
            "error": None,
        }
    except Exception as e:
        return {**state, "error": f"舌诊 JSON 解析失败: {e}"}


def build_tongue_graph():
    graph = StateGraph(GraphState)
    graph.add_node("load_image", load_image)
    graph.add_node("analyze_tongue", analyze_tongue)
    graph.add_node("parse_result", parse_result)
    graph.add_edge(START, "load_image")
    graph.add_edge("load_image", "analyze_tongue")
    graph.add_edge("analyze_tongue", "parse_result")
    graph.add_edge("parse_result", END)
    return graph.compile()


_tongue_graph = None


def get_tongue_graph():
    global _tongue_graph
    if _tongue_graph is None:
        _tongue_graph = build_tongue_graph()
    return _tongue_graph