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,  # todo
        """
你是一个图像分析助手。现在给你一张道路图片，你需要观察远离你的第二根车道上，画面中心的车辆，请你从分析该车辆。

要求：
1. have_animal 字段中填写 true
2. livestock_type 字段中填写 “货物种类”，例如 牛、羊、猪、鸡、鸭、鹅、钢管、土渣等任何你观察到的
3. remark 字段 需要你简短的描述该车辆车身情况，例如什么颜色的车身，你需要完整的尽你所能形容一下。
3. 不允许输出多余文字，直接返回 JSON。

JSON 示例格式：
   {
       "have_animal": true,
       "livestock_type": "“测试数据：” + 5位随机数",   // 例如 测试数据34223
       "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