import asyncio
from io import BytesIO
import unittest
from zipfile import ZipFile

import docx
from docx.oxml import OxmlElement, parse_xml
from docx.oxml.ns import nsdecls, qn
from docx.shared import Inches
from lxml import etree

from docutranslate.ir.document import Document
from docutranslate.translator.ai_translator.docx_translator import (
    DocxTranslator,
    DocxTranslatorConfig,
)


NS = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}


def _docx_bytes(document):
    stream = BytesIO()
    document.save(stream)
    return stream.getvalue()


def _text_values(docx_bytes):
    with ZipFile(BytesIO(docx_bytes)) as zf:
        tree = etree.fromstring(zf.read("word/document.xml"))
    return tree.xpath("//w:t/text()", namespaces=NS)


def _comments_text_values(docx_bytes):
    with ZipFile(BytesIO(docx_bytes)) as zf:
        if "word/comments.xml" not in zf.namelist():
            return []
        tree = etree.fromstring(zf.read("word/comments.xml"))
    return tree.xpath("//w:t/text()", namespaces=NS)


def _header_text_values(docx_bytes):
    values = []
    with ZipFile(BytesIO(docx_bytes)) as zf:
        for name in zf.namelist():
            if name.startswith("word/header") and name.endswith(".xml"):
                tree = etree.fromstring(zf.read(name))
                values.extend(tree.xpath("//w:t/text()", namespaces=NS))
    return values


def _paragraph_and_table_counts(docx_bytes):
    with ZipFile(BytesIO(docx_bytes)) as zf:
        tree = etree.fromstring(zf.read("word/document.xml"))
    return (
        len(tree.xpath("//w:p", namespaces=NS)),
        len(tree.xpath("//w:tbl", namespaces=NS)),
    )


def _section_type_counts(docx_bytes):
    with ZipFile(BytesIO(docx_bytes)) as zf:
        tree = etree.fromstring(zf.read("word/document.xml"))
    counts = {}
    for sect_pr in tree.xpath("//w:sectPr", namespaces=NS):
        value = sect_pr.xpath("./w:type/@w:val", namespaces=NS)
        key = value[0] if value else "nextPage(default)"
        counts[key] = counts.get(key, 0) + 1
    return counts


def _add_toc_paragraph(document):
    paragraph = document.add_paragraph()
    p_style = OxmlElement("w:pStyle")
    p_style.set(qn("w:val"), "TOC1")
    paragraph._p.get_or_add_pPr().append(p_style)

    hyperlink = OxmlElement("w:hyperlink")
    hyperlink.set(qn("w:anchor"), "_TocFixture")

    title_run = OxmlElement("w:r")
    title_text = OxmlElement("w:t")
    title_text.text = "Purpose"
    title_run.append(title_text)

    tab_run = OxmlElement("w:r")
    tab_run.append(OxmlElement("w:tab"))

    page_run = OxmlElement("w:r")
    page_text = OxmlElement("w:t")
    page_text.text = "3"
    page_run.append(page_text)

    hyperlink.extend([title_run, tab_run, page_run])
    paragraph._p.append(hyperlink)


def _add_internal_hyperlink_paragraph(document):
    paragraph = document.add_paragraph()
    paragraph.add_run("Read ")
    hyperlink = OxmlElement("w:hyperlink")
    hyperlink.set(qn("w:anchor"), "_Manual")
    link_run = OxmlElement("w:r")
    link_text = OxmlElement("w:t")
    link_text.text = "manual"
    link_run.append(link_text)
    hyperlink.append(link_run)
    paragraph._p.append(hyperlink)
    paragraph.add_run(" now")


def _add_textbox_paragraph(document):
    paragraph = document.add_paragraph()
    paragraph.add_run("Diagram: ")
    textbox_run = parse_xml(
        f"""
        <w:r {nsdecls("w", "wp", "a")} xmlns:wps="http://schemas.microsoft.com/office/word/2010/wordprocessingShape">
          <w:drawing>
            <wp:inline>
              <a:graphic>
                <a:graphicData uri="http://schemas.microsoft.com/office/word/2010/wordprocessingShape">
                  <wps:wsp>
                    <wps:txbx>
                      <w:txbxContent>
                        <w:p>
                          <w:r><w:t>Flow label</w:t></w:r>
                        </w:p>
                      </w:txbxContent>
                    </wps:txbx>
                  </wps:wsp>
                </a:graphicData>
              </a:graphic>
            </wp:inline>
          </w:drawing>
        </w:r>
        """
    )
    paragraph._p.append(textbox_run)


def _add_inline_sdt(paragraph, text):
    sdt = OxmlElement("w:sdt")
    sdt_content = OxmlElement("w:sdtContent")
    run = OxmlElement("w:r")
    text_element = OxmlElement("w:t")
    text_element.text = text
    run.append(text_element)
    sdt_content.append(run)
    sdt.append(sdt_content)
    paragraph._p.append(sdt)


def _add_block_sdt(document, text):
    sdt = OxmlElement("w:sdt")
    sdt_content = OxmlElement("w:sdtContent")
    paragraph = OxmlElement("w:p")
    run = OxmlElement("w:r")
    text_element = OxmlElement("w:t")
    text_element.text = text
    run.append(text_element)
    paragraph.append(run)
    sdt_content.append(paragraph)
    sdt.append(sdt_content)
    document._body._element.append(sdt)


def _add_page_field_paragraph(document):
    paragraph = document.add_paragraph("Page ")
    for fld_type in ("begin", "separate", "end"):
        run = OxmlElement("w:r")
        if fld_type == "begin":
            fld = OxmlElement("w:fldChar")
            fld.set(qn("w:fldCharType"), "begin")
            run.append(fld)
            instr_run = OxmlElement("w:r")
            instr = OxmlElement("w:instrText")
            instr.text = " PAGE "
            instr_run.append(instr)
            paragraph._p.append(run)
            paragraph._p.append(instr_run)
            continue
        if fld_type == "separate":
            fld = OxmlElement("w:fldChar")
            fld.set(qn("w:fldCharType"), "separate")
            run.append(fld)
            cached_run = OxmlElement("w:r")
            cached_text = OxmlElement("w:t")
            cached_text.text = "1"
            cached_run.append(cached_text)
            paragraph._p.append(run)
            paragraph._p.append(cached_run)
            continue
        fld = OxmlElement("w:fldChar")
        fld.set(qn("w:fldCharType"), "end")
        run.append(fld)
        paragraph._p.append(run)


def _add_tracked_changes_paragraph(document):
    paragraph = document.add_paragraph("Revision ")

    deleted = OxmlElement("w:del")
    deleted.set(qn("w:id"), "1")
    deleted_run = OxmlElement("w:r")
    deleted_text = OxmlElement("w:delText")
    deleted_text.text = "old text"
    deleted_run.append(deleted_text)
    deleted.append(deleted_run)
    paragraph._p.append(deleted)

    inserted = OxmlElement("w:ins")
    inserted.set(qn("w:id"), "2")
    inserted_run = OxmlElement("w:r")
    inserted_text = OxmlElement("w:t")
    inserted_text.text = "new text"
    inserted_run.append(inserted_text)
    inserted.append(inserted_run)
    paragraph._p.append(inserted)


def _add_redundant_section_break(document):
    document.add_paragraph("Section One")
    for _ in range(2):
        section = document.add_section()
        section.start_type = docx.enum.section.WD_SECTION.NEW_PAGE
        section.top_margin = Inches(1)
        section.bottom_margin = Inches(1)
        section.left_margin = Inches(1)
        section.right_margin = Inches(1)
    document.add_paragraph("Section Two")


def _build_fixture_docx():
    document = docx.Document()
    for section in document.sections:
        section.top_margin = Inches(1)
        section.bottom_margin = Inches(1)
        section.left_margin = Inches(1)
        section.right_margin = Inches(1)
    document.add_heading("Fixture", level=1)
    document.add_paragraph("Normal paragraph")
    _add_toc_paragraph(document)
    _add_internal_hyperlink_paragraph(document)
    _add_textbox_paragraph(document)
    inline_sdt_paragraph = document.add_paragraph("Controlled ")
    _add_inline_sdt(inline_sdt_paragraph, "Approval Required")
    _add_block_sdt(document, "Block Controlled Text")
    _add_page_field_paragraph(document)
    _add_tracked_changes_paragraph(document)

    table = document.add_table(rows=2, cols=2)
    merged = table.cell(0, 0).merge(table.cell(0, 1))
    merged.text = "Merged Cell"
    table.cell(1, 0).text = "Left Cell"
    table.cell(1, 1).text = "Right Cell"
    nested_table = table.cell(1, 0).add_table(rows=1, cols=1)
    nested_table.cell(0, 0).text = "Nested Cell"

    comment_anchor = document.add_paragraph().add_run("Comment anchor")
    document.add_comment(comment_anchor, text="Review this note", author="QA")

    _add_redundant_section_break(document)
    for sect_pr in document._element.xpath(".//w:sectPr"):
        for child in list(sect_pr):
            if child.tag in {qn("w:headerReference"), qn("w:footerReference")}:
                sect_pr.remove(child)
    return _docx_bytes(document)


def _build_signature_table_docx():
    document = docx.Document()
    table = document.add_table(rows=2, cols=11)
    table.cell(0, 0).text = "日期"
    paragraph = table.cell(1, 0).paragraphs[0]
    paragraph.add_run("姓名&签名")
    paragraph.add_run("  ")
    return _docx_bytes(document)


class ScriptedAgent:
    def __init__(self, responses):
        self.responses = iter(responses)
        self.calls = []

    def send_segments(self, segments, chunk_size):
        self.calls.append(list(segments))
        return next(self.responses)

    async def send_segments_async(self, segments, chunk_size):
        return self.send_segments(segments, chunk_size)


class DocxTranslatorTest(unittest.TestCase):
    def _translator(self, skip_translate=True):
        return DocxTranslator(DocxTranslatorConfig(skip_translate=skip_translate))

    def _quality_translator(self, responses, **options):
        translator = DocxTranslator(DocxTranslatorConfig(
            api_key="dummy", base_url="http://127.0.0.1", model_id="dummy",
            **options,
        ))
        translator.translate_agent = ScriptedAgent(responses)
        return translator

    def test_only_incomplete_segments_are_retried(self):
        for use_async in (False, True):
            for incomplete in ("姓名&签名  ", "Nama & 签名", "", None):
                with self.subTest(use_async=use_async, incomplete=incomplete):
                    translator = self._quality_translator(
                        [["Tanggal", incomplete], ["Nama & Tanda Tangan"]],
                        to_lang="Bahasa Indonesia",
                    )
                    document = Document(content=_build_signature_table_docx(), suffix=".docx")
                    if use_async:
                        asyncio.run(translator.translate_async(document))
                    else:
                        translator.translate(document)
                    self.assertEqual(translator.translate_agent.calls, [
                        ["日期", "姓名&签名  "], ["姓名&签名  "],
                    ])
                    result = docx.Document(BytesIO(document.content))
                    self.assertEqual(result.tables[0].cell(1, 0).text, "Nama & Tanda Tangan")

    def test_incomplete_translation_fails_after_two_repairs_without_changing_input(self):
        for use_async in (False, True):
            with self.subTest(use_async=use_async):
                translator = self._quality_translator(
                    [["Tanggal", "姓名&签名"], ["Nama & 签名"], ["Nama & 签名"]],
                    to_lang="Indonesian",
                )
                source = _build_signature_table_docx()
                document = Document(content=source, suffix=".docx")
                with self.assertRaisesRegex(ValueError, "已定向重试 2 次"):
                    if use_async:
                        asyncio.run(translator.translate_async(document))
                    else:
                        translator.translate(document)
                self.assertEqual(document.content, source)
                self.assertEqual(len(translator.translate_agent.calls), 3)

    def test_second_repair_can_complete_translation(self):
        translator = self._quality_translator(
            [["Tanggal", "姓名&签名"], ["Nama & 签名"], ["Nama & Tanda Tangan"]],
            to_lang="印尼语",
        )
        document = Document(content=_build_signature_table_docx(), suffix=".docx")
        translator.translate(document)
        self.assertEqual(len(translator.translate_agent.calls), 3)
        self.assertIn("Nama & Tanda Tangan", "".join(_text_values(document.content)))

    def test_translation_quality_respects_target_language_and_glossary(self):
        for language in ("English", "en-US", "英文", "Indonesian", "id", "印尼语"):
            with self.subTest(language=language):
                translator = self._quality_translator([], to_lang=language)
                self.assertEqual(translator._incomplete_translation_indices(
                    ["姓名&签名"], ["Nama & 签名"]
                ), [0])
        for language in ("简体中文", "繁体中文", "Japanese"):
            with self.subTest(language=language):
                translator = self._quality_translator([], to_lang=language)
                self.assertEqual(translator._incomplete_translation_indices(
                    ["姓名&签名"], ["姓名&签名"]
                ), [])
        translator = self._quality_translator(
            [], to_lang="English", glossary_dict={"华宝": "华宝"},
        )
        self.assertEqual(translator._incomplete_translation_indices(
            ["华宝签名"], ["华宝 signature"]
        ), [])
        self.assertEqual(translator._incomplete_translation_indices(
            ["华宝签名"], ["华宝签名"]
        ), [0])

    def test_wrong_translation_count_is_not_silently_truncated(self):
        for use_async in (False, True):
            for responses in ([["Tanggal"]], [["Tanggal", "姓名&签名"], []]):
                with self.subTest(use_async=use_async, responses=responses):
                    translator = self._quality_translator(responses, to_lang="Indonesian")
                    source = _build_signature_table_docx()
                    document = Document(content=source, suffix=".docx")
                    with self.assertRaisesRegex(ValueError, "译文数量与原文不一致"):
                        if use_async:
                            asyncio.run(translator.translate_async(document))
                        else:
                            translator.translate(document)
                    self.assertEqual(document.content, source)

    def test_signature_label_after_empty_cells_is_always_extracted(self):
        input_bytes = _build_signature_table_docx()
        translator = self._translator()
        for iteration in range(50):
            with self.subTest(iteration=iteration):
                _, _, originals = translator._pre_translate(
                    Document(content=input_bytes, suffix=".docx")
                )
                self.assertEqual(originals, ["日期", "姓名&签名  "])

    def test_empty_and_merged_cells_do_not_hide_later_labels(self):
        document = docx.Document()
        table = document.add_table(rows=5, cols=11)
        table.cell(0, 0).merge(table.cell(0, 1)).text = "日期"
        table.cell(1, 0).merge(table.cell(2, 0)).text = "姓名&签名"
        table.cell(4, 0).text = "姓名&签名"
        table.cell(4, 10).text = "审批<通过>"
        input_bytes = _docx_bytes(document)

        for iteration in range(20):
            with self.subTest(iteration=iteration):
                _, _, originals = self._translator()._pre_translate(
                    Document(content=input_bytes, suffix=".docx")
                )
                self.assertEqual(
                    originals, ["日期", "姓名&签名", "姓名&签名", "审批<通过>"]
                )

    def test_signature_label_round_trip_sync_and_async(self):
        class FakeAgent:
            def send_segments(self, segments, chunk_size):
                translations = {
                    "日期": "Tanggal",
                    "姓名&签名  ": "Nama & Tanda Tangan  ",
                }
                return [translations[segment] for segment in segments]

            async def send_segments_async(self, segments, chunk_size):
                return self.send_segments(segments, chunk_size)

        input_bytes = _build_signature_table_docx()
        for use_async in (False, True):
            for mode in ("replace", "append", "prepend"):
                with self.subTest(use_async=use_async, mode=mode):
                    translator = DocxTranslator(DocxTranslatorConfig(
                        skip_translate=False, insert_mode=mode,
                        api_key="dummy", base_url="http://127.0.0.1",
                        model_id="dummy", to_lang="Indonesian",
                    ))
                    translator.translate_agent = FakeAgent()
                    document = Document(content=input_bytes, suffix=".docx")
                    if use_async:
                        asyncio.run(translator.translate_async(document))
                    else:
                        translator.translate(document)
                    result = docx.Document(BytesIO(document.content))
                    self.assertEqual(len(result.tables), 1)
                    self.assertEqual(len(result.tables[0].rows), 2)
                    self.assertEqual(len(result.tables[0].columns), 11)
                    self.assertIn("Tanggal", result.tables[0].cell(0, 0).text)
                    label = result.tables[0].cell(1, 0).text
                    self.assertIn("Nama & Tanda Tangan", label)
                    self.assertEqual("姓名&签名" in label, mode != "replace")
                    for row in result.tables[0].rows:
                        self.assertTrue(all(not cell.text for cell in row.cells[1:]))

    def test_pre_translate_extracts_docx_boundary_text(self):
        translator = self._translator()
        _, _, originals = translator._pre_translate(Document(content=_build_fixture_docx(), suffix=".docx"))

        self.assertIn("Purpose", originals)
        self.assertIn("Flow label", originals)
        self.assertIn("Controlled Approval Required", originals)
        self.assertIn("Block Controlled Text", originals)
        self.assertIn("Revision new text", originals)
        self.assertNotIn("Revision old textnew text", originals)
        self.assertTrue(all("old text" not in original for original in originals))
        self.assertIn("Merged Cell", originals)
        self.assertIn("Nested Cell", originals)
        self.assertIn("Review this note", originals)
        self.assertNotIn("Purpose3", originals)
        self.assertNotIn("PAGE", originals)
        self.assertIn("Page ", originals)
        self.assertNotIn("Page 1", originals)

    def test_replace_mode_writes_boundary_text_without_structural_loss(self):
        input_bytes = _build_fixture_docx()
        translator = self._translator()
        doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
        translations = {
            "Purpose": "Tujuan",
            "Flow label": "Label alur",
            "Controlled Approval Required": "Persetujuan diperlukan",
            "Block Controlled Text": "Teks kontrol blok",
            "Revision new text": "Revisi teks baru",
            "Merged Cell": "Sel gabungan",
            "Nested Cell": "Sel bersarang",
            "Review this note": "Tinjau catatan ini",
        }
        translated = [translations.get(text, text) for text in originals]

        output_bytes = translator._after_translate(doc, elements, translated, originals)

        body_text = _text_values(output_bytes)
        comment_text = _comments_text_values(output_bytes)
        self.assertIn("Tujuan", body_text)
        self.assertIn("3", body_text)
        self.assertIn("Label alur", body_text)
        self.assertIn("Persetujuan diperlukan", "".join(body_text))
        self.assertIn("Teks kontrol blok", body_text)
        self.assertIn("Revisi teks baru", "".join(body_text))
        self.assertIn("Sel gabungan", body_text)
        self.assertIn("Sel bersarang", body_text)
        self.assertIn("Tinjau catatan ini", comment_text)
        self.assertEqual(_paragraph_and_table_counts(input_bytes), _paragraph_and_table_counts(output_bytes))

    def test_append_mode_keeps_textbox_mapping(self):
        input_bytes = _build_fixture_docx()
        translator = DocxTranslator(DocxTranslatorConfig(skip_translate=True, insert_mode="append"))
        doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
        translated = ["Label alur" if text == "Flow label" else text for text in originals]

        output_bytes = translator._after_translate(doc, elements, translated, originals)

        body_text = _text_values(output_bytes)
        self.assertIn("Flow label", body_text)
        self.assertIn("Label alur", body_text)

    def test_real_translation_mode_normalizes_redundant_section_breaks(self):
        input_bytes = _build_fixture_docx()
        translator = DocxTranslator(
            DocxTranslatorConfig(
                skip_translate=False,
                api_key="dummy",
                base_url="http://127.0.0.1",
                model_id="dummy",
            )
        )
        doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
        for sect_pr in doc._element.xpath(".//w:sectPr"):
            for child in list(sect_pr):
                if child.tag in {qn("w:headerReference"), qn("w:footerReference")}:
                    sect_pr.remove(child)

        output_bytes = translator._after_translate(doc, elements, originals, originals)

        source_counts = _section_type_counts(input_bytes)
        output_counts = _section_type_counts(output_bytes)
        self.assertGreater(source_counts.get("nextPage(default)", 0), output_counts.get("nextPage(default)", 0))
        self.assertGreater(output_counts.get("continuous", 0), source_counts.get("continuous", 0))

    def test_hyperlink_runs_are_not_merged_across_parent_boundary(self):
        input_bytes = _build_fixture_docx()
        translator = self._translator()
        doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
        translated = ["Baca manual sekarang" if text == "Read manual now" else text for text in originals]

        output_bytes = translator._after_translate(doc, elements, translated, originals)

        with ZipFile(BytesIO(output_bytes)) as zf:
            tree = etree.fromstring(zf.read("word/document.xml"))
        self.assertEqual(len(tree.xpath("//w:hyperlink", namespaces=NS)), 2)

    def test_shared_headers_are_extracted_once_across_sections(self):
        document = docx.Document()
        document.sections[0].header.paragraphs[0].text = "Shared Header"
        document.add_paragraph("Body One")
        document.add_section()
        document.add_paragraph("Body Two")
        input_bytes = _docx_bytes(document)

        translator = self._translator()
        _, _, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))

        self.assertEqual(originals.count("Shared Header"), 1)

    def test_shared_header_is_written_once_and_preserved(self):
        document = docx.Document()
        document.sections[0].header.paragraphs[0].text = "Shared Header"
        document.add_paragraph("Body One")
        document.add_section()
        document.add_paragraph("Body Two")
        input_bytes = _docx_bytes(document)

        translator = self._translator()
        doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
        translated = ["Header Bersama" if text == "Shared Header" else text for text in originals]
        output_bytes = translator._after_translate(doc, elements, translated, originals)

        self.assertEqual(_header_text_values(output_bytes).count("Header Bersama"), 1)

    def test_translate_method_end_to_end_with_fake_agent(self):
        class FakeAgent:
            def send_segments(self, segments, chunk_size):
                translations = {
                    "Purpose": "Tujuan",
                    "Flow label": "Label alur",
                    "Controlled Approval Required": "Persetujuan diperlukan",
                    "Block Controlled Text": "Teks kontrol blok",
                    "Revision new text": "Revisi teks baru",
                    "Nested Cell": "Sel bersarang",
                    "Review this note": "Tinjau catatan ini",
                }
                return [translations.get(segment, segment) for segment in segments]

        document = Document(content=_build_fixture_docx(), suffix=".docx")
        translator = DocxTranslator(
            DocxTranslatorConfig(
                skip_translate=False,
                api_key="dummy",
                base_url="http://127.0.0.1",
                model_id="dummy",
            )
        )
        translator.translate_agent = FakeAgent()

        translator.translate(document)

        body_text = _text_values(document.content)
        comments_text = _comments_text_values(document.content)
        self.assertIn("Tujuan", body_text)
        self.assertIn("Label alur", body_text)
        self.assertIn("Persetujuan diperlukan", "".join(body_text))
        self.assertIn("Teks kontrol blok", body_text)
        self.assertIn("Revisi teks baru", "".join(body_text))
        self.assertIn("Sel bersarang", body_text)
        self.assertIn("Tinjau catatan ini", comments_text)


if __name__ == "__main__":
    unittest.main()