"""
修复版PDF试卷解析器 - 专门解决数学符号提取问题
针对数学符号丢失的问题进行特殊处理
"""

import re
import json
import math
import time
import unicodedata
from collections import Counter, defaultdict
from typing import List, Dict, Optional, Tuple
import fitz  # PyMuPDF
import pdfplumber

class PDFExamParserFixed:
    def __init__(self):
        # 全角转半角映射表
        self.fullwidth_trans = str.maketrans(
            "０１２３４５６７８９．（）ＡＢＣＤＥＦＧＨＩＪＫＬＭＮＯＰＱＲＳＴＵＶＷＸＹＺａｂｃｄｅｆｇｈｉｊｋｌｍｎｏｐｑｒｓｔｕｖｗｘｙｚ",
            "0123456789.()ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
        )
        
        # 大题分区正则
        self.heading_pattern = re.compile(r"^\s*[（(]?\s*([一二三四五六七八九十]{1,3})\s*[、.)]\s*(\S*题)")
        
        # 题号正则 - 更宽松的匹配
        self.question_start_pattern = re.compile(r"^\s*([0-9]{1,2})\s*[.．、)]")
        
        # 选项正则 - 更灵活的匹配
        self.option_label = r"([A-H])\s*[\.．、\)）]"
        self.option_line_begin = re.compile(r"^\s*" + self.option_label + r"\s*(.*)")
        self.option_inline_finder = re.compile(self.option_label)

    def normalize_text(self, text: str) -> str:
        """
        文本标准化：全角转半角、Unicode规范化等
        """
        if not text:
            return ""
        
        # 全角转半角
        text = text.translate(self.fullwidth_trans)
        
        # Unicode规范化
        text = unicodedata.normalize("NFKC", text)
        
        # 替换特殊空格
        text = text.replace("\u00A0", " ")
        
        # 标点符号统一
        text = text.replace("．", ".").replace("：", ":")
        
        # 清理空行并去除首尾空格
        text = "\n".join([line.strip() for line in text.splitlines() if line.strip()])
        
        return text

    def extract_text_with_coords(self, pdf_path: str) -> List[Dict]:
        """
        提取带坐标信息的文本，用于更好地理解布局
        """
        doc = fitz.open(pdf_path)
        all_text_items = []
        
        for page_num, page in enumerate(doc):
            # 使用字典方式获取文本和坐标
            text_dict = page.get_text("dict")
            
            for block in text_dict["blocks"]:
                if "lines" in block:  # 文本块
                    for line in block["lines"]:
                        for span in line["spans"]:
                            text = span["text"].strip()
                            if text:
                                all_text_items.append({
                                    "page": page_num + 1,
                                    "x0": span["bbox"][0],
                                    "y0": span["bbox"][1], 
                                    "x1": span["bbox"][2],
                                    "y1": span["bbox"][3],
                                    "text": text,
                                    "font": span.get("font", ""),
                                    "size": span.get("size", 0)
                                })
        
        doc.close()
        
        # 按页面和位置排序
        all_text_items.sort(key=lambda x: (x["page"], x["y0"], x["x0"]))
        
        return all_text_items

    def reconstruct_lines(self, text_items: List[Dict]) -> List[Dict]:
        """
        重构文本行，合并同一行的文本片段
        """
        lines = []
        current_line = None
        
        for item in text_items:
            # 如果是新行或者第一个项目
            if (current_line is None or 
                item["page"] != current_line["page"] or
                abs(item["y0"] - current_line["y0"]) > 3):  # 减小容差到3像素
                
                # 保存上一行
                if current_line:
                    lines.append(current_line)
                
                # 开始新行
                current_line = {
                    "page": item["page"],
                    "y0": item["y0"],
                    "x0": item["x0"],
                    "text": item["text"],
                    "items": [item]
                }
            else:
                # 继续当前行 - 检查是否需要空格
                last_item = current_line["items"][-1] if current_line["items"] else None
                gap = item["x0"] - (last_item["x1"] if last_item else 0)
                separator = " " if gap > 5 else ""  # 如果间隔大于5像素才加空格
                current_line["text"] += separator + item["text"]
                current_line["items"].append(item)
                current_line["x0"] = min(current_line["x0"], item["x0"])
        
        # 保存最后一行
        if current_line:
            lines.append(current_line)
        
        return lines

    def extract_options_improved(self, question_text: str, text_items: List[Dict]) -> Tuple[str, List[Dict]]:
        """
        改进的选项提取，处理数学符号缺失问题
        """
        lines = question_text.splitlines()
        options = []
        stem_buffer = []
        
        # 首先尝试标准方法
        current_option = None
        
        for line in lines:
            line = line.strip()
            if not line:
                continue
            
            # 检测选项标签
            option_match = self.option_line_begin.match(line)
            if option_match:
                # 保存前一个选项
                if current_option:
                    options.append(current_option)
                
                label = option_match.group(1)
                text = option_match.group(2).strip()
                
                # 如果选项文本为空或很短，尝试从坐标信息中补充
                if not text or len(text) < 2:
                    text = self.find_option_content_by_coords(label, text_items)
                
                current_option = {"label": label, "text": text}
            else:
                # 如果当前在处理选项，继续添加到选项内容
                if current_option:
                    current_option["text"] += " " + line
                else:
                    stem_buffer.append(line)
        
        # 保存最后一个选项
        if current_option:
            options.append(current_option)
        
        # 如果没有找到选项，尝试从完整文本中解析
        if not options:
            options = self.extract_options_fallback(question_text)
        
        # 清理题干
        stem = "\n".join(stem_buffer).strip()
        if not stem:
            # 从完整文本中移除选项部分作为题干
            stem = self.extract_stem_from_full_text(question_text, options)
        
        return stem, options

    def find_option_content_by_coords(self, label: str, text_items: List[Dict]) -> str:
        """
        通过坐标信息查找选项内容，用于处理符号缺失
        """
        # 查找选项标签的位置
        label_item = None
        for item in text_items:
            if f"{label}." in item["text"] or f"{label}．" in item["text"]:
                label_item = item
                break
        
        if not label_item:
            return ""
        
        # 查找同一行或邻近的其他文本
        content_items = []
        for item in text_items:
            if (item["page"] == label_item["page"] and
                abs(item["y0"] - label_item["y0"]) <= 3 and  # 同一行
                item["x0"] > label_item["x1"]):  # 在标签右侧
                content_items.append(item)
        
        # 按x坐标排序并合并
        content_items.sort(key=lambda x: x["x0"])
        content = " ".join([item["text"] for item in content_items])
        
        return content.strip()

    def extract_options_fallback(self, question_text: str) -> List[Dict]:
        """
        备用选项提取方法
        """
        options = []
        
        # 分割选项的方法：查找A. B. C. D.模式
        pattern = r"([A-D])\s*[\.．]\s*([^A-D]*?)(?=\s*[A-D]\s*[\.．]|$)"
        matches = re.findall(pattern, question_text, re.DOTALL)
        
        for label, content in matches:
            content = re.sub(r'\s+', ' ', content.strip())
            if content:
                options.append({"label": label, "text": content})
        
        return options

    def extract_stem_from_full_text(self, question_text: str, options: List[Dict]) -> str:
        """
        从完整文本中提取题干
        """
        stem = question_text
        
        # 移除题号前缀
        stem = re.sub(r'^\s*\d+\s*[\.．]\s*\([^)]*\)\s*', '', stem)
        
        # 移除选项内容
        for option in options:
            pattern = f"{option['label']}\\s*[\\.．]\\s*{re.escape(option['text'])}"
            stem = re.sub(pattern, '', stem, flags=re.IGNORECASE)
        
        # 清理多余空白
        stem = re.sub(r'\s+', ' ', stem).strip()
        
        return stem

    def process_pdf(self, pdf_path: str) -> Dict:
        """
        主处理函数
        """
        print(f"🎯 开始修复版PDF解析: {pdf_path}")
        start_time = time.time()
        
        try:
            # 1. 提取带坐标的文本
            text_items = self.extract_text_with_coords(pdf_path)
            print(f"📝 提取到{len(text_items)}个文本片段")
            
            # 2. 重构文本行
            lines = self.reconstruct_lines(text_items)
            print(f"📄 重构为{len(lines)}行文本")
            
            # 3. 简化题目识别
            questions = []
            current_section = "未知"
            current_question = None
            
            for line in lines:
                text = self.normalize_text(line["text"])
                
                # 检测分区
                heading_match = self.heading_pattern.match(text)
                if heading_match:
                    current_section = f"{heading_match.group(1)}、{heading_match.group(2)}"
                    print(f"📚 发现分区: {current_section}")
                    continue
                
                # 检测题号
                question_match = self.question_start_pattern.match(text)
                if question_match:
                    # 保存上一题
                    if current_question:
                        questions.append(current_question)
                    
                    question_num = question_match.group(1)
                    print(f"📝 发现题目: {question_num}")
                    
                    # 收集这道题的所有文本（包括后续行直到下一题）
                    question_text = text
                    question_items = line["items"]
                    
                    current_question = {
                        "number": question_num,
                        "section": current_section,
                        "text": question_text,
                        "items": question_items,
                        "page": line["page"]
                    }
                elif current_question:
                    # 继续当前题目
                    current_question["text"] += "\n" + text
                    current_question["items"].extend(line["items"])
            
            # 保存最后一题
            if current_question:
                questions.append(current_question)
            
            # 4. 处理每道题的选项
            formatted_questions = []
            for q in questions:
                prefer_choice = "选择" in q["section"]
                
                if prefer_choice:
                    stem, options = self.extract_options_improved(q["text"], q["items"])
                    
                    # 转换选项格式
                    options_dict = {}
                    for opt in options:
                        options_dict[opt["label"]] = opt["text"] if opt["text"] else f"[数学符号-{opt['label']}]"
                else:
                    stem = q["text"]
                    options_dict = None
                
                # 推断题目类型
                question_type = "subjective"
                if "选择" in q["section"]:
                    question_type = "single_choice"
                elif "填空" in q["section"]:
                    question_type = "fill_blank"
                
                formatted_questions.append({
                    'number': q['number'],
                    'type': question_type,
                    'text': stem,
                    'options': options_dict,
                    'page': q['page'],
                    'bbox': [],
                    'confidence': 1.0,
                    'page_spans': [q['page']],
                    'source_text': stem,
                    'images': [],
                    'auto_extracted': True,
                    'section': q['section']
                })
            
            total_time = time.time() - start_time
            print(f"🎉 修复版解析完成！总耗时: {total_time:.2f}秒")
            
            return {
                "success": True,
                "questions": formatted_questions,
                "total_time": total_time,
                "pages_processed": max([line["page"] for line in lines]) if lines else 0,
                "questions_found": len(formatted_questions),
                "stats": {
                    "total_time": total_time,
                    "pages_processed": max([line["page"] for line in lines]) if lines else 0,
                    "questions_found": len(formatted_questions),
                    "avg_time_per_page": total_time / max([line["page"] for line in lines]) if lines else 0,
                    "method": "fixed_text_parsing"
                }
            }
            
        except Exception as e:
            print(f"❌ 修复版解析失败: {str(e)}")
            return {"success": False, "error": str(e)}

# 测试函数
def test_fixed_parser():
    """
    测试修复版解析器
    """
    import os
    
    pdf_path = "/Users/brunogao/Downloads/深圳市八年级（上）期中数学试卷1.pdf"
    
    if not os.path.exists(pdf_path):
        print(f"❌ 测试文件不存在: {pdf_path}")
        return
    
    parser = PDFExamParserFixed()
    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("前3道选择题详细信息:")
        for i, q in enumerate(result['questions'][:3]):
            if q.get('type') == 'single_choice':
                print(f"\\n题目{q['number']}:")
                print(f"题干: {q['text']}")
                if q.get('options'):
                    print(f"选项:")
                    for k, v in q['options'].items():
                        print(f"  {k}. {v}")
    else:
        print(f"❌ 测试失败: {result['error']}")

if __name__ == "__main__":
    test_fixed_parser()