import json
import re
from typing import TypedDict

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

from config.llm import llmVision


# -------- 定义状态 --------
class State(TypedDict):
    image_url: str  # 图像
    content: str  # 最终内容


def send_analyze(state: State, prompt_text: str):
    messages = [
        HumanMessage(
            content=[
                {"type": "text", "text": prompt_text},
                {"type": "image_url", "image_url": {"url": state["image_url"]}},
            ]
        )
    ]
    return llmVision.invoke(messages).content


def analysis(state: State):
    state["content"] = send_analyze(
        state,
        """
提示词示例
你是一个图像分析助手。现在给你一张车的侧身照片，请你从图中分析车上运输的牲畜种类。

要求：
1. 牲畜种类可能是：牛、羊、猪、鸡、鸭、鹅。
2. 如果图中无法判断牲畜类型，请在备注字段 remark 中写明“无法识别”或你观察到的情况。
3. 不允许输出多余文字，直接返回 JSON。

JSON 示例格式：
   {
       "livestock_type": "<牲畜种类>",   // 如果能识别就填牛/羊/猪/鸡/鸭/鹅
       "remark": "<备注>"                  // 如果无法识别，写明原因；否则可留空
   }
请确保输出的 JSON 可以被严格解析。
""",
    )
    return state


# ------------------------------------------------------------------------ 构建有向图 --------
workflow = StateGraph(State)
# 必须先从 START 指向 analysis
workflow.add_node("analysis", analysis)
workflow.set_entry_point("analysis")
workflow.add_edge("analysis", END)
graph = workflow.compile()


# 执行函数


async def get_vehicle_response(image_url: str):
    final_state = graph.invoke(
        {
            "image_url": image_url,
        }
    )
    # 去掉 ```json 和 ``` 包裹
    content_str = re.sub(r"^```json\s*|\s*```$", "", final_state["content"].strip())
    # 把 JSON 字符串转为字典
    try:
        content_dict = json.loads(content_str)
    except json.JSONDecodeError:
        print("JSON解析失败")
        content_dict = {}

    return content_dict