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  # 证件照
    type: int  # 证件类别
    content: str  # 最终内容


def analyze_card(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 decide_source(state: State, max_retry=3):
    choice = analyze_card(
        state,
        """
        请根据图片内容判断其中的证件类型，并输出对应代号。
        代号说明：
        -1: 并非证件
        0：身份证
        1：银行卡
        
        只输出代号数字，不要输出任何文字、标点或解释。
        你的输出是：
    """,
    )
    print("图片类型是", choice)
    try:
        choice = int(choice)
        state["type"] = choice
        if choice == -1:
            state["content"] = '{"result": "暂不支持此证件"}'
    except ValueError:
        state["type"] = -1
        state["content"] = '{"result": "暂不支持此证件"}'
    return state


def idcard(state: State):
    state["content"] = analyze_card(
        state,
        """
    你是一个 OCR 信息提取专家，请从图片中识别出身份证的全部可见信息，并输出严格的 JSON 格式。
    要求：
    1. 无论是正面还是反面，都要识别所有字段。
    2. 如果某个字段无法辨认，请将值设为 null。
    3. 只输出 JSON，不要输出任何解释或多余内容。
    4. 严格保持 JSON 格式，键名统一为英文。
    5. "side" 字段为身份证方向，填写：人像面 或 国徽面。
    
    JSON 示例格式：
    {
      "side": "国徽面",
      "name": "持证人姓名",
      "gender": "性别",
      "ethnicity": "民族",
      "id_number": "身份证号",
      "birth_date": "出生日期",
      "address": "住址",
      "issuing_authority": "签发机关",
      "valid_period_start": "有效期开始日期",
      "valid_period_end": "有效期结束日期",
      "notes": "其他可见文字或备注"
    }
    
    请确保输出的 JSON 可以被严格解析。
""",
    )
    return state


def bankcard(state: State):
    state["content"] = analyze_card(
        state,
        """
    你是一个 OCR 信息提取专家，请从图片中识别出银行卡上的全部可见信息，并输出严格的 JSON 格式。
    要求：
    1. 识别卡号、发卡行、持卡人、有效期等关键信息。
    2. 如果某个字段无法辨认，请将值设为 null。
    3. 只输出 JSON，不要输出任何解释或附加内容。
    4. 严格保持 JSON 格式，键名统一为英文。
    
    JSON 示例格式：
    {
      "bank_name": "发卡行名称",
      "card_number": "卡号",
      "card_holder": "持卡人姓名",
      "expiry_date": "有效期",
      "card_type": "卡种类型（如借记卡/信用卡/Visa/MasterCard）",
      "issuer_country": "发卡国家",
      "notes": "其他可见文字或符号"
    }
    
    请确保输出的 JSON 可以被严格解析。
""",
    )
    return state


# ------------------------------------------------------------------------ 构建有向图 --------
workflow = StateGraph(State)
workflow.add_node("decide", decide_source)
workflow.add_node("idcard", idcard)
workflow.add_node("bankcard", bankcard)
workflow.set_entry_point("decide")
# 条件边：根据 path 决定走向
workflow.add_conditional_edges(
    "decide",
    lambda state: state["type"],
    {
        -1: END,
        0: "idcard",
        1: "bankcard",
    },
)
workflow.add_edge("idcard", END)
workflow.add_edge("bankcard", END)
graph = workflow.compile()

import re


# 执行函数
def get_license_response(image_url: str):
    final_state = graph.invoke(
        {
            "image_url": image_url,
        }
    )
    # 去掉 ```json 和 ``` 包裹
    final_state["content"] = final_state["content"]
    final_state["content"] = re.sub(
        r"^```json\s*|\s*```$", "", final_state["content"].strip()
    )
    print("最终识别内容：", final_state["content"])
    return final_state