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

# ===============================================
# 编制依据目录抽取
#
# @author :ZengChao
# @date   :2025-12-30 15:13
# @logo   :(=^･w･^=)
# ===============================================

from __future__ import annotations

import json  # ✅ 最小修改：新增
import re
import asyncio
from typing import List

from pydantic import BaseModel, Field, ValidationError  # ✅ 最小修改：新增 ValidationError
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from foundation.ai.agent.generate.model_generate import generate_model_client
from foundation.observability.logger.loggering import review_logger as logger


# --------- 1) 结构定义 ---------
class BasisItem(BaseModel):
    """编制依据条目"""
    title: str = Field(..., description="法规/标准名称，不含书名号《》")
    suffix: str = Field("", description="《》后面的版本/日期/修订说明，可为空")
    raw: str = Field(..., description="该条目原文（尽量原样保留）")


class BasisItems(BaseModel):
    """编制依据条目列表"""
    items: List[BasisItem]


# --------- 2) Prompt（强约束：只抽《》条目，输出可解析结构）---------
SYSTEM = """
你是信息抽取助手。请从中文文本中抽取"编制依据/法律法规/规范标准"等条目。
规则：
1) 只抽取包含书名号《 》的条目。
2) 每条条目包括：title（《》内名称，去掉书名号）、suffix（《》后面的版本/日期/修订说明，可为空）、raw（该条目原文）。
3) 忽略标题行、段落说明、无《》的行。
4) 输出必须严格符合格式要求，不要输出任何额外文字。
"""
HUMAN ="""
文本如下：
{input_text}
【输出格式要求】
{format_instructions}
/no_think
"""

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


# --------- 3) 兜底：模型失败时，用简单规则捞《》行 ---------
def fallback_regex(text: str) -> BasisItems:
    """
    兜底方案：使用正则表达式提取编制依据
    改进：处理跨行的编制依据（如编号被换行分割的情况）
    优化：限制匹配长度，避免提取过多非编制依据内容
    """
    items: List[BasisItem] = []
    
    # 标准化换行符：将换行符替换为空格
    text = text.replace("\r\n", " ").replace("\r", " ").replace("\n", " ")
    
    # 使用正则表达式匹配所有编制依据条目
    # 优化后的模式：
    # 1. 限制书名长度1-60字符（避免匹配过长内容）
    # 2. 后缀限制：优先匹配括号内的编号，其次匹配最多30个字符的后缀
    # 3. 使用正向前瞻确保正确截断
    
    # 模式1：匹配《名称》（编号）格式
    pattern1 = r'《([^《》]{1,60})》[\s]*[（(]([^)）]{1,30})[）)]'
    
    # 模式2：匹配《名称》编号 格式（编号无括号）
    pattern2 = r'《([^《》]{1,60})》[\s]*([A-Za-z]{0,6}[/-]?[0-9]{2,6}(?:-[0-9]{4})?)'
    
    # 模式3：匹配只有《名称》的格式（在列表结尾或分号前）
    pattern3 = r'《([^《》]{1,60})》(?=[\s]*[；;]|\s*$|\s+[（(]\d)'
    
    # 合并模式，按优先级匹配
    matched_positions = set()
    
    for pattern in [pattern1, pattern2, pattern3]:
        for match in re.finditer(pattern, text):
            # 检查是否已匹配过此位置
            start_pos = match.start()
            if start_pos in matched_positions:
                continue
            
            title = match.group(1).strip()
            
            # 提取后缀（模式3可能没有group2）
            if len(match.groups()) >= 2:
                suffix = match.group(2).strip() if match.group(2) else ""
            else:
                suffix = ""
            
            # 清理 suffix
            next_book = suffix.find("《")
            if next_book != -1:
                suffix = suffix[:next_book].strip()
            suffix = suffix.rstrip("；;").strip()
            
            # 提取原始文本（限制长度）
            raw = match.group(0).strip()
            # 如果 raw 太长，尝试截断到合理长度
            if len(raw) > 100:
                # 尝试在书名和编号后截断
                book_end = raw.find("》")
                if book_end != -1:
                    # 包含书名号和后面的编号（如果有）
                    suffix_end = book_end + 1
                    bracket_start = raw.find("（", book_end)
                    bracket_end = raw.find("）", bracket_start) if bracket_start != -1 else -1
                    
                    if bracket_start != -1 and bracket_end != -1:
                        suffix_end = bracket_end + 1
                    
                    raw = raw[:suffix_end].strip()
            
            if title and len(title) <= 60:
                items.append(BasisItem(
                    title=title,
                    suffix=suffix,
                    raw=raw
                ))
                # 记录已匹配的位置
                matched_positions.add(start_pos)

    logger.info(f"[编制依据提取] 兜底方案提取到 {len(items)} 条")
    return BasisItems(items=items)


# --------- 3.5) 根据chapter_code过滤，避免LLM误抽 ---------
def _filter_by_chapter_code(items: List[BasisItem], chapter_code: str) -> List[BasisItem]:
    """
    根据章节代码过滤编制依据条目

    规则：
    - 如果 chapter_code == "basis"：编制依据章节，不过滤，返回所有条目
    - 如果 chapter_code != "basis"：非编制依据章节，过滤掉raw中没有书名号的条目

    Args:
        items: 待过滤的条目列表
        chapter_code: 章节代码（"basis" 表示编制依据章节）

    Returns:
        过滤后的条目列表
    """
    # 编制依据章节，不过滤
    if chapter_code == "basis":
        logger.debug(f"[编制依据提取] 当前为编制依据章节(basis)，不过滤，共 {len(items)} 条")
        return items

    # 非编制依据章节，过滤掉没有书名号的条目（避免LLM误抽）
    before_count = len(items)
    filtered_items = [
        item for item in items
        if "《" in item.raw and "》" in item.raw
    ]
    after_count = len(filtered_items)

    # 如果有过滤，记录日志
    if before_count != after_count:
        logger.info(
            f"[编制依据提取] 非编制依据章节(chapter_code={chapter_code})，"
            f"过滤掉无书名号条目: {before_count} -> {after_count} "
            f"(过滤 {before_count - after_count} 条)"
        )

    return filtered_items


# --------- 4) 最小修改：只做一件事，提取第一个 JSON ---------
def extract_first_json(text: str) -> dict:
    """
    从任意模型输出中提取第一个完整 JSON 对象 { ... }
    不关心 think、不关心多余文本
    """
    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 花括号未闭合")


# --------- 5) 主函数：使用 LangChain 抽取（最小改动版） ---------
async def extract_basis_with_langchain_qwen(progress_manager,callback_task_id:str,text: str,chapter_code: str) -> BasisItems:
    """
    使用 LangChain + LLM 提取编制依据信息（流式输出版本）
    """
    # 标准化文本（无论是否走LLM，都先做）
    text = text.replace("\r\n", "\n").replace("\r", "\n")

    try:
        # 创建解析器（✅ 保留：仅用于 format_instructions）
        parser = PydanticOutputParser(pydantic_object=BasisItems)

        # 构建消息
        messages = prompt.format_messages(
            input_text=text,
            format_instructions=parser.get_format_instructions()
        )

        logger.info(f"[编制依据提取] 开始使用 LLM 提取，文本长度: {len(text)}")

        # 推送开始消息（使用独立命名空间，避免与主流程进度冲突）
        if progress_manager and callback_task_id:
            try:
                await progress_manager.update_stage_progress(
                    callback_task_id=callback_task_id,
                    stage_name="编制依据提取-子任务",  # 独立命名空间
                    status="processing",
                    message=f"开始编制依据提取",
                    overall_task_status="processing",
                    event_type="processing"
                    # 不设置 current，避免覆盖主流程进度
                )
            except Exception as e:
                logger.error(f"SSE推送开始消息失败: {e}")

        # 流式调用模型，每检测到5个}符号推送一次进度
        raw_out = ""
        brace_count = 0

        # 使用通用模型底座的流式调用（通过 function_name 从 model_setting.yaml 加载配置）
        for chunk in generate_model_client.get_model_generate_stream(
            trace_id=callback_task_id or "directory_extract",
            messages=messages,
            function_name="directory_extraction"
        ):
            raw_out += chunk
            
            # 统计}符号出现次数
            for char in chunk:
                if char == "}":
                    brace_count += 1
                    
                    # 每5个}推送一次进度（使用独立命名空间）
                    if brace_count % 5 == 0:
                        if progress_manager and callback_task_id:
                            try:
                                await progress_manager.update_stage_progress(
                                    callback_task_id=callback_task_id,
                                    stage_name="编制依据提取-子任务",  # 独立命名空间
                                    status="processing",
                                    message=f"编制依据提取中... (已处理 {brace_count} 个结构)",
                                    overall_task_status="processing",
                                    event_type="processing"
                                    # 不设置 current，避免覆盖主流程进度
                                )
                            except Exception as e:
                                logger.error(f"SSE推送进度失败: {e}")
        
        logger.info(f"编制依据提取调试：{raw_out}")

        # ✅ 最小修改：不再 parser.parse / 不再 _parse_with_retry
        data = extract_first_json(raw_out)
        parsed = BasisItems.model_validate(data)

        # 清洗：title 去书名号，补 raw
        cleaned: List[BasisItem] = []
        for it in parsed.items:
            title = (it.title or "").replace("《", "").replace("》", "").strip()
            suffix = (it.suffix or "").strip()
            raw = (it.raw or "").strip()

            if not title:
                continue
            if not raw:
                raw = f"《{title}》{suffix}".strip()

            cleaned.append(BasisItem(title=title, suffix=suffix, raw=raw))

        # ✅ 新增：根据chapter_code过滤，避免LLM在非编制依据章节误抽
        cleaned = _filter_by_chapter_code(cleaned, chapter_code)

        if cleaned:
            logger.info(f"[编制依据提取] LLM 提取成功，共 {len(cleaned)} 条")
            return BasisItems(items=cleaned)

        logger.warning("[编制依据提取] LLM 未提取到内容，使用兜底方案")
        # ✅ 修改：兜底方案也要经过过滤
        fallback_result = fallback_regex(text)
        filtered_items = _filter_by_chapter_code(fallback_result.items, chapter_code)
        return BasisItems(items=filtered_items)

    except (json.JSONDecodeError, ValidationError, ValueError) as e:
        fallback_result = fallback_regex(text)
        filtered_items = _filter_by_chapter_code(fallback_result.items, chapter_code)
        return BasisItems(items=filtered_items)
    except Exception as e:
        fallback_result = fallback_regex(text)
        filtered_items = _filter_by_chapter_code(fallback_result.items, chapter_code)
        return BasisItems(items=filtered_items)







if __name__ == "__main__":
    demo = """一）编制依据
（1）相关法律法规
1）《中华人民共和国安全生产法》2021年修订版
2）《中华人民共和国环境保护法》2021年修订版
3）《建设工程安全生产管理条例》2023年最新修正
4）《中华人民共和国道路交通安全法》2021年4月29日修订
5）《中华人民共和国水土保持法》2010年12月25日修订
6）《公路水运危险性较大工程专项施工方案编制审查规程》（JT/T 1495-2024）
7）《公路水运工程临时用电技术规程》（JT/T 1499—2024）
8）《建设工程质量管理条例》2019年4月23日修订
9）《建设工程安全生产管理条例》2023年最新修正
10）《四川省安全生产条例》2023年5月25日修订
11）《危险性较大的分部分项工程安全管理规定》住建部令第37号
12）《建筑机械使用安全技术规程》JGJ 33—2012
13）《起重机—手势信号》GB/T 5082—2019
14）《架桥机通用技术条件》GB/T 26470—2011
15）《施工现场机械设备检查技术规范》JGJ 160—2016
16）《粗直径钢丝绳》GB/T 20067—2017
17）《建筑施工起重吊装工程安全技术规范》JGJ 276—2012
18）《架桥机安全规程》GB 26469—2011
19）《起重机械安全规程》GB 6067.1—2010
20）《电气装置安装工程起重机电气装置施工及验收规范》GB 50256—2014
21）《起重设备安装工程施工及验收规范》GB 50278—2010
22）《公路工程施工安全技术规范》JTG F90—2015
23）《建筑施工高处作业安全技术规范》JGJ 80—2016
24）《电力高处作业防坠器》DL/T 1147-2009
（2）项目文件
1）《S81线会理至禄劝（四川境）高速公路两阶段施工图设计》；
2）《S81线会理至禄劝（四川境）高速公路ZCB1-3标段实施性施工组织设计》；
3）《S81线会理至禄劝（四川境）高速公路ZCB1-3标段T梁预制、运输及安装专项施工方案》。
（3）我公司现有可投入工程的施工技术力量和机械设备；
（4）近年来，我公司参加类似工程的经验；
（5）本合同段工程现场踏勘、调查所获得的现场情况、自然环境、人文环境、市场环境等参考资料；
"""

    async def _demo_run():
        result = await extract_basis_with_langchain_qwen(None, None, demo)
        print(f"\n提取到 {len(result.items)} 条编制依据：")
        for idx, item in enumerate(result.items, 1):
            print(f"\n{idx}. {item.model_dump()}")

    asyncio.run(_demo_run())
