import fitz  # PyMuPDF
import re
import json
import os
from datetime import datetime

def extract_and_split_construction_plan(pdf_path):
    # 打开PDF文件
    doc = fitz.open(pdf_path)
    
    # 编译正则表达式
    chapter_pattern = re.compile(r'^第[一二三四五六七八九十百]+章\s*.*')
    section_pattern = re.compile(r'^[一二三四五六七八九十百]+、\s*.*')
    # 用于识别目录的特征：连续的三个以上小数点或省略号
    toc_pattern = re.compile(r'\.{3,}|…{2,}') 
    
    structured_data = {}
    current_chapter = "未分类前言"
    current_section = "默认部分"
    
    in_body = False  # 状态机：标记是否已经跳过目录，正式进入正文
    
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        
        # 1. 清理页眉页脚：利用 clip 裁剪页面提取区域
        # 默认A4纸高度约842磅，裁剪掉顶部和底部各60磅的区域（可根据实际PDF微调）
        rect = page.rect
        clip_box = fitz.Rect(0, 60, rect.width, rect.height - 60)
        
        # 仅提取裁剪框内的纯文本
        text = page.get_text("text", clip=clip_box)
        lines = text.split('\n')
        
        for line in lines:
            line = line.strip()
            # 跳过空行
            if not line:
                continue
            
            # 双保险：过滤掉可能因排版偏移漏掉的页眉页脚特征词或孤立的页码
            if "四川路桥建设集团股份有限公司" in line or "T梁运输及安装专项施工方案" in line or line.isdigit():
                continue
            
            # 2. 删除目录逻辑：判断是否正式进入正文
            if not in_body:
                if chapter_pattern.match(line) and not toc_pattern.search(line):
                    in_body = True
                else:
                    continue  # 还在目录页，直接跳过
            
            # 进入正文后的防干扰处理：跳过残余目录格式
            if toc_pattern.search(line):
                continue
            
            # 匹配到一级标题
            if chapter_pattern.match(line):
                current_chapter = line
                current_section = "章节前言" 
                if current_chapter not in structured_data:
                    structured_data[current_chapter] = {current_section: []}
                continue
            
            # 匹配到二级标题
            if section_pattern.match(line):
                current_section = line
                if current_chapter not in structured_data:
                    structured_data[current_chapter] = {}
                if current_section not in structured_data[current_chapter]:
                    structured_data[current_chapter][current_section] = []
                continue
            
            # 容错处理：确保基础字典结构存在
            if current_chapter not in structured_data:
                structured_data[current_chapter] = {current_section: []}
            if current_section not in structured_data[current_chapter]:
                structured_data[current_chapter][current_section] = []
                
            # 3. 将正文内容累加到对应的层级下
            structured_data[current_chapter][current_section].append(line)
    
    # 将列表拼接成完整的文本块
    for chap in structured_data:
        for sec in structured_data[chap]:
            structured_data[chap][sec] = '\n'.join(structured_data[chap][sec])
            
    return structured_data

if __name__ == "__main__":
    # 获取用户输入的路径
    user_input = input("请输入需要提取的PDF文件路径（支持直接拖入文件或粘贴路径）：")
    
    # 清理路径两端可能存在的引号和空格（应对“复制文件地址”或拖拽文件带来的双引号）
    pdf_file_path = user_input.strip('\'" ')
    
    # 检查文件是否存在
    if not os.path.exists(pdf_file_path):
        print(f"\n[错误] 找不到文件，请检查路径是否正确：{pdf_file_path}")
    else:
        print("\n开始提取施工方案，请稍候...")
        try:
            result_data = extract_and_split_construction_plan(pdf_file_path)
            
            # 4. 保存为本地JSON，名称为：文件名+当前时间（到秒）
            base_name = os.path.splitext(os.path.basename(pdf_file_path))[0]
            current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
            
            # 将输出文件保存在与原PDF相同的目录下
            output_dir = os.path.dirname(pdf_file_path)
            output_filename = os.path.join(output_dir, f"{base_name}_{current_time}.json")
            
            with open(output_filename, 'w', encoding='utf-8') as json_file:
                json.dump(result_data, json_file, ensure_ascii=False, indent=4)
                
            print(f"\n[成功] 提取完成！")
            print(f"结构化数据已保存至: {output_filename}")
            
        except Exception as e:
            print(f"\n[失败] 提取过程中发生错误: {e}")