#!/usr/bin/env python
# -*- coding: utf-8 -*-

import json
from typing import List, Literal

from pydantic import BaseModel, Field, ValidationError
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser
from foundation.ai.agent.generate.model_generate import generate_model_client


# ===== 1) 定义结构 =====
RiskLevel = Literal["无风险", "高风险"]


class TimelinessResult(BaseModel):
    issue_point: str = Field(..., description="时效性判定类型")
    location: str = Field(..., description="审查的规范原文，与输入完全一致")
    suggestion: str = Field(..., description="建议")
    reason: str = Field(..., description="原因")
    risk_level: RiskLevel = Field(..., description='风险水平，只能是 "无风险" / "高风险"')


class TimelinessResults(BaseModel):
    items: List[TimelinessResult]


# ===== 2) SYSTEM Prompt =====
SYSTEM = """
你是【规范时效性判定助手】。

【任务】
根据规范匹配结果，判定每个审查规范的时效性问题类型。

【重要说明（必须严格遵守）】
- 如果有具体的文件将回答中的"XXX"替换为具体文件，如果没有具体文件回答内容中不要出现"XXX"字样

【输出要求】
- 为每个审查规范输出一个判定结果
- 确保输出数量与输入一致
- review_item 必须与输入完全一致
- 严格按照判定规则的优先级顺序进行判断
"""

HUMAN = """
请根据以下规范匹配结果，判定每个审查规范的时效性问题类型：

【判定规则（按优先级从高到低）】

1. **无参考规范**（无风险）
   - 条件：has_related_file = false
   - 原因：在参考规范库中完全找不到相关文件
   - 建议：当前引用未在参考规范库中发现，建议人工核实其有效性

2. **规范编号错误**（高风险）
   - 条件：has_related_file = true 且 has_exact_match = false
   - 原因：与参考文件XXX编号不一致
   - 建议：建议核实并更正为参考库中的正确编号XXX

3. **规范编号正确**（无风险）
   - 条件：has_exact_match = true 且 exact_match_info 中状态为"现行"
   - 原因：与参考文件XXX名称编号一致，且文件状态为现行
   - 建议：引用规范为现行有效版本，无需调整

4. **引用已废止的规范**（高风险）
   - 条件：has_exact_match = true 且 exact_match_info 中状态为"废止" 且 same_name_current 为空
   - 原因：参考文件显示XXX已废止，且无明确替代版本
   - 建议：建议删除该引用或咨询最新替代规范

5. **引用已被替代的规范**（高风险）
   - 条件：has_exact_match = true 且 exact_match_info 中状态为"废止" 且 same_name_current 不为空
   - 原因：参考文件显示XXX已废止，但存在XXX现行版本
   - 建议：建议更新为现行替代标准

【规范匹配结果】
{match_results}

【输出格式要求】
{format_instructions}
/no_think
"""

# ===== 3) Output Parser =====
parser = PydanticOutputParser(pydantic_object=TimelinessResults)

# ===== 4) Prompt =====
prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM),
    ("human", HUMAN)
])

# ===== 5) LLM Client (通用模型底座) =====
model_client = generate_model_client


# ===== 6) 提取第一个 JSON =====
def extract_first_json(text: str) -> dict:
    """从任意模型输出中提取第一个完整 JSON 对象 { ... }"""
    start = text.find("{")
    if start == -1:
        raise ValueError("未找到 JSON 起始 '{'")

    depth = 0
    for i in range(start, len(text)):
        ch = text[i]
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return json.loads(text[start:i + 1])

    raise ValueError("JSON 花括号未闭合")


# ===== 7) 核心方法 =====
async def determine_timeliness_issue(match_results: str) -> str:
    """
    根据规范匹配结果判定时效性问题类型
    
    Args:
        match_results: match_reference_files 的返回结果（JSON字符串）
        
    Returns:
        时效性判定结果的JSON字符串
    """
    format_instructions = parser.get_format_instructions()

    # 构建消息
    messages = prompt.format_messages(
        match_results=match_results,
        format_instructions=format_instructions
    )

    last_err = None

    for _ in range(2):
        try:
            raw = await model_client.get_model_generate_invoke(
                trace_id="timeliness_determine",
                messages=messages,
                model_name="qwen3_30b"
            )
            print(f"[时效性判定] 模型输出: {raw}...")
            data = extract_first_json(raw)
            findings = TimelinessResults.model_validate(data)
            result = [x.model_dump() for x in findings.items]
            return json.dumps(result, ensure_ascii=False, indent=2)
        except (Exception, ValidationError, json.JSONDecodeError) as e:
            last_err = e

    raise RuntimeError(f"时效性判定失败：{last_err}") from last_err


# ===== 8) 示例 =====
if __name__ == "__main__":
    import asyncio

    # 模拟 match_reference_files 的返回结果
    match_results = json.dumps([
        {
            "review_item": "《混凝土结构设计规范》（GB 50010-2010）",
            "has_related_file": True,
            "has_exact_match": True,
            "exact_match_info": "《混凝土结构设计规范》（GB 50010-2010）状态为现行",
            "same_name_current": "《混凝土结构设计规范》（GB 50010-2010）状态为现行"
        },
        {
            "review_item": "《建筑施工组织设计规范》（GB/T 50502-2015）",
            "has_related_file": True,
            "has_exact_match": True,
            "exact_match_info": "《建筑施工组织设计规范》（GB/T 50502-2015）状态为废止",
            "same_name_current": "《建筑施工组织设计规范》（GB/T 50502-2020）状态为现行"
        },
        {
            "review_item": "《建筑抗震设计规范》（GB 50011-2001）",
            "has_related_file": True,
            "has_exact_match": True,
            "exact_match_info": "《建筑抗震设计规范》（GB 50011-2001）状态为废止",
            "same_name_current": ""
        },
        {
            "review_item": "《城市道路工程设计规范》（CJJ 37-2012）",
            "has_related_file": False,
            "has_exact_match": False,
            "exact_match_info": "",
            "same_name_current": ""
        }
    ], ensure_ascii=False)

    result = asyncio.run(determine_timeliness_issue(match_results))
    print("\n时效性判定结果：")
    print(result)
