import asyncio
from core.config import root_path
from utils.doc_util import DocBase
from spire.pdf import PdfDocument, PdfTextExtractOptions, PdfTextExtractor, License, PdfTextFinder, TextFindParameter, \
    PdfTextMarkupAnnotation, PdfRGBColor, Color, PdfPolyLineAnnotation, PointF, RectangleF, PdfPopupAnnotation, \
    PdfPopupIcon
from loguru import logger
from rapidfuzz import process
from utils.ocr_util import OCRUtil
import copy
from utils.common_util import adjust_single_chunk_size
import re
import os

ocr_thresholds = 1000

def is_messy_text(text: str,
                  min_chars=40,
                  chinese_ratio_thresh=0.20,
                  printable_ratio_thresh=0.70,
                  symbol_ratio_thresh=0.30,
                  longest_non_word_run_thresh=10,
                  english_word_density_thresh=0.03):
    """
    判断单页或一段 text 是否 '乱码'。如果返回 True 表示该 text 被判为乱码（需要 OCR）。
    参数可调整：
      - min_chars: 少于这个字符数直接认为质量差（比如抽出来是碎行）
      - chinese_ratio_thresh: 中文字符占比阈值（若低于且其它判据也成立则可能乱码）
      - printable_ratio_thresh: 可打印字符占比（太低说明很多不可见/非标准字符）
      - symbol_ratio_thresh: 标点/特殊符号占比阈值（高于则更像乱码）
      - longest_non_word_run_thresh: 最长连续特殊符号串长度阈值（如 "~~##@@@... "）
      - english_word_density_thresh: 英文单词密度（单词数/行数或单词数/字符数近似）；若很低且无中文，则文本不可读
    """
    if not text:
        return True

    text_len = len(text)
    if text_len < min_chars:
        # 太短的片段通常是拆分后的乱码或空白
        return True

    # 统计中文、字母/数字、可打印、空白
    chinese_count = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
    alnum_count = sum(1 for c in text if c.isalnum())
    printable_count = sum(1 for c in text if c.isprintable())
    space_count = sum(1 for c in text if c.isspace())

    # 标点/特殊符号（既不是中文也不是字母数字也不是空白，视为符号）
    symbol_count = sum(1 for c in text if not (('\u4e00' <= c <= '\u9fff') or c.isalnum() or c.isspace()))

    chinese_ratio = chinese_count / text_len
    printable_ratio = printable_count / text_len
    symbol_ratio = symbol_count / text_len

    # 连续的非字/非中文/非数字串（比如 "~~~##@@@"）
    non_word_runs = re.findall(r'[^0-9A-Za-z\u4e00-\u9fff\s]+', text)
    longest_non_word_run = max((len(s) for s in non_word_runs), default=0)

    # 英文单词密度（粗略）：单词数 / 总字符数，单词定义为至少两个字母
    english_words = re.findall(r'\b[a-zA-Z]{2,}\b', text)
    english_word_density = len(english_words) / max(1, text_len)

    # 启发式判断：多条规则投票（其中任意一条强指标触发即可）
    # 强指标：非常高的符号占比、非常长的连续符号、几乎不可打印
    if printable_ratio < printable_ratio_thresh * 0.6:
        return True
    if symbol_ratio > max(0.5, symbol_ratio_thresh):
        return True
    if longest_non_word_run >= longest_non_word_run_thresh * 1.5:
        return True

    # 组合判断：如果中文占比很低且英文单词密度也很低，并且符号占比高 -> 乱码
    if chinese_ratio < chinese_ratio_thresh and english_word_density < english_word_density_thresh and symbol_ratio > symbol_ratio_thresh:
        return True

    # 如果中文占比极低且打印字符不多，也判为乱码
    if chinese_ratio < (chinese_ratio_thresh * 0.5) and printable_ratio < printable_ratio_thresh:
        return True

    # 其它情况视为可读
    return False


class SpirePdfDoc(DocBase):
    def get_chunk_location(self, chunk_idx):
        chunk_item = self._chunk_list[chunk_idx]
        return '、'.join(str(page + 1) for page in chunk_item['page_idx'])

    # 调整块大小
    def adjust_chunk_size(self):
        all_text_len = len(self.get_all_text())
        self._max_single_chunk_size = adjust_single_chunk_size(all_text_len)
        logger.info(f'SpirePdfDoc adjust _max_single_chunk_size to {self._max_single_chunk_size}')
        self.reset()
        return self._max_single_chunk_size

    async def get_from_ocr(self):
        """
        新策略：
          - 先用 Spire 的提取结果（self._chunk_list 已由 reset/init_chunks 填充）
          - 对每一页单独用 is_messy_text 判断
          - 如果被判为乱码的页占比超过 page_messy_ratio_thresh（比如 0.4），则触发 OCR（按页OCR）
          - 若被判为乱码的页数量少（但某些页极差），也可触发 OCR（按需要）
        """
        # 先从当前 chunk_list 拼出每页文本（init_chunks 保证 chunk_list 基于 PdfTextExtractor）
        # 注意 chunk_list 中 page 是每个 chunk 的 page 文本列表，我们要逐页统计
        all_pages = []
        for chunk in self._chunk_list:
            # chunk['page'] 是字符串列表，对应 chunk['page_idx'] 顺序一致
            all_pages.extend(chunk.get('page', []))

        # 如果没有任何页面文本（极端情况），强制 OCR
        if len(all_pages) == 0:
            logger.info('Spire extract returned 0 pages text, forcing OCR.')
            do_ocr = True
        else:
            # 针对每页判断是否为乱码
            page_messy_flags = []
            for i, page_text in enumerate(all_pages):
                messy = is_messy_text(page_text)
                page_messy_flags.append(messy)
                logger.debug(f'page {i + 1}: len={len(page_text)} messy={messy}')

            messy_count = sum(1 for f in page_messy_flags if f)
            total_pages = len(all_pages)
            messy_ratio = messy_count / total_pages
            logger.info(f'spire extract pages: {total_pages}, messy pages: {messy_count}, ratio: {messy_ratio:.2f}')

            # 触发 OCR 的条件（可调）：
            # 1) 若多数页为乱码 -> OCR
            # 2) 若至少一页非常差（例如 is_messy_text 在内部检测到极端情况） -> OCR
            page_messy_ratio_thresh = 0.40  # 如果超过 40% 页被判为乱码则 OCR
            # 若某页文本长度非常短（如 <20）且很多页属于这种短页 -> OCR
            many_short_pages = sum(1 for p in all_pages if len(p) < 20) / total_pages > 0.3

            do_ocr = (messy_ratio > page_messy_ratio_thresh) or many_short_pages

            # 作为后备：如果全文字符数低于 ocr_thresholds（原逻辑），同样 OCR
            all_text = '\n'.join(all_pages)
            if len(all_text) < ocr_thresholds:
                do_ocr = True

        if do_ocr:
            page_result_list = await self.ocr_util.ocr_result_pdf(self._doc_path)
            logger.info(f'ocr success: len[{len(page_result_list)}]')
            # 用 OCR 结果重新构建 chunk_list（与原逻辑一致）
            chunk_list = []
            chunk_item = {
                'page_objs': [],
                'page': [],
                'page_idx': []
            }
            cur_text_len = 0
            for i in range(0, len(page_result_list)):
                text = page_result_list[i]
                if cur_text_len + len(text) > self._max_single_chunk_size and len(chunk_item['page']):
                    chunk_list.append(copy.deepcopy(chunk_item))
                    chunk_item = {
                        'page_objs': [],
                        'page': [],
                        'page_idx': []
                    }
                    cur_text_len = 0
                chunk_item['page'].append(text)
                chunk_item['page_idx'].append(i)
                cur_text_len += len(text)
            if len(chunk_item['page']):
                chunk_list.append(copy.deepcopy(chunk_item))
            self._chunk_list = chunk_list
        else:
            logger.info('Spire extract judged readable — skip OCR.')

        return ''

    def __init__(self, **kwargs):
        super(SpirePdfDoc, self).__init__(**kwargs)
        self.ocr_util = OCRUtil()

    def load(self,doc_path):
        self._doc_path = doc_path
        self._doc_name = os.path.basename(doc_path)
        self._doc = PdfDocument()
        self._doc.LoadFromFile(self._doc_path)
        self._chunk_list = self.init_chunks()

    def init_chunks(self):
        chunk_list = []
        extract_options = PdfTextExtractOptions()
        extract_options.IsExtractAllText = True
        chunk_item = {
            'page_objs': [],
            'page': [],
            'page_idx': []
        }
        cur_text_len = 0
        for i in range(0, self._doc.Pages.Count):
            page = self._doc.Pages[i]
            text = PdfTextExtractor(page).ExtractText(extract_options)
            if cur_text_len + len(text) > self._max_single_chunk_size and len(chunk_item['page_objs']):
                chunk_list.append(chunk_item)
                chunk_item = {
                    'page_objs': [],
                    'page': [],
                    'page_idx': []
                }
                cur_text_len = 0
            chunk_item['page_objs'].append(page)
            chunk_item['page'].append(text)
            chunk_item['page_idx'].append(i)
            cur_text_len += len(text)
        if len(chunk_item['page_objs']):
            chunk_list.append(chunk_item)
        return chunk_list

    def get_chunk_item(self, chunk_id):
        chunk = self._chunk_list[chunk_id]
        return '\n'.join([page for page in chunk['page']])

    def get_chunk_info(self, chunk_id):
        chunk = self._chunk_list[chunk_id]
        from_location = f'[页面{chunk["page_idx"][0] + 1}]'
        to_location = f'[页面{chunk["page_idx"][-1] + 1}]'
        chunk_content_tips = '[' + chunk['page'][0][:20] + ']...到...[' + chunk['page'][-1][-20:] + ']'
        return f'文件块id: {chunk_id + 1}\n文件块位置: 从{from_location}到{to_location}\n文件块简述: {chunk_content_tips}\n'

    def format_comment_author(self, comment):
        return '{}|{}'.format(str(comment['id']), comment['key_points'])

    def add_text_markup_anno(self, page_obj, text_fragment, author, review_suggest):
        # 遍历文本边界
        for j in range(len(text_fragment.Bounds)):
            # 获取边界
            rect = text_fragment.Bounds[j]
            # 创建文本标记注释
            annotation = PdfTextMarkupAnnotation(author, review_suggest, rect)
            annotation.Name = author
            annotation.TextMarkupColor = PdfRGBColor(Color.get_Green())
            # 将注释添加到注释集合中
            page_obj.AnnotationsWidget.Add(annotation)

    # TODO 好像注释删不掉
    def add_poly_line_anno(self, page_obj, text_fragment, author, review_suggest):
        # 获取文本边界
        rect = text_fragment.Bounds[0]
        # 获取文本边界的坐标以添加注释
        left = rect.Left
        top = rect.Top
        right = rect.Right
        bottom = rect.Bottom
        polyLineAnnotation = PdfPolyLineAnnotation(page_obj,
                                                   [PointF(left, top), PointF(right, top), PointF(right, bottom),
                                                    PointF(left, bottom), PointF(left, top)])
        polyLineAnnotation.Name = author
        # 自定义注释内容
        polyLineAnnotation.Text = review_suggest
        # 将注释添加到文档的注释集合中
        page_obj.AnnotationsWidget.Add(polyLineAnnotation)

    # TODO 批注好像删不掉
    def add_popup_anno(self, page_obj, text_fragment, author, review_suggest):
        # 获取文本边界
        rect = text_fragment.Bounds[0]
        # 获取文本边界的坐标以添加注释
        right = rect.Right
        top = rect.Top
        # 创建注释
        rectangle = RectangleF(right, top - 15., 15., 15.)
        popupAnnotation = PdfPopupAnnotation(rectangle)
        # 自定义注释内容
        popupAnnotation.Name = author
        popupAnnotation.Text = review_suggest
        # 设置注释的图标和颜色
        popupAnnotation.Icon = PdfPopupIcon.Comment
        popupAnnotation.Color = PdfRGBColor(Color.get_Red())
        # 将注释添加到文档的注释集合中
        page_obj.AnnotationsWidget.Add(popupAnnotation)

    def add_anno(self, chunk, find_key, author, review_suggest):
        accurate_find = False
        for i in range(0, len(chunk['page_objs'])):
            page_obj = chunk['page_objs'][i]
            finder = PdfTextFinder(page_obj)
            finder.Options.Parameter = TextFindParameter.WholeWord
            fragments = finder.Find(find_key)
            # 未精确匹配
            if fragments is None or len(fragments) == 0:
                pass
            else:
                textFragment = fragments[0]
                # logger.info(f'[add_anno markup] {find_key} {author} {review_suggest}')
                # textFragment.HighLight(Color.get_Yellow())
                # self.add_text_markup_anno(page_obj, textFragment, author, review_suggest)
                # self.add_poly_line_anno(page_obj, textFragment, author, review_suggest)
                self.add_popup_anno(page_obj, textFragment, author, review_suggest)
                accurate_find = True
        return accurate_find

    def add_chunk_comment(self, chunk_id, comments):
        # logger.info(f'add_chunk_comment: {chunk_id} {comments}')
        chunk = self._chunk_list[chunk_id]
        for comment in comments:
            if comment['result'] != '不合格':
                continue
            review_suggest = comment.get('suggest', '')
            author = self.format_comment_author(comment)
            find_key = comment['original_text'].strip() if comment['original_text'].strip() else comment['key_points']
            # 先进行精确查找
            accurate_find = self.add_anno(chunk, find_key, author, review_suggest)
            # 模糊查找
            if not accurate_find:
                sub_chunk_texts = []
                for page in chunk['page']:
                    sub_chunk_texts.extend([item.strip() for item in page.split('\n') if len(item.strip())])
                top_match = process.extract(find_key, sub_chunk_texts, limit=1)
                sub_chunk_text, score, sub_chunk_idx = top_match[0]
                # logger.info(f'add_chunk_comment find_key:{find_key} fuzz match: {sub_chunk_text}')
                # 在进行模糊查找
                self.add_anno(chunk, sub_chunk_text, author, review_suggest)

    def print_all_anno(self):
        for j in range(0, self._doc.Pages.Count):
            page_obj = self._doc.Pages.get_Item(j)
            for i in range(page_obj.AnnotationsWidget.Count):
                anno = page_obj.AnnotationsWidget.get_Item(i)
                name = ''
                try:
                    name = anno.Name
                except:
                    name = ''
                logger.info(f'anno批注:{j + 1}页 第{i + 1}个 [{name}]:[{anno.Text}]')
            # page_obj.AnnotationsWidget.Clear()

    def find_anno_idx(self, page_obj, author):
        for i in range(page_obj.AnnotationsWidget.Count):
            anno = page_obj.AnnotationsWidget.get_Item(i)
            anno_name = ''
            try:
                anno_name = anno.Name
            except Exception as e:
                pass
            # logger.info(f'{type(anno)} {anno_name}')
            if anno_name == author:
                # logger.info(f'find_anno_idx:{author} at {i}')
                return i
        return -1

    def edit_chunk_comment(self, comments):
        for comment in comments:
            review_answer = comment['result']
            if review_answer == '不合格':
                idx, page, delete_comment = self.delete_single_chunk_comment_once(comment)
                # 不存在存在删除项，新增
                if idx == -1:
                    self.add_chunk_comment(comment['chunk_id'], [comment])
                # 存在删除项，编辑
                else:
                    review_suggest = comment.get('suggest', '')
                    all_pages = []
                    all_delete_comment = []
                    all_idx = []
                    while idx != -1:
                        delete_comment.Text = review_suggest
                        all_pages.append(page)
                        all_delete_comment.append(delete_comment)
                        all_idx.append(idx)
                        idx, page, delete_comment = self.delete_single_chunk_comment_once(comment)
                    for page, delete_comment, idx in zip(all_pages, all_delete_comment, all_idx):
                        page.AnnotationsWidget.Insert(idx, delete_comment)
            # 合格则删除
            else:
                self.delete_single_chunk_comment(comment)

    def delete_chunk_comment(self, comments):
        for comment in comments:
            self.delete_single_chunk_comment(comment)

    def delete_single_chunk_comment(self, comment):
        while True:
            idx, _, _ = self.delete_single_chunk_comment_once(comment)
            if idx == -1:
                break

    # 删除一次批注
    def delete_single_chunk_comment_once(self, comment):
        author = self.format_comment_author(comment)
        chunk_id = comment['chunk_id']
        chunk = self._chunk_list[chunk_id]
        for page_obj in chunk['page_objs']:
            idx = self.find_anno_idx(page_obj, author)
            if idx != -1:
                # logger.info(f'delete_single_chunk_comment_once:{author} {comment}')
                delete_comment = page_obj.AnnotationsWidget.get_Item(idx)
                page_obj.AnnotationsWidget.RemoveAt(idx)
                # logger.info(page_obj.AnnotationsWidget.Contains(delete_comment))
                return idx, page_obj, delete_comment
        return -1, None, None

    def get_chunk_id_list(self, step=1):
        return [idx for idx in range(0, self.get_chunk_num(), step)]

    def get_chunk_num(self):
        return len(self._chunk_list)

    def get_all_text(self):
        return '\n'.join([page for chunk in self._chunk_list for page in chunk['page']])

    def to_file(self, path, **kwargs):
        self._doc.SaveToFile(path)

    def re_comment_by_memory(self, comments):
        self._doc.Close()
        self.reset()
        for comment in comments:
            self.add_chunk_comment(comment['chunk_id'], [comment])

    def release(self):
        self._doc.Close()
        super().release()


if __name__ == '__main__':
    tool = SpirePdfDoc()
    tool.load('/home/ccran/lufa-contract/tmp/财03 2023年国科京东方评估合同.pdf')
    print('', tool.get_all_text())
    print(asyncio.run(tool.get_from_ocr()))
    print(tool.get_chunk_id_list())
