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

"""
@author: bidikeji
@time: 2025/9/29 11:17
"""
from lxml import etree
import re
from typing import List, Dict, Tuple
import time
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class OutlineExtractor:
    """
    大纲提取工具类，支持处理不规范HTML文档，能识别短大纲和编号被拆分的情况
    可多次调用，适用于批量处理多个HTML文件
    """

    def __init__(self, max_merge_short=2, max_merge_normal=5):
        """初始化提取器，预编译正则表达式"""
        # 配置参数
        self.max_merge_short = max_merge_short  # 短标题最大合并数量
        self.max_merge_normal = max_merge_normal  # 正常标题最大合并数量

        # 基础编号模式（可能被拆分的部分）
        self.number_patterns = {
            'cn_number': re.compile(r'^[一二三四五六七八九十壹贰叁肆伍陆柒捌玖拾]+、?$'),
            'cn_number_with_paren': re.compile(r'^（[一二三四五六七八九十壹贰叁肆伍陆柒捌玖拾]+）$'),
            'digit': re.compile(r'^\d+$'),
            'digit_with_paren': re.compile(r'^（\d+）$'),
        }

        # 符号模式（可能与编号分离的符号）
        self.symbol_patterns = {
            'colon': re.compile(r'^[:：]$'),
            'comma': re.compile(r'^[、,.]'),
            'dot': re.compile(r'^\.$'),
            'brackets ': re.compile(r'^[（）]$'),
        }

        # 完整编号规则（用于匹配合并后的完整标题）
        self.full_outline_patterns = [
            (re.compile(r'^[一二三四五六七八九十壹贰叁拾]+[、.].*$'), 1, 6),
            (re.compile(r'^（[一二三四五六七八九十壹贰叁拾]+）.*$'), 2, 6),
            (re.compile(r'^\d+[、.\s].*$'), 3, 5),
            (re.compile(r'^（\d+）.*$'), 4, 5),
            (re.compile(r'^\d+\.\d+.*$'), 4, 4),
        ]

        # 预编译完整模式
        self.compiled_full_patterns = [
            (re.compile(p.pattern), l, prio) for p, l, prio in self.full_outline_patterns
        ]

        # 排除规则使用的模式
        self.money_patterns = [
            re.compile(r'^[￥$]\d+[,.\d]*'),
            re.compile(r'^\d+[,.\d]*\s*元'),
            re.compile(r'^\d+[,.\d]*\s*（[万亿]元）'),
            re.compile(r'^\d+[,.\d]*\s*[万亿]'),
            re.compile(r'^0+\.'),
        ]

        self.date_patterns = [
            re.compile(r'^\d{4}[年./-]\d{1,2}[月./-]\d{1,2}[日]?'),
            re.compile(r'^\d{2}[/-]\d{2}[/-]\d{4}'),
            re.compile(r'^\d{4}年\d{1,2}月'),
            re.compile(r'^\d{2,4}[年月日]'),
            re.compile(r'^（\d{4}）')
        ]

        self.quantity_patterns = [
            re.compile(r'^\d+[个只台件套次]'),
            re.compile(r'^\d+[.,]\d+[米千克吨亩平立方mkgtMKGT]'),
            re.compile(r'^[\d.]+%')
        ]

        # 缓存
        self.level_cache = {}
        self.identifier_cache = {}

    def _get_node_identifier(self, elem) -> str:
        """生成节点的唯一标识字符串，带缓存"""
        if elem in self.identifier_cache:
            return self.identifier_cache[elem]

        if elem is None:
            result = "None"
            self.identifier_cache[elem] = result
            return result

        tag = elem.tag
        if tag.startswith('{'):
            tag = tag.split('}')[-1]

        parent_id = self._get_node_identifier(elem.getparent())
        siblings = list(elem.getparent()) if elem.getparent() is not None else []
        index = siblings.index(elem) if elem in siblings else -1

        result = f"{parent_id}->{tag}[{index}]"
        self.identifier_cache[elem] = result
        return result

    def _get_level(self, elem) -> int:
        """计算节点层级，带缓存"""
        if elem in self.level_cache:
            return self.level_cache[elem]

        level = 0
        current = elem
        while current is not None and current.tag != 'body':
            level += 1
            current = current.getparent()

        self.level_cache[elem] = level
        return level

    def _is_non_outline(self, text: str, is_final: bool = False) -> bool:
        """判断是否为非大纲内容（金额、日期、数量等）"""
        text = text.strip()
        # if len(text) < 2:  # 过滤过短文本
        #     return True

        # 检查金额模式
        for p in self.money_patterns:
            if p.match(text):
                return True

        # 检查日期模式
        for p in self.date_patterns:
            if p.match(text):
                return True

        # 检查数量模式
        for p in self.quantity_patterns:
            if p.match(text):
                return True

        if is_final:
            # 检查最终句子是否包含中文
            if re.search('[\u4e00-\u9fa5]{2,}', text) == None:
                # print('检查最终句子是否包含中文', text)
                return True

            # 检查数字是否3位数以上
            elif re.search('^（?\d{3,}|^(\d{1,}\.){3,}\d|^\d{1,}\.\d{3,}', text):
                # print('检查数字是否3位数以上', text)
                return True
        return False

    def _collect_text_nodes(self, tree) -> List[Dict]:
        """收集所有文本节点及其相关信息，优化查询范围"""
        text_nodes = []
        # 只查询可能包含标题的标签，减少节点数量
        # elements = tree.xpath(
        #     '//*[self::p or self::div or self::h1 or self::h2 or self::h3 or self::h4 or self::h5 or self::h6 or self::li]')
        elements = tree.xpath('//*[normalize-space(text()) != ""]')

        for elem in elements:
            parent = elem.getparent()
            parent_identifier = self._get_node_identifier(parent)

            # 获取元素在父节点中的位置及相邻元素
            parent_children = list(parent) if parent is not None else []
            elem_index = parent_children.index(elem) if elem in parent_children else -1
            prev_sibling_elem = parent_children[elem_index - 1] if elem_index > 0 else None
            next_sibling_elem = parent_children[elem_index + 1] if elem_index != -1 and elem_index + 1 < len(
                parent_children) else None

            # 提取文本节点
            for node in elem.xpath('./text()[normalize-space() != ""]'):
                text = re.sub(r'\s+', ' ', node.strip())
                text = text.replace('(', '（').replace(')', '）')
                if text:
                    # print('文本节点：', elem.tag,elem_index, self._get_level(elem) ,parent_identifier, text)
                    text_nodes.append({
                        'text': text,
                        'level': self._get_level(elem),
                        'position': len(text_nodes), #  + 1
                        'element': elem,
                        'parent_identifier': parent_identifier,
                        'prev_sibling_elem': prev_sibling_elem,
                        'next_sibling_elem': next_sibling_elem,
                        'elem_index': elem_index
                    })

        logger.info(f"收集到有效文本节点: {len(text_nodes)}个")
        return text_nodes

    def _identify_candidates(self, text_nodes: List[Dict]) -> List[Dict]:
        """识别潜在的标题片段，包括被拆分的情况"""
        candidate_fragments = []
        # for i, item in enumerate(text_nodes):
        i = 0
        while i < len(text_nodes):
            item = text_nodes[i]
            text = item['text']

            # 初步过滤非大纲内容（双重保险）
            if self._is_non_outline(text):
                # print('初步过滤非大纲内容text', text)
                if len(candidate_fragments)>0:
                    candidate_fragments.append({
                        **item,
                        'outline_level': None,
                        'is_start': False,
                        'pattern': None,
                        'original_position': (i, i)
                    })
                i += 1
                continue

            # 检查是否为编号部分（可能被拆分）
            is_number_part = False
            # print('检查是否为编号部分: ', text)
            for _, pattern in self.number_patterns.items():
                if pattern.match(text):
                    is_number_part = True
                    combined_text = text
                    current_pos = i
                    # print('编号部分：', text)

                    # 向前查找（如果当前是符号，前面可能有编号）
                    if i > 0:
                        prev_item = text_nodes[i - 1]
                        for _, s_pattern in self.symbol_patterns.items():
                            if s_pattern.match(prev_item['text']):
                                combined_text = prev_item['text'] + combined_text
                                current_pos = i - 1

                    # 向后查找（如果当前是编号，后面可能有符号和文本）
                    j = i + 1
                    max_lookahead = min(i + 3, len(text_nodes))  # 最多向后找2个节点
                    while j < max_lookahead:
                        next_item = text_nodes[j]
                        # 检查是否是相邻元素
                        if (item['next_sibling_elem'] == next_item['element'] or
                                item['parent_identifier'] == next_item['parent_identifier']):

                            # 检查是否是符号
                            is_symbol = False
                            for _, s_pattern in self.symbol_patterns.items():
                                if s_pattern.match(next_item['text']):
                                    combined_text += next_item['text']
                                    is_symbol = True
                                    break
                            if is_symbol:
                                i = j
                                j += 1
                                continue

                            # 合并文本后停止
                            combined_text += next_item['text']
                            i = j
                            j += 1
                            break
                        elif re.match('^[一二三四五六七八九十]+$', item['text']) and re.match('、', next_item['text']):
                            combined_text += next_item['text']
                            i = j
                            j += 1
                            break
                        elif re.match('^[一二三四五六七八九十]+、$', item['text']) and re.match('[\u4e00-\u9fa5]', next_item['text']): # 优化 573063077 被多重标签包裹，不是直接相邻标签
                            combined_text += next_item['text']
                            i = j
                            j += 1
                            break
                        else:
                            break
                        j += 1

                    # 检查合并后的文本是否符合大纲模式
                    for pattern, level, priority in self.compiled_full_patterns:
                        if pattern.match(combined_text):
                            candidate_fragments.append({
                                **item,
                                'text': combined_text,
                                'outline_level': level,
                                'is_start': True,
                                'pattern': pattern,
                                'original_position': (current_pos, j - 1)
                            })
                            break
                    break

            # 如果不是编号部分，检查是否为完整标题或普通文本
            if not is_number_part:
                # print('不是编号一部分：', text)
                # 检查是否为完整标题
                best_match = None
                highest_priority = -1
                for pattern, level, priority in self.compiled_full_patterns:
                    if pattern.match(text):
                        if priority > highest_priority:
                            highest_priority = priority
                            best_match = (pattern, level)
                        break

                if best_match:
                    candidate_fragments.append({**item,
                                                'outline_level': best_match[1],
                                                'is_start': True,
                                                'pattern': best_match[0],
                                                'original_position': (i, i)
                                                })
                else:
                    # 可能是标题的一部分
                    if len(text) > 1 and len(candidate_fragments)>0: # not re.match(r'^[\d\W]+$', text)
                        candidate_fragments.append({
                            **item,
                            'outline_level': None,
                            'is_start': False,
                            'pattern': None,
                            'original_position': (i, i)
                        })

            i += 1
        return candidate_fragments

    def _should_merge(self, current_frag, next_frag, current_level, current_parent_id):
        """判断是否应该合并两个片段"""
        if next_frag['is_start'] and not re.search(r'^\d+\.$|^\d+$', current_frag['text']):
            return False
        if abs(next_frag['level'] - current_level) > 1:
            return False
        if next_frag['parent_identifier'] != current_parent_id:
            if not (current_frag['elem_index'] + 1 == next_frag['elem_index'] and
                    current_frag['level'] == next_frag['level']):
                return False
        if re.search('[\u4e00-\u9fa5]{2,}', current_frag['text']) and next_frag['element'].tag in ['div', 'p' , 'h1', 'h2', 'h3', 'td', 'th']:
            return False
        return True

    def _merge_fragments(self, candidates: List[Dict]) -> List[Dict]:
        """合并标题片段，处理短标题情况"""
        merged_outlines = []
        i = 0
        frag_count = len(candidates)

        while i < frag_count:
            fragment = candidates[i]

            if fragment['is_start']:
                merged_text = fragment['text']
                current_level = fragment['level']
                current_parent_id = fragment['parent_identifier']
                # start_position = fragment['position']

                start_position = fragment['original_position'][0]
                end_position = fragment['original_position'][1]

                elem_index = fragment['elem_index']
                elem_tag = fragment['element'].tag

                # 根据标题长度决定最大合并数量
                max_merge = self.max_merge_short if len(merged_text) < 5 else self.max_merge_normal
                j = i + 1
                merge_count = 0

                # 优化合并循环
                while j < frag_count and merge_count < max_merge:
                    next_frag = candidates[j]
                    if self._should_merge(fragment, next_frag, current_level, current_parent_id):
                        print('需合并大纲： ', merged_text, next_frag['text'])
                        merged_text += next_frag['text']
                        end_position = next_frag['original_position'][1]
                        j += 1
                        merge_count += 1
                    else:
                        break

                # 接受有效标题
                if len(merged_text) >= 2 and not self._is_non_outline(merged_text, is_final=True):
                    merged_outlines.append({
                        'text': merged_text,
                        'level': fragment['level'],
                        'position': start_position,
                        'start_position': start_position,
                        'end_position': end_position,
                        'original_fragments': (i, j - 1),
                        'outline_level': fragment['outline_level'],
                        'pattern': fragment['pattern']
                    })

                i = j
            else:
                i += 1

        return merged_outlines

    def wrap_free_text_with_span(self, tree):
        """处理直接文本内容，减少重复判断"""

        # 去除多列表格
        tables = tree.xpath('//table')
        # print('表格数量：', len(tables))
        for table in tables:
            trs = table.xpath('./tr|./tbody/tr')
            if trs:
                tds = trs[0].xpath('./td|./th')
                # print('表格,行数： %d, 列数：%d'%(len(trs), len(tds)))
                # print('第一行', trs[0].xpath('string(.)'))
                if len(tds) > 2:
                    parent = table.getparent()
                    if parent is not None:
                        parent.remove(table)
                        # print('去除多列表格,行数： %d, 列数：%d'%(len(trs), len(tds)))

        skip_tags = {'script', 'style'}
        for element in tree.iter():
            if element.tag in skip_tags:
                continue

            # 处理元素的直接文本
            if element.text and element.text.strip():
                # 有子元素时文本是游离文本
                if len(element) > 0:
                    span = etree.Element("span")
                    span.text = element.text
                    element.text = None
                    element.insert(0, span)

            # 处理tail文本（总是游离文本）
            if element.tail and element.tail.strip():
                parent = element.getparent()
                if parent is not None:
                    # span = etree.Element("span")
                    span = etree.Element("p") if element.tag == "br" else etree.Element("span")
                    span.text = element.tail
                    element.tail = None
                    index = parent.index(element) + 1
                    parent.insert(index, span)

    def extract(self, html_content: str) -> List[Dict]:
        """
        提取HTML内容中的大纲

        Args:
            html_content: HTML字符串内容

        Returns:
            包含大纲信息的列表，每个元素包含标题、分级、位置等信息
        """
        start_time = time.time()
        self.level_cache.clear()  # 清空缓存
        self.identifier_cache.clear()

        try:
            # 解析HTML
            tree = etree.HTML(html_content)
            # 修复游离文本标签
            self.wrap_free_text_with_span(tree)

            # 1. 收集文本节点
            text_nodes = self._collect_text_nodes(tree)
            if not text_nodes:
                logger.warning("未找到有效文本节点")
                return []
            # print('所有文本节点：', [it['text'] for it in text_nodes])

            # 2. 识别候选标题片段
            candidates = self._identify_candidates(text_nodes)
            logger.info(f"识别到候选标题片段: {len(candidates)}个")
            # print('候选标题：', [(it['text'], it['is_start']) for it in candidates])

            # 3. 合并标题片段
            merged_outlines = self._merge_fragments(candidates)
            logger.info(f"合并后得到标题: {len(merged_outlines)}个")
            # print('合并后标题：', [(it['text'],it['level'], it['outline_level']) for it in merged_outlines])

            # 4. 优化分级
            optimized_outlines = []
            for outline in merged_outlines:
                base_level = outline['outline_level']

                if outline['level'] < 4:
                    adjusted_level = max(1, base_level - 1)
                elif 4 <= outline['level'] <= 7:
                    adjusted_level = base_level
                else:
                    adjusted_level = min(5, base_level + 1)

                optimized_outlines.append({**outline, 'final_level': adjusted_level})
            # print('优化分级后大纲：', [(it['text'],it['final_level']) for it in optimized_outlines])
            # 5. 确定作用范围和下属文本
            result = []
            outline_count = len(optimized_outlines)
            for i in range(outline_count):
                outline = optimized_outlines[i]
                end_pos = len(text_nodes)

                # 寻找当前标题的结束位置
                for j in range(i + 1, outline_count):
                    if optimized_outlines[j]['final_level'] <= outline['final_level']:
                        # end_pos = optimized_outlines[j]['position'] - 1
                        end_pos = optimized_outlines[j]['start_position'] - 1
                        break

                # 提取下属文本（优化列表推导性能）
                sub_text = []

                # 拆分大纲中包含的内容 例子：一、采购单编号：P-XJ-24-00044544
                if re.search('[:：]', outline['text']):
                    text1, text2 = re.split('[:：]', outline['text'], maxsplit=1)
                    outline['text'] = text1
                    if text2 != "":
                        sub_text.append(text2)

                # pos = outline['position']
                pos = outline['end_position']
                for item in text_nodes:
                    if pos < item['position'] <= end_pos:
                        sub_text.append(item['text'])
                sub_text = ''.join(sub_text)

                result.append({
                    '标题': outline['text'],
                    '分级': outline['final_level'],
                    '位置': outline['position'],
                    '作用范围': f"{outline['position']}-{end_pos}",
                    '下属文本': sub_text,
                    '合并信息': f"合并了{outline['original_fragments'][1]-outline['original_fragments'][0]+1}个片段"
                })
                # result.append((outline['text'], sub_text))

            # 输出性能指标
            end_time = time.time()
            logger.info(f"提取完成，耗时: {end_time - start_time:.4f}秒，处理文本节点: {len(text_nodes)}个")

            return result

        except Exception as e:
            logger.error(f"提取过程出错: {str(e)}", exc_info=True)
            return []

    def extract_from_file(self, file_path: str, encoding: str = 'utf-8') -> List[Dict]:
        """
        从HTML文件中提取大纲

        Args:
            file_path: HTML文件路径
            encoding: 文件编码，默认为utf-8

        Returns:
            包含大纲信息的列表
        """
        try:
            with open(file_path, 'r', encoding=encoding) as f:
                html_content = f.read()
            return self.extract(html_content)
        except Exception as e:
            logger.error(f"从文件提取失败 {file_path}: {str(e)}")
            return []


# 使用示例
if __name__ == "__main__":
    # 创建提取器实例（可重复使用）
    extractor = OutlineExtractor()

    # # 示例1：从文件提取
    # outline = extractor.extract_from_file('d:/html/2.html')  # YC2023-10-0027中标招标错误  大纲范例1表格内容大纲
    # print('大纲：', outline)
    # if outline and isinstance(outline[0], tuple):
    #     print([it[0] for it in outline])
    #     print('内容：', outline)

    # # 打印结果
    # print(f"共提取到 {len(outline)} 个大纲标题\n")
    # for item in outline:
    #     print(f"**{item['分级']}级标题: {item['标题']}")
    #     print(f"下属文本: {item['下属文本'][:100]}...")  # 只显示前100字符
    #     print(f"分级: {item['分级']}")
    #     print(f"位置: {item['位置']}")
    #     print(f"作用范围: {item['作用范围']}")
    #     print("-" * 80)

    import pandas as pd
    import time


    def get_out_line_text(html_str):
        outline = extractor.extract(html_str)
        outline = [it['标题'][:20] for it in outline]
        # print(outline)
        return outline


    # df = pd.read_csv('E:\待检查数据/待检查_一周待筛选数据的公告_类别站源筛选数据_html.csv')
    # print('公告数量：', len(df))
    # datas = []
    # t1 = time.time()
    # for docid, html_str in zip(df['docid'], df['dochtmlcon']):
    #     outline = get_out_line_text(html_str)
    #     if outline:
    #         # print(outline)
    #         datas.append((docid, '\n'.join(outline)))
    #         # break
    #
    # df = pd.DataFrame(datas, columns=['docid', '大纲'])
    # print('len(df)', len(df))
    # print('耗时：', time.time()-t1)
    # df.to_excel('E:\待检查数据/待检查_一周待筛选数据的公告_类别站源筛选数据_大纲提取结果0928.xlsx', index=False)

    df = pd.read_csv('E:/地区匹配/20241213-19导出数据_站源地区筛选数据_html.csv')[:]
    # df2 = pd.read_excel('E:\大纲提取内容/20241213-19导出数据_站源地区筛选数据_大纲提取0425.xlsx')
    df2 = pd.read_excel('E:\大纲提取内容/20241213-19导出数据_站源地区筛选数据_大纲提取新旧方法对比0928.xlsx')
    print('公告数量：', len(df), len(df2))
    datas = []
    t1 = time.time()
    for docid, html_str in zip(df['docid'], df['dochtmlcon']):
        outline = get_out_line_text(html_str)
        # if outline:
            # print(outline)
        datas.append((docid, '\n'.join(outline), len(outline)))

    df = pd.DataFrame(datas, columns=['docid', '调整候选', '大纲数量4'])
    print('len(df)', len(df))
    print('耗时：', time.time()-t1)

    df = df.merge(df2, on='docid', how='inner')
    df.to_excel('E:\大纲提取内容/20241213-19导出数据_站源地区筛选数据_大纲提取新旧方法对比0928.xlsx', index=False)

    # df = pd.read_csv('E:/模版提取/中国南方电网-站源公告_html.csv')
    # print('公告数量：', len(df))
    # datas = []
    # t1 = time.time()
    # for docid, html_str in zip(df['docid'], df['dochtmlcon']):
    #     outline = get_out_line_text(html_str)
    #     if outline:
    #         # print(outline)
    #         datas.append((docid, '\n'.join(outline[:5])))
    #         # break
    #
    # df = pd.DataFrame(datas, columns=['docid', '大纲'])
    # print('len(df)', len(df))
    # print('耗时：', time.time()-t1)
    #
    # from collections import Counter
    # c = Counter(df['大纲'])
    # print(c.most_common(200))

    # df.to_excel('E:\模版提取/中国南方电网-站源公告_大纲提取结果.xlsx', index=False)

    from copy import deepcopy

    html_content = """

<div>
	<p>正文内容：</p>
	<p> 项目名称：</p>

</div>
<div class='fujian'>
附件文本， 相关附件内容
</div>

    """

    # def extract_div_with_metadata(html_content, div_xpath):
    #     """提取div并保留元数据信息"""
    #     doc = etree.HTML(html_content)
    #
    #     target_div = doc.xpath(div_xpath)
    #     if not target_div:
    #         return None, None
    #
    #     target_div = target_div[0]
    #
    #     # 创建新树
    #     new_doc = etree.Element("div") # extracted-content
    #     new_doc.append(deepcopy(target_div))
    #
    #     # # 添加元数据
    #     # metadata = etree.SubElement(new_doc, "metadata")
    #     # etree.SubElement(metadata, "original-xpath").text = div_xpath
    #     # etree.SubElement(metadata, "extraction-time").text = "2024-01-01"
    #
    #     # # 添加内容
    #     # content = etree.SubElement(new_doc, "content")
    #     # content.append(deepcopy(target_div))
    #
    #     # 从原文档移除
    #     parent = target_div.getparent()
    #     if parent is not None:
    #         parent.remove(target_div)
    #
    #     return etree.ElementTree(new_doc), etree.ElementTree(doc)
    #
    #
    # # 使用示例
    # new_tree, original_tree = extract_div_with_metadata(
    #     html_content,
    #     '//div[@class="fujian"]'
    # )
    # if new_tree:
    #     print('附件：', new_tree.xpath('string(.)'))
    #     print('正文：', original_tree.xpath('string(.)'))
    # # if new_tree:
    # #     print("带元数据的提取树:")
    # #     print(etree.tostring(new_tree, encoding='unicode', pretty_print=True))
    # # print('original_tree： ', etree.tostring(original_tree, encoding='unicode', pretty_print=True))