import fitz  # PyMuPDF
import argparse
import os

class PDFTextFinder:
    def __init__(self, pdf_path):
        """
        初始化PDF文本查找器
        
        Args:
            pdf_path: PDF文件路径
        """
        self.pdf_path = pdf_path
        # 检查文件是否存在
        if not os.path.exists(pdf_path):
            raise FileNotFoundError(f"错误: 文件 {pdf_path} 不存在!")
        
        # 打开PDF文件
        self.doc = fitz.open(pdf_path)
        
    def __del__(self):
        """
        析构函数，确保关闭PDF文件
        """
        if hasattr(self, 'doc') and self.doc:
            self.doc.close()
    
    def find_text(self, search_text):
        """
        在PDF中查找文本并返回所在页码
        
        Args:
            search_text: 要查找的文本
            
        Returns:
            list: 包含文本出现的页码列表，页码从1开始
        """
        result_pages = []
        
        # 遍历所有页面
        for page_num in range(len(self.doc)):
            page = self.doc.load_page(page_num)
            text_instances = page.search_for(search_text)  # 搜索文本
            
            if text_instances:  # 如果找到了文本
                # 页码从0开始，所以需要+1
                result_pages.append(page_num + 1)
        
        return result_pages
    
    def find_text_with_context(self, search_text, context_words=20):
        """
        在PDF中查找文本并返回所在页码及上下文
        
        Args:
            search_text: 要查找的文本
            context_words: 返回的上下文单词数量
            
        Returns:
            list: 包含字典的列表，每个字典包含页码、文本位置和上下文
        """
        results = []
        
        # 遍历所有页面
        for page_num in range(len(self.doc)):
            page = self.doc.load_page(page_num)
            text_instances = page.search_for(search_text)  # 搜索文本
            
            if text_instances:  # 如果找到了文本
                # 获取页面文本
                page_text = page.get_text()
                
                # 对于每个找到的实例，提取上下文
                for inst in text_instances:
                    # 找到文本在页面中的位置
                    text_pos = page_text.find(search_text)
                    if text_pos >= 0:
                        # 计算上下文的起始和结束位置
                        start_pos = max(0, text_pos - context_words * 5)  # 估算每个单词平均5个字符
                        end_pos = min(len(page_text), text_pos + len(search_text) + context_words * 5)
                        
                        # 提取上下文
                        context = page_text[start_pos:end_pos]
                        
                        # 添加到结果列表
                        results.append({
                            'page': page_num + 1,  # 页码从0开始，所以需要+1
                            'position': inst,  # 文本的位置信息
                            'context': context  # 上下文
                        })
        
        return results

def find_text_in_pdf(pdf_path, search_text):
    """
    在PDF中查找文本并打印所在页码
    
    Args:
        pdf_path: PDF文件路径
        search_text: 要查找的文本
    """
    try:
        finder = PDFTextFinder(pdf_path)
        pages = finder.find_text(search_text)
        
        if pages:
            print(f"在文件 '{pdf_path}' 中找到文本 '{search_text}' 的页码:")
            for page in pages:
                print(f"第 {page} 页")
            print(f"总共在 {len(pages)} 个页面中找到匹配文本。")
        else:
            print(f"在文件 '{pdf_path}' 中没有找到文本 '{search_text}'。")
            
        return pages
    except Exception as e:
        print(f"查找文本时出错: {str(e)}")
        return []

def find_text_with_context_in_pdf(pdf_path, search_text, context_words=20):
    """
    在PDF中查找文本并打印所在页码及上下文
    
    Args:
        pdf_path: PDF文件路径
        search_text: 要查找的文本
        context_words: 返回的上下文单词数量
    """
    try:
        finder = PDFTextFinder(pdf_path)
        results = finder.find_text_with_context(search_text, context_words)
        
        if results:
            print(f"在文件 '{pdf_path}' 中找到文本 '{search_text}' 的位置:")
            for result in results:
                print(f"第 {result['page']} 页:")
                print(f"上下文: {result['context']}")
                print("-" * 50)
            print(f"总共找到 {len(results)} 处匹配。")
        else:
            print(f"在文件 '{pdf_path}' 中没有找到文本 '{search_text}'。")
            
        return results
    except Exception as e:
        print(f"查找文本时出错: {str(e)}")
        return []

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='在PDF中查找文本并返回页码')
    parser.add_argument('--pdf', required=True, help='PDF文件路径')
    parser.add_argument('--text', required=True, help='要查找的文本')
    parser.add_argument('--context', action='store_true', help='是否显示上下文')
    parser.add_argument('--context_words', type=int, default=20, help='上下文单词数量')
    
    args = parser.parse_args()
    
    # 检查文件是否存在
    if not os.path.exists(args.pdf):
        print(f"错误: 文件 {args.pdf} 不存在!")
        return
    
    print(f"开始在文件 '{args.pdf}' 中查找文本 '{args.text}'...")
    
    # 根据是否需要上下文调用不同的函数
    if args.context:
        find_text_with_context_in_pdf(args.pdf, args.text, args.context_words)
    else:
        find_text_in_pdf(args.pdf, args.text)

if __name__ == "__main__":
    main()