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)


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

    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()