# SPDX-FileCopyrightText: 2025 QinHan
# SPDX-License-Identifier: MPL-2.0
import asyncio
from collections import defaultdict
from copy import deepcopy
from dataclasses import dataclass
from io import BytesIO
from typing import Self, Literal, List, Dict, Any, Tuple, Optional

import docx
from docx.document import Document as DocumentObject
from docx.opc.part import Part
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.oxml.text.run import CT_R
from docx.section import _Header, _Footer
from docx.text.paragraph import Paragraph
from docx.text.run import Run
from docx.table import _Cell, Table

from docutranslate.agents.segments_agent import SegmentsTranslateAgentConfig, SegmentsTranslateAgent
from docutranslate.ir.document import Document
from docutranslate.translator.ai_translator.base import AiTranslatorConfig, AiTranslator

# ---------------- 辅助函数 ----------------



def is_image_run(run: Run) -> bool:
    """检查一个 Run 是否包含图片。"""
    xml = getattr(run.element, 'xml', '')
    return '<w:drawing' in xml or '<w:pict' in xml


def is_instr_text_run(run: Run) -> bool:
    """
    检查 Run 是否包含域指令文本 (w:instrText)。
    目录(TOC)、页码等功能的指令代码存储在此标签中。
    必须跳过这些 Run，否则写入 text 会破坏域结构。
    """
    return run.element.find(qn('w:instrText')) is not None


# ---------------- 配置类 ----------------
@dataclass
class DocxTranslatorConfig(AiTranslatorConfig):
    insert_mode: Literal["replace", "append", "prepend"] = "replace"
    separator: str = "\n"
    office_password: Optional[str] = None


# ---------------- 主类 ----------------
class DocxTranslator(AiTranslator):
    """
    一个基于高级结构化解析的 .docx 文件翻译器。
    它能高精度保留样式，并正确处理正文、表格、页眉/脚、脚注/尾注、超链接和目录(TOC)等复杂元素。

    [v6.4 - 形状文本翻译]
    - 新增对形状（Shapes）内文本的提取与翻译支持。现在可以正确翻译标注、流程图等图形元素中的文字。
    - 改进了核心遍历逻辑，使其能够深入 <w:drawing> 元素，解析内部的 <w:txbxContent>（文本框内容），
      同时保持对无文本的普通图片的原有处理方式，实现鲁棒性和兼容性。

    [v6.3 - 上下文感知格式切分]
    - 根本性地改进了文本切分逻辑，解决了因带格式的空格（如下划线空格）导致连续短语被错误拆分的问题。
    - 新逻辑不再将任何带格式的空格视为固定的边界，而是通过比较相邻文本片段的“显著格式”是否一致来决定是否合并。
    - 只要格式（如下划线、高亮等）保持不变，即使中间有空格，文本也会被视为一个连续的翻译单元，
      极大地提升了对标题、合同条款等格式化文本的翻译准确性和流畅性。

    [v6.2 - 精确格式保留]
    - 根据用户反馈，精确定义了构成“重要格式”的样式范围。
    - 现在仅在空白 Run 包含下划线、背景色、边框等显著视觉样式时才将其视为边界。
    - 忽略了加粗、倾斜、字体/字号变化等对空格本身无视觉效果或不重要的样式，
      避免了因此对翻译句子造成不必要的切分，提升了翻译质量。

    [v6.1 - 格式保留修复版]
    - 修复了因合并 Run 导致下划线等格式在翻译后丢失的问题。
    - 通过引入 is_styled_whitespace_run 检查，将仅包含空格但带有样式的 Run（如下划线空格）
      视为与图片类似的不可翻译边界。

    [v6.0 - 语义切分重构版]
    - 重构核心逻辑，不再跳过域结果，而是将其作为语义边界来切分文本，增强了鲁棒性。
    """

    def __init__(self, config: DocxTranslatorConfig):
        super().__init__(config=config)
        self.chunk_size = config.chunk_size
        self.translate_agent = None
        glossary_dict = self.glossary.glossary_dict if self.glossary else None
        if not self.skip_translate:
            agent_config = SegmentsTranslateAgentConfig(
                custom_prompt=config.custom_prompt, to_lang=config.to_lang, base_url=config.base_url,
                api_key=config.api_key, model_id=config.model_id, temperature=config.temperature,
                thinking=config.thinking, concurrent=config.concurrent, timeout=config.timeout,
                logger=self.logger, glossary_dict=glossary_dict, retry=config.retry,
                system_proxy_enable=config.system_proxy_enable, force_json=config.force_json,
                rpm=config.rpm,
                tpm=config.tpm,
                provider=config.provider,
            )
            self.translate_agent = SegmentsTranslateAgent(agent_config)
        self.insert_mode = config.insert_mode
        self.separator = config.separator
        self.office_password = config.office_password

    def _decrypt_if_needed(self, content: bytes) -> bytes:
        """如果文件加密则解密，否则返回原内容。"""
        try:
            import msoffcrypto
            from io import BytesIO as BIO
            file_stream = BIO(content)
            try:
                office_file = msoffcrypto.OfficeFile(file_stream)
                if office_file.is_encrypted():
                    if not self.office_password:
                        raise ValueError("此DOCX文件已加密，但未提供密码。")
                    decrypted = BIO()
                    office_file.load_key(password=self.office_password)
                    office_file.decrypt(decrypted)
                    return decrypted.getvalue()
                return content
            finally:
                file_stream.close()
        except ImportError:
            return content

    @staticmethod
    def _run_format_key(run: Run):
        """生成 Run 的格式签名，用于合并相同格式的 Run。"""
        return (
            run.bold,
            run.italic,
            run.underline,
            run.font.size if run.font.size else None,
            str(run.font.color.rgb) if run.font.color and run.font.color.rgb else None,
            run.font.name or None,
        )

    def _merge_adjacent_runs(self, runs: List[Run]) -> List[Run]:
        """
        合并相邻的、格式完全相同的 Run，同时累积文本到第一个 Run。
        解决 Word 因修订历史/变更追踪产生的微观 Run 碎片问题（单个字符一个 Run）。
        """
        if len(runs) <= 1:
            return runs

        merged = []
        group_start = 0
        for i in range(1, len(runs)):
            if self._run_format_key(runs[i]) != self._run_format_key(runs[group_start]):
                # Format boundary: finalize the current group
                if i - group_start > 1:
                    # Merge: accumulate all text into first run, delete the rest
                    runs[group_start].text = "".join(r.text for r in runs[group_start:i])
                    for r in runs[group_start + 1:i]:
                        self._remove_run_element(r)
                merged.append(runs[group_start])
                group_start = i

        # Final group
        if len(runs) - group_start > 1:
            runs[group_start].text = "".join(r.text for r in runs[group_start:])
            for r in runs[group_start + 1:]:
                self._remove_run_element(r)
        merged.append(runs[group_start])

        return merged

    def _process_paragraph(self, para: Paragraph, elements: List[Dict[str, Any]], texts: List[str],
                         top_level_para: Paragraph = None):
        """
        段落级翻译处理：收集所有文本 Run → 合并相邻同格式 Run → 整段翻译 → 按比例分配。
        """
        if top_level_para is None:
            top_level_para = para

        text_runs = []
        for run in para.runs:
            if is_image_run(run) or is_instr_text_run(run):
                continue
            if not run.text.strip():
                continue
            text_runs.append(run)

        if text_runs:
            # Merge adjacent runs with identical formatting to reduce fragmentation
            text_runs = self._merge_adjacent_runs(text_runs)
            full_text = "".join(r.text for r in text_runs)
            if full_text.strip():
                elements.append({
                    "type": "text_runs",
                    "runs": list(text_runs),
                    "paragraph": para,
                    "top_level_paragraph": top_level_para
                })
                texts.append(full_text)

    # ---------------------- 修改结束 ----------------------

    def _process_body_elements(self, parent_element, container, elements: List[Dict[str, Any]], texts: List[str]):
        """ 遍历一个容器内的所有顶级元素（段落、表格、内容控件等） """
        for child_element in parent_element:
            if child_element.tag.endswith('p'):
                self._process_paragraph(Paragraph(child_element, container), elements, texts)
            elif child_element.tag.endswith('tbl'):
                table = Table(child_element, container)
                for row in table.rows:
                    for cell in row.cells:
                        self._traverse_container(cell, elements, texts)
            elif child_element.tag.endswith('sdt'):
                sdt_content = child_element.find(qn('w:sdtContent'))
                if sdt_content is not None:
                    self._process_body_elements(sdt_content, container, elements, texts)

    def _traverse_container(self, container: Any, elements: List[Dict[str, Any]], texts: List[str]):
        """
        [核心导航员 v2.0 - 增强版] 健壮地遍历任何文本容器，
        特别为具有额外嵌套层的 Part (如脚注/尾注) 提供了专门处理逻辑。
        """
        if container is None:
            return

        parent_element = None
        if isinstance(container, (DocumentObject, Part)):
            parent_element = container.element.body if hasattr(container.element, 'body') else container.element
        elif isinstance(container, (_Cell, _Header, _Footer)):
            parent_element = container._element
        else:
            self.logger.warning(f"跳过未知类型的容器: {type(container)}")
            return

        if parent_element is not None and parent_element.tag in [qn('w:footnotes'), qn('w:endnotes')]:
            for note_element in parent_element:
                self._process_body_elements(note_element, container, elements, texts)
        elif parent_element is not None:
            self._process_body_elements(parent_element, container, elements, texts)

    def _pre_translate(self, document: Document) -> Tuple[DocumentObject, List[Dict[str, Any]], List[str]]:
        content = self._decrypt_if_needed(document.content)
        doc = docx.Document(BytesIO(content))
        elements, texts = [], []

        self._traverse_container(doc, elements, texts)

        for section in doc.sections:
            self._traverse_container(section.header, elements, texts)
            self._traverse_container(section.first_page_header, elements, texts)
            self._traverse_container(section.even_page_header, elements, texts)
            self._traverse_container(section.footer, elements, texts)
            self._traverse_container(section.first_page_footer, elements, texts)
            self._traverse_container(section.even_page_footer, elements, texts)

        if hasattr(doc.part, 'footnotes_part') and doc.part.footnotes_part is not None:
            self._traverse_container(doc.part.footnotes_part, elements, texts)
        if hasattr(doc.part, 'endnotes_part') and doc.part.endnotes_part is not None:
            self._traverse_container(doc.part.endnotes_part, elements, texts)

        return doc, elements, texts

    def _apply_translation(self, element_info: Dict[str, Any], final_text: str):
        if element_info["type"] == "text_runs":
            runs = element_info["runs"]
            if not runs: return

            # Filter to runs that are still attached to the document
            valid_runs = []
            for run in runs:
                if run.element.getparent() is not None:
                    run._parent = element_info["paragraph"]
                    valid_runs.append(run)

            if not valid_runs:
                self.logger.warning(f"无法应用翻译 '{final_text}'，因为找不到有效的run。")
                return

            if len(valid_runs) == 1:
                # Single run: just write the translation
                valid_runs[0].text = final_text
                return

            # Multiple runs: proportionally distribute translated text to preserve formatting
            orig_lengths = [len(r.text) for r in valid_runs]
            total_orig = sum(orig_lengths)
            final_len = len(final_text)

            if total_orig == 0:
                valid_runs[0].text = final_text
                for run in valid_runs[1:]:
                    self._remove_run_element(run)
                return

            # Distribute characters proportionally
            char_pos = 0
            for i, run in enumerate(valid_runs):
                if i == len(valid_runs) - 1:
                    # Last run gets all remaining text
                    run.text = final_text[char_pos:]
                else:
                    ratio = orig_lengths[i] / total_orig
                    run_char_count = max(1, round(final_len * ratio))
                    run_char_count = min(run_char_count, final_len - char_pos - (len(valid_runs) - i - 1))
                    if run_char_count <= 0:
                        # Remove runs that would get zero characters
                        self._remove_run_element(run)
                        continue
                    run.text = final_text[char_pos:char_pos + run_char_count]
                    char_pos += run_char_count

    def _remove_run_element(self, run) -> None:
        """Safely remove a run element from its parent."""
        parent_element = run.element.getparent()
        if parent_element is not None:
            try:
                parent_element.remove(run.element)
            except ValueError:
                self.logger.debug(f"尝试删除一个不存在的run元素。这通常是安全的。")

    # ---------- FIX START: 新增用于清理副本段落的辅助方法 ----------
    def _prune_unwanted_elements_from_copy(self, p_element: OxmlElement):
        """
        从复制的段落元素中移除包含图片或页码字段的 Run。
        这可以防止在“append”和“prepend”模式下出现重复。
        """
        runs_to_remove = []
        runs = p_element.findall(qn('w:r'))

        i = 0
        while i < len(runs):
            run_element = runs[i]

            # 检查图片
            if run_element.find(qn('w:drawing')) is not None or run_element.find(qn('w:pict')) is not None:
                runs_to_remove.append(run_element)
                i += 1
                continue

            # 检查页码字段
            fldChar = run_element.find(qn('w:fldChar'))
            if fldChar is not None and fldChar.get(qn('w:fldCharType')) == 'begin':
                is_page_field = False
                field_end_index = -1

                # 向前查找以确定是否是页码字段
                for j in range(i + 1, len(runs)):
                    next_run = runs[j]
                    instrText = next_run.find(qn('w:instrText'))
                    if instrText is not None and instrText.text is not None:
                        text = instrText.text.strip().upper()
                        if 'PAGE' in text or 'NUMPAGES' in text:
                            is_page_field = True
                        break

                    next_fldChar = next_run.find(qn('w:fldChar'))
                    if next_fldChar is not None and next_fldChar.get(qn('w:fldCharType')) == 'begin':
                        break

                if is_page_field:
                    # 如果是页码字段，则找到其结束标记并标记整个字段的 runs
                    field_runs_to_remove = [run_element]
                    end_found = False
                    for j in range(i + 1, len(runs)):
                        field_run = runs[j]
                        field_runs_to_remove.append(field_run)
                        end_fldChar = field_run.find(qn('w:fldChar'))
                        if end_fldChar is not None and end_fldChar.get(qn('w:fldCharType')) == 'end':
                            end_found = True
                            field_end_index = j
                            break

                    if end_found:
                        runs_to_remove.extend(field_runs_to_remove)
                        i = field_end_index + 1
                        continue
            i += 1

        # 从 XML 树中实际移除被标记的 runs
        for run_to_remove in runs_to_remove:
            if run_to_remove.getparent() is not None:
                p_element.remove(run_to_remove)

    # ---------- FIX END ----------

    def _after_translate(self, doc: DocumentObject, elements: List[Dict[str, Any]], translated: List[str],
                         originals: List[str]) -> bytes:
        if len(elements) != len(translated):
            self.logger.error(
                f"翻译数量不匹配！原文: {len(originals)}, 译文: {len(translated)}. 将只处理公共部分。")
            min_len = min(len(elements), len(translated), len(originals))
            elements, translated, originals = elements[:min_len], translated[:min_len], originals[:min_len]

        if self.insert_mode == "replace":
            for info, trans in zip(elements, translated):
                self._apply_translation(info, trans)
        else:
            paragraph_segments = defaultdict(list)
            # [FIX] 按顶级段落对所有片段进行分组，以确保形状等嵌套内容与主段落一起处理
            for i, info in enumerate(elements):
                top_level_paragraph = info["top_level_paragraph"]
                paragraph_segments[id(top_level_paragraph._p)].append({
                    "index": i,
                    "translation": translated[i],
                    "paragraph_obj": top_level_paragraph
                })

            for para_id, segments_for_this_para in paragraph_segments.items():
                # 从该组的第一个片段中获取唯一的顶级段落对象
                top_level_paragraph_orig = segments_for_this_para[0]["paragraph_obj"]
                p_element_orig = top_level_paragraph_orig._p

                # 创建顶级段落的深拷贝，所有翻译将应用于此副本
                translated_p_element = deepcopy(p_element_orig)
                self._prune_unwanted_elements_from_copy(translated_p_element)
                top_level_paragraph_copy = Paragraph(translated_p_element, top_level_paragraph_orig._parent)

                # [FIX] 创建从原始元素到复制元素的映射，包括嵌套的段落和所有runs
                # 1. 嵌套段落映射
                para_map = {id(p_element_orig): top_level_paragraph_copy}
                orig_nested_ps = p_element_orig.iter(qn('w:p'))
                copy_nested_ps = translated_p_element.iter(qn('w:p'))
                for o, c in zip(orig_nested_ps, copy_nested_ps):
                    para_map[id(o)] = Paragraph(c, top_level_paragraph_copy)

                # 2. Run元素映射
                run_element_map = {
                    id(orig_r): copied_r
                    for orig_r, copied_r in zip(p_element_orig.iter(qn('w:r')), translated_p_element.iter(qn('w:r')))
                }

                for seg_info in segments_for_this_para:
                    element_index = seg_info["index"]
                    translation = seg_info["translation"]
                    original_element_info = elements[element_index]

                    # [FIX] 从映射中查找当前片段对应的、位于副本中的父段落对象
                    original_para_id = id(original_element_info["paragraph"]._p)
                    translated_paragraph_obj = para_map.get(original_para_id)

                    if not translated_paragraph_obj:
                        self.logger.warning(
                            f"无法在段落副本中找到对应的嵌套段落，跳过翻译应用: '{translation}'")
                        continue

                    # [FIX] 查找位于副本中的Run对象
                    runs_from_copy = []
                    for r in original_element_info["runs"]:
                        copied_r_element = run_element_map.get(id(r.element))
                        if copied_r_element is not None:
                            # 使用正确的、位于副本中的父段落对象来创建Run
                            new_run = Run(copied_r_element, translated_paragraph_obj)
                            runs_from_copy.append(new_run)

                    if not runs_from_copy:
                        # self.logger.warning("在副本段落中找不到对应的Runs，跳过翻译应用。")
                        continue

                    translated_element_info = {
                        "type": "text_runs",
                        "runs": runs_from_copy,
                        "paragraph": translated_paragraph_obj
                    }
                    self._apply_translation(translated_element_info, translation)

                # --- 分隔符和插入逻辑 (保持不变) ---
                separator_p_element = None
                if self.separator:
                    separator_p_element = OxmlElement('w:p')
                    run_element = OxmlElement('w:r')
                    lines = self.separator.split('\n')
                    for i, line in enumerate(lines):
                        text_element = OxmlElement('w:t')
                        text_element.set(qn('xml:space'), 'preserve')
                        text_element.text = line
                        run_element.append(text_element)
                        if i < len(lines) - 1:
                            run_element.append(OxmlElement('w:br'))
                    separator_p_element.append(run_element)

                if self.insert_mode == "append":
                    current_element = p_element_orig
                    if separator_p_element is not None:
                        current_element.addnext(separator_p_element)
                        current_element = separator_p_element
                    current_element.addnext(translated_p_element)
                elif self.insert_mode == "prepend":
                    p_element_orig.addprevious(translated_p_element)
                    if separator_p_element is not None:
                        translated_p_element.addnext(separator_p_element)

        doc_output_stream = BytesIO()
        doc.save(doc_output_stream)
        return doc_output_stream.getvalue()

    def translate(self, document: Document) -> Self:
        doc, elements, originals = self._pre_translate(document)
        if not originals:
            self.logger.info("\n文档中未找到可翻译的文本内容。")
            document.content = self._after_translate(doc, elements, [], [])
            return self

        if self.glossary_agent:
            # 1. 获取增量
            glossary_dict_gen = self.glossary_agent.send_segments(originals, self.chunk_size)

            # 2. 在 Translator 层统一合并 (SSOT)
            if self.glossary:
                self.glossary.update(glossary_dict_gen)

            # 3. 将合并后的【完整字典】传给 Agent
            if self.translate_agent and self.glossary:
                self.translate_agent.update_glossary_dict(self.glossary.glossary_dict)

        translated = self.translate_agent.send_segments(originals,
                                                        self.chunk_size) if self.translate_agent else originals
        document.content = self._after_translate(doc, elements, translated, originals)
        return self

    async def translate_async(self, document: Document) -> Self:
        doc, elements, originals = await asyncio.to_thread(self._pre_translate, document)
        if not originals:
            self.logger.info("\n文档中未找到可翻译的文本内容。")
            document.content = await asyncio.to_thread(self._after_translate, doc, elements, [], [])
            return self

        if self.glossary_agent:
            # 1. 获取增量
            glossary_dict_gen = await self.glossary_agent.send_segments_async(originals, self.chunk_size)

            # 2. 在 Translator 层统一合并 (SSOT)
            if self.glossary:
                self.glossary.update(glossary_dict_gen)

            # 3. 将合并后的【完整字典】传给 Agent
            if self.translate_agent and self.glossary:
                self.translate_agent.update_glossary_dict(self.glossary.glossary_dict)

        translated = await self.translate_agent.send_segments_async(originals,
                                                                    self.chunk_size) if self.translate_agent else originals
        document.content = await asyncio.to_thread(self._after_translate, doc, elements, translated, originals)
        return self