"""
基于PDF文本层的智能题目解析器
借鉴ChatGPT的策略：直接提取PDF文本层，避免OCR误差
"""

import re
import json
import time
from typing import List, Dict, Tuple, Optional
import fitz  # PyMuPDF
import pdfplumber
from dataclasses import dataclass

@dataclass
class Question:
    section: str
    qno: str
    type: str
    stem: str
    options: Dict[str, str]
    page_start: int
    page_end: int
    confidence: float = 1.0

class PDFTextParser:
    def __init__(self):
        # 大题分区正则
        self.section_patterns = [
            re.compile(r'^\s*[（(]?([一二三四五六七八九十]{1,3})[、.)]\s*(\S*题)', re.MULTILINE),
            re.compile(r'^\s*(\d+)\s*[、.]\s*(\S*题)', re.MULTILINE),
        ]
        
        # 题号正则 - 支持多种格式
        self.question_patterns = [
            re.compile(r'^\s*(?:第\s*)?([0-9]{1,3}|[一二三四五六七八九十]{1,3})\s*[.)]'),
            re.compile(r'^\s*([0-9]{1,3})\s*[．、]\s*'),
            re.compile(r'^\s*\(([0-9]{1,3})\)\s*'),
            re.compile(r'^\s*（([0-9]{1,3})）\s*'),
        ]
        
        # 选项正则
        self.option_patterns = [
            re.compile(r'^\s*([A-H])\s*[.)]\s*(.*)'),
            re.compile(r'^\s*（([A-H])）\s*(.*)'),
            re.compile(r'^\s*\(([A-H])\)\s*(.*)'),
        ]
        
        # 中文数字转换
        self.chinese_numbers = {
            '一': '1', '二': '2', '三': '3', '四': '4', '五': '5',
            '六': '6', '七': '7', '八': '8', '九': '9', '十': '10',
            '十一': '11', '十二': '12', '十三': '13', '十四': '14', '十五': '15',
            '十六': '16', '十七': '17', '十八': '18', '十九': '19', '二十': '20'
        }

    def extract_text_with_pymupdf(self, pdf_path: str) -> List[Dict]:
        """
        使用PyMuPDF提取文本（优先，速度快，顺序准确）
        """
        try:
            doc = fitz.open(pdf_path)
            pages_text = []
            
            for page_num in range(len(doc)):
                page = doc[page_num]
                text = page.get_text()
                
                # 按行分割并保持行信息
                lines = []
                for line in text.split('\n'):
                    line = line.strip()
                    if line:  # 过滤空行
                        lines.append(line)
                
                pages_text.append({
                    'page_num': page_num + 1,
                    'lines': lines,
                    'full_text': text
                })
            
            doc.close()
            print(f"✅ PyMuPDF提取完成: {len(pages_text)}页")
            return pages_text
            
        except Exception as e:
            print(f"❌ PyMuPDF提取失败: {e}")
            return []

    def extract_text_with_pdfplumber(self, pdf_path: str) -> List[Dict]:
        """
        使用pdfplumber提取文本（备用，处理复杂布局）
        """
        try:
            pages_text = []
            
            with pdfplumber.open(pdf_path) as pdf:
                for page_num, page in enumerate(pdf.pages):
                    text = page.extract_text()
                    
                    if text:
                        lines = []
                        for line in text.split('\n'):
                            line = line.strip()
                            if line:
                                lines.append(line)
                        
                        pages_text.append({
                            'page_num': page_num + 1,
                            'lines': lines,
                            'full_text': text
                        })
            
            print(f"✅ pdfplumber提取完成: {len(pages_text)}页")
            return pages_text
            
        except Exception as e:
            print(f"❌ pdfplumber提取失败: {e}")
            return []

    def extract_pdf_text(self, pdf_path: str) -> List[Dict]:
        """
        智能文本提取：优先PyMuPDF，失败时回退pdfplumber
        """
        print(f"🔄 开始PDF文本提取: {pdf_path}")
        start_time = time.time()
        
        # 优先使用PyMuPDF
        pages_text = self.extract_text_with_pymupdf(pdf_path)
        
        # 如果PyMuPDF失败或结果为空，回退到pdfplumber
        if not pages_text:
            print("⚠️ PyMuPDF结果为空，回退到pdfplumber")
            pages_text = self.extract_text_with_pdfplumber(pdf_path)
        
        print(f"⚡ 文本提取完成: 耗时{time.time() - start_time:.2f}秒")
        return pages_text

    def normalize_text(self, text: str) -> str:
        """
        文本标准化：处理全角半角、中文数字等
        """
        # 全角转半角
        text = text.replace('（', '(').replace('）', ')')
        text = text.replace('．', '.').replace('、', '.')
        text = text.replace('：', ':').replace('；', ';')
        
        # 去除多余空格
        text = re.sub(r'\s+', ' ', text).strip()
        
        return text

    def detect_sections(self, pages_text: List[Dict]) -> List[Dict]:
        """
        检测大题分区（一、选择题，二、填空题等）
        """
        sections = []
        current_section = None
        
        for page_data in pages_text:
            page_num = page_data['page_num']
            
            for line in page_data['lines']:
                normalized_line = self.normalize_text(line)
                
                # 检测分区标题
                for pattern in self.section_patterns:
                    match = pattern.match(normalized_line)
                    if match:
                        section_num = match.group(1)
                        section_name = match.group(2)
                        
                        # 转换中文数字
                        if section_num in self.chinese_numbers:
                            section_num = self.chinese_numbers[section_num]
                        
                        current_section = {
                            'section_num': section_num,
                            'section_name': section_name,
                            'start_page': page_num,
                            'start_line': normalized_line
                        }
                        sections.append(current_section)
                        print(f"📚 检测到分区: {section_num}、{section_name} (第{page_num}页)")
                        break
        
        return sections

    def convert_chinese_number(self, num_str: str) -> str:
        """
        转换中文数字为阿拉伯数字
        """
        if num_str in self.chinese_numbers:
            return self.chinese_numbers[num_str]
        return num_str

    def detect_question_number(self, line: str) -> Optional[str]:
        """
        检测题号
        """
        normalized_line = self.normalize_text(line)
        
        for pattern in self.question_patterns:
            match = pattern.match(normalized_line)
            if match:
                qno = match.group(1)
                # 转换中文数字
                qno = self.convert_chinese_number(qno)
                try:
                    # 验证是否为合理的题号
                    num = int(qno)
                    if 1 <= num <= 100:  # 假设题目数量在1-100之间
                        return str(num)
                except ValueError:
                    continue
        
        return None

    def detect_option(self, line: str) -> Optional[Tuple[str, str]]:
        """
        检测选项（A、B、C、D等）
        """
        normalized_line = self.normalize_text(line)
        
        for pattern in self.option_patterns:
            match = pattern.match(normalized_line)
            if match:
                option_letter = match.group(1)
                option_content = match.group(2).strip()
                return option_letter, option_content
        
        return None

    def parse_questions(self, pages_text: List[Dict]) -> List[Question]:
        """
        解析题目：核心算法
        """
        print("🔄 开始智能题目解析")
        start_time = time.time()
        
        questions = []
        current_question = None
        current_section = "未知"
        all_lines = []
        
        # 合并所有页面的行，保持页码信息
        for page_data in pages_text:
            for line in page_data['lines']:
                all_lines.append({
                    'text': line,
                    'page': page_data['page_num']
                })
        
        i = 0
        while i < len(all_lines):
            line_data = all_lines[i]
            line = line_data['text']
            page_num = line_data['page']
            
            # 检测分区
            for pattern in self.section_patterns:
                match = pattern.match(self.normalize_text(line))
                if match:
                    section_name = match.group(2)
                    current_section = section_name
                    print(f"📚 进入分区: {section_name}")
                    break
            
            # 检测题号
            qno = self.detect_question_number(line)
            if qno:
                # 保存上一题
                if current_question:
                    questions.append(current_question)
                
                # 开始新题
                question_type = self.infer_question_type(current_section, line)
                current_question = Question(
                    section=current_section,
                    qno=qno,
                    type=question_type,
                    stem=line,
                    options={},
                    page_start=page_num,
                    page_end=page_num
                )
                print(f"📝 题目{qno}: {line[:50]}...")
                
            elif current_question:
                # 检测选项
                option_result = self.detect_option(line)
                if option_result:
                    option_letter, option_content = option_result
                    current_question.options[option_letter] = option_content
                    current_question.page_end = page_num
                else:
                    # 继续当前题目内容
                    if len(current_question.stem) < 1000:  # 防止过长
                        current_question.stem += ' ' + line
                        current_question.page_end = page_num
            
            i += 1
        
        # 保存最后一题
        if current_question:
            questions.append(current_question)
        
        print(f"✅ 题目解析完成: {len(questions)}道题, 耗时{time.time() - start_time:.2f}秒")
        return questions

    def infer_question_type(self, section: str, line: str) -> str:
        """
        推断题目类型
        """
        if '选择' in section:
            return 'multiple_choice'
        elif '填空' in section:
            return 'fill_blank'
        elif '解答' in section or '计算' in section:
            return 'solve'
        elif '判断' in section:
            return 'true_false'
        else:
            # 根据内容推断
            if re.search(r'[A-D][.)）]', line):
                return 'multiple_choice'
            elif '______' in line or '___' in line:
                return 'fill_blank'
            else:
                return 'unknown'

    def process_pdf(self, pdf_path: str) -> Dict:
        """
        主处理函数
        """
        total_start = time.time()
        print(f"🎯 开始基于文本层的PDF解析: {pdf_path}")
        
        try:
            # 1. 提取PDF文本
            pages_text = self.extract_pdf_text(pdf_path)
            
            if not pages_text:
                return {"success": False, "error": "PDF文本提取失败"}
            
            # 2. 解析题目
            questions = self.parse_questions(pages_text)
            
            # 3. 转换为API格式
            formatted_questions = []
            for q in questions:
                formatted_questions.append({
                    'number': q.qno,
                    'type': q.type,
                    'text': q.stem,
                    'options': q.options if q.options else None,
                    'page': q.page_start,
                    'bbox': [],  # 文本层解析无bbox信息
                    'confidence': q.confidence,
                    'page_spans': list(range(q.page_start, q.page_end + 1)),
                    'source_text': q.stem,
                    'images': [],
                    'auto_extracted': True,
                    'section': q.section
                })
            
            total_time = time.time() - total_start
            print(f"🎉 文本层解析完成！总耗时: {total_time:.2f}秒")
            
            return {
                "success": True,
                "questions": formatted_questions,
                "total_time": total_time,
                "pages_processed": len(pages_text),
                "questions_found": len(formatted_questions),
                "stats": {
                    "total_time": total_time,
                    "pages_processed": len(pages_text),
                    "questions_found": len(formatted_questions),
                    "avg_time_per_page": total_time / len(pages_text) if pages_text else 0,
                    "method": "text_layer_parsing"
                }
            }
            
        except Exception as e:
            print(f"❌ 文本层解析失败: {str(e)}")
            return {"success": False, "error": str(e)}

    def export_to_json(self, questions: List[Question], output_path: str):
        """
        导出为JSON格式
        """
        data = []
        for q in questions:
            data.append({
                'section': q.section,
                'qno': q.qno,
                'type': q.type,
                'stem': q.stem,
                'options': q.options,
                'page_start': q.page_start,
                'page_end': q.page_end,
                'confidence': q.confidence
            })
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        print(f"✅ JSON导出完成: {output_path}")

# 测试函数
def test_text_parser():
    """
    测试文本解析器
    """
    pdf_path = "/Users/brunogao/Downloads/深圳市八年级（上）期中数学试卷1.pdf"
    
    if not os.path.exists(pdf_path):
        print(f"❌ 测试文件不存在: {pdf_path}")
        return
    
    parser = PDFTextParser()
    result = parser.process_pdf(pdf_path)
    
    if result["success"]:
        print(f"""
📊 文本层解析测试结果:
- 总时间: {result['total_time']:.2f}秒
- 页数: {result['pages_processed']}
- 题目: {result['questions_found']}道
- 方法: {result['stats']['method']}
        """)
        
        print("识别的题目:")
        for i, q in enumerate(result['questions'][:8]):
            options_str = f" 选项: {list(q.get('options', {}).keys())}" if q.get('options') else ""
            print(f"题目{q['number']} ({q['type']}): {q['text'][:60]}...{options_str}")
    else:
        print(f"❌ 测试失败: {result['error']}")

if __name__ == "__main__":
    import os
    test_text_parser()