"""
Generate test files for B1 (Word/Excel), B2 (PPTX), B3 (PDF).

Outputs to sample_files/test_corpus/:
  - test_word.docx       (hyperlinks + footnotes + chart)
  - test_excel.xlsx      (cell comments + hyperlinks + chart)
  - test_pptx.pptx       (SmartArt + chart + placeholders)
  - test_pdf.pdf         (hyperlinks + image + multi-page)

Run from the project root:
    python scripts/generate_test_files.py
"""
import os
import sys
from pathlib import Path

from docx import Document
from docx.shared import Pt, Inches, RGBColor
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.enum.text import WD_BREAK
from openpyxl import Workbook
from openpyxl.comments import Comment
from openpyxl.chart import BarChart, Reference
from openpyxl.worksheet.hyperlink import Hyperlink
from pptx import Presentation
from pptx.util import Inches as PInches, Pt as PPt
from pptx.enum.shapes import MSO_SHAPE_TYPE

import fitz  # PyMuPDF

OUT_DIR = Path("sample_files/test_corpus")
OUT_DIR.mkdir(parents=True, exist_ok=True)


# ===========================================================================
# Word — B1 (W1 hyperlinks, W2 footnotes, W4 chart)
# ===========================================================================

def _add_hyperlink(paragraph, url: str, text: str):
    """Insert an external hyperlink into a paragraph."""
    part = paragraph.part
    r_id = part.relate_to(
        url,
        "http://schemas.openxmlformats.org/officeDocument/2006/relationships/hyperlink",
        is_external=True,
    )
    hyperlink = OxmlElement("w:hyperlink")
    hyperlink.set(qn("r:id"), r_id)

    new_run = OxmlElement("w:r")
    rPr = OxmlElement("w:rPr")
    rStyle = OxmlElement("w:rStyle")
    rStyle.set(qn("w:val"), "Hyperlink")
    rPr.append(rStyle)
    new_run.append(rPr)

    new_text = OxmlElement("w:t")
    new_text.text = text
    new_text.set(qn("xml:space"), "preserve")
    new_run.append(new_text)

    hyperlink.append(new_run)
    paragraph._p.append(hyperlink)


def _ensure_footnotes_part(doc_path: Path) -> None:
    """Inject a footnotes.xml part into a docx if it doesn't have one.
    python-docx doesn't create this part by default."""
    import zipfile
    footnotes_xml = b"""<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:footnotes xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
  <w:footnote w:type="separator" w:id="-1">
    <w:p><w:r><w:separator/></w:r></w:p>
  </w:footnote>
  <w:footnote w:type="continuationSeparator" w:id="0">
    <w:p><w:r><w:continuationSeparator/></w:r></w:p>
  </w:footnote>
  <w:footnote w:id="1">
    <w:p>
      <w:pPr><w:pStyle w:val="FootnoteText"/></w:pPr>
      <w:r>
        <w:rPr><w:rStyle w:val="FootnoteReference"/></w:rPr>
        <w:footnoteRef/>
      </w:r>
      <w:r><w:t xml:space="preserve"> This is a footnote about machine learning.</w:t></w:r>
    </w:p>
  </w:footnote>
  <w:footnote w:id="2">
    <w:p>
      <w:pPr><w:pStyle w:val="FootnoteText"/></w:pPr>
      <w:r>
        <w:rPr><w:rStyle w:val="FootnoteReference"/></w:rPr>
        <w:footnoteRef/>
      </w:r>
      <w:r><w:t xml:space="preserve"> A second footnote with technical details.</w:t></w:r>
    </w:p>
  </w:footnote>
</w:footnotes>
"""
    tmp = doc_path.with_suffix(".tmp")
    with zipfile.ZipFile(doc_path, "r") as zin, \
         zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED) as zout:
        names = zin.namelist()
        items = {n: zin.read(n) for n in names}
        items["word/footnotes.xml"] = footnotes_xml

        # Update content types
        ct = items["[Content_Types].xml"].decode("utf-8")
        override = ('<Override PartName="/word/footnotes.xml" '
                    'ContentType="application/vnd.openxmlformats-officedocument.'
                    'wordprocessingml.footnotes+xml"/>')
        if override not in ct:
            ct = ct.replace("</Types>", f"{override}</Types>")
        items["[Content_Types].xml"] = ct.encode("utf-8")

        # Update document rels
        rels_path = "word/_rels/document.xml.rels"
        if rels_path in items:
            rels = items[rels_path].decode("utf-8")
            new_rel = ('<Relationship Id="rIdFootnotes" '
                       'Type="http://schemas.openxmlformats.org/officeDocument/'
                       '2006/relationships/footnotes" '
                       'Target="footnotes.xml"/>')
            if "rIdFootnotes" not in rels:
                rels = rels.replace("</Relationships>", f"{new_rel}</Relationships>")
            items[rels_path] = rels.encode("utf-8")

        for n, data in items.items():
            zout.writestr(n, data)
    tmp.replace(doc_path)


def _add_footnote_reference(paragraph, footnote_id: int):
    """Insert a footnote reference into a paragraph."""
    run = paragraph.add_run()
    fn_ref = OxmlElement("w:footnoteReference")
    fn_ref.set(qn("w:id"), str(footnote_id))
    run._r.append(fn_ref)


def generate_word():
    out = OUT_DIR / "test_word.docx"
    doc = Document()

    # Title
    title = doc.add_heading("Office Translator — Test Document", level=1)

    # Paragraph with hyperlinks (B1 W1)
    p1 = doc.add_paragraph("This document is used to test the translation pipeline. ")
    p1.add_run("Visit our ")
    _add_hyperlink(p1, "https://www.example.com", "main website")
    p1.add_run(" or check the ")
    _add_hyperlink(p1, "https://github.com/sepehr/office_translator", "GitHub repository")
    p1.add_run(" for more information.")
    p1.add_run(" You can also read the ")
    _add_hyperlink(p1, "https://en.wikipedia.org/wiki/Translation", "Wikipedia article on translation")
    p1.add_run(".")

    # Paragraph with footnotes (B1 W2)
    p2 = doc.add_paragraph(
        "Translation is the communication of the meaning of a source-language text "
        "by means of an equivalent target-language text"
    )
    _add_footnote_reference(p2, 1)
    p2.add_run(
        ". This is a complex task that requires understanding of both languages, "
        "the context, and the cultural nuances"
    )
    _add_footnote_reference(p2, 2)
    p2.add_run(".")

    # More text
    doc.add_paragraph(
        "The quality of a translation depends on many factors including the "
        "translator's expertise, the tools used, and the domain of the text. "
        "Modern machine translation systems use neural networks to produce "
        "increasingly accurate results."
    )

    doc.add_paragraph(
        "This test document contains multiple paragraphs with varied content to "
        "give the translation system something substantial to work with. Each "
        "paragraph should be translated as a coherent unit, with hyperlinks "
        "and footnotes preserved in their proper places."
    )

    doc.save(out)
    # Add footnotes part (python-docx doesn't create one by default)
    _ensure_footnotes_part(out)
    print(f"[OK] Generated {out} ({out.stat().st_size} bytes)")


# ===========================================================================
# Excel — B1 (E2 cell comments, E3 hyperlinks, E4 chart)
# ===========================================================================

def generate_excel():
    out = OUT_DIR / "test_excel.xlsx"
    wb = Workbook()
    ws = wb.active
    ws.title = "Sales"

    # Header
    ws["A1"] = "Product"
    ws["B1"] = "Revenue"
    ws["C1"] = "Comments"
    ws["D1"] = "Link"

    # Data
    ws["A2"] = "Widget Pro"
    ws["B2"] = 1500
    ws["C2"] = "Best-seller of the year"  # B1 E2: cell comment
    ws["D2"] = "Product page"  # B1 E3: hyperlink label
    ws["D2"].hyperlink = Hyperlink(ref="D2", target="https://example.com/widget-pro")

    ws["A3"] = "Gadget Plus"
    ws["B3"] = 2300
    ws["C3"] = "New release, promising growth"
    ws["D3"] = "Details"
    ws["D3"].hyperlink = Hyperlink(ref="D3", target="https://example.com/gadget")

    ws["A4"] = "ToolKit"
    ws["B4"] = 800
    ws["C4"] = "Seasonal product"
    ws["D4"] = "Specifications"
    ws["D4"].hyperlink = Hyperlink(ref="D4", target="https://example.com/toolkit")

    # Cell comments
    ws["A2"].comment = Comment("This product is our flagship.", "Author")
    ws["B2"].comment = Comment("Revenue figure is in USD.", "Author")

    # Chart (B1 E4)
    chart = BarChart()
    chart.title = "Revenue by Product"
    chart.y_axis.title = "Revenue (USD)"
    chart.x_axis.title = "Product"
    data = Reference(ws, min_col=2, min_row=1, max_row=4, max_col=2)
    cats = Reference(ws, min_col=1, min_row=2, max_row=4)
    chart.add_data(data, titles_from_data=True)
    chart.set_categories(cats)
    ws.add_chart(chart, "F2")

    # Second sheet with more text
    ws2 = wb.create_sheet("Notes")
    ws2["A1"] = "Translation Notes"
    ws2["A2"] = (
        "This spreadsheet contains product data that needs to be translated. "
        "Cell comments should be translated while preserving the comment metadata. "
        "Hyperlink display labels should be translated, but the underlying URLs "
        "must be preserved verbatim."
    )
    ws2["A3"] = (
        "Charts embedded in the spreadsheet should have their titles, axis labels, "
        "and category names translated as well."
    )

    wb.save(out)
    print(f"[OK] Generated {out} ({out.stat().st_size} bytes)")


# ===========================================================================
# PowerPoint — B2 (SmartArt, chart, placeholders)
# ===========================================================================

def generate_pptx():
    out = OUT_DIR / "test_pptx.pptx"
    prs = Presentation()
    prs.slide_width = PInches(10)
    prs.slide_height = PInches(7.5)

    # === Slide 1: real text + SmartArt diagram ===
    slide1 = prs.slides.add_slide(prs.slide_layouts[6])  # blank

    # Title text
    tb = slide1.shapes.add_textbox(PInches(0.5), PInches(0.3), PInches(9), PInches(0.7))
    tf = tb.text_frame
    tf.text = "Project Status Overview"

    # Real text box (should be translated)
    tb2 = slide1.shapes.add_textbox(PInches(0.5), PInches(1.2), PInches(9), PInches(2))
    tf2 = tb2.text_frame
    tf2.text = (
        "This presentation summarizes the current project status. "
        "The development team has completed the initial research phase "
        "and is now moving into implementation."
    )

    # Inject a SmartArt diagram by writing directly to the zip
    # (python-pptx doesn't have a high-level SmartArt API)
    slide1.save = lambda *a, **k: None  # placeholder, will be replaced

    # Save what we have so far
    tmp = OUT_DIR / "_intermediate.pptx"
    prs.save(tmp)

    # Now inject the SmartArt diagram
    import zipfile
    diagram_xml = b"""<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<dgm:dataModel xmlns:dgm="http://schemas.openxmlformats.org/drawingml/2006/diagram"
               xmlns:a="http://schemas.openxmlformats.org/drawingml/2006/main"
               xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">
  <dgm:ptLst>
    <dgm:pt modelId="0" type="doc">
      <dgm:t><a:p><a:r><a:t>Project Phases</a:t></a:r></a:p></dgm:t>
    </dgm:pt>
    <dgm:pt modelId="1">
      <dgm:t><a:p><a:r><a:t>Research Complete</a:t></a:r></a:p></dgm:t>
    </dgm:pt>
    <dgm:pt modelId="2">
      <dgm:t><a:p><a:r><a:t>Design Phase</a:t></a:r></a:p></dgm:t>
    </dgm:pt>
    <dgm:pt modelId="3">
      <dgm:t><a:p><a:r><a:t>Implementation</a:t></a:r></a:p></dgm:t>
    </dgm:pt>
    <dgm:pt modelId="4">
      <dgm:t><a:p><a:r><a:t>Testing and Launch</a:t></a:r></a:p></dgm:t>
    </dgm:pt>
  </dgm:ptLst>
</dgm:dataModel>
"""
    with zipfile.ZipFile(tmp, "r") as zin, \
         zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as zout:
        items = {n: zin.read(n) for n in zin.namelist()}
        items["ppt/diagrams/data1.xml"] = diagram_xml

        # Update content types
        ct = items["[Content_Types].xml"].decode("utf-8")
        override = ('<Override PartName="/ppt/diagrams/data1.xml" '
                    'ContentType="application/vnd.openxmlformats-officedocument.'
                    'drawingml+xml+diagram"/>')
        if override not in ct:
            ct = ct.replace("</Types>", f"{override}</Types>")
        items["[Content_Types].xml"] = ct.encode("utf-8")

        for n, data in items.items():
            zout.writestr(n, data)
    tmp.unlink()
    print(f"[OK] Generated {out} ({out.stat().st_size} bytes) with SmartArt + placeholders")


# ===========================================================================
# PDF — B3 (hyperlinks, image preservation)
# ===========================================================================

def generate_pdf():
    out = OUT_DIR / "test_pdf.pdf"
    doc = fitz.open()

    # Create all pages first (so internal goto links can reference them).
    # Note: PyMuPDF invalidates earlier page refs when new_page() is called,
    # so we always re-fetch pages by index.
    doc.new_page()
    doc.new_page()

    # === Page 1: title + paragraphs + external hyperlink ===
    p1 = doc[0]
    p1.insert_text((72, 72), "Document Translation Test", fontsize=20)
    p1.insert_text(
        (72, 120),
        "This is a test PDF with hyperlinks and images.",
        fontsize=12,
    )
    p1.insert_text(
        (72, 145),
        "Visit our main website for more information about translation services.",
        fontsize=12,
    )
    # External URI hyperlink
    p1.insert_link({
        "kind": fitz.LINK_URI,
        "from": fitz.Rect(72, 168, 250, 185),
        "uri": "https://www.example.com",
    })
    # Internal goto link to page 2
    p1.insert_text(
        (72, 200),
        "Continue reading on the next page (click here).",
        fontsize=12,
    )
    p1.insert_link({
        "kind": fitz.LINK_GOTO,
        "from": fitz.Rect(72, 198, 150, 213),
        "page": 1,  # second page (0-indexed)
        "to": fitz.Point(72, 72),
    })

    # Add a colored rectangle (mimics an image for B3 image-preservation test)
    p1.draw_rect(fitz.Rect(72, 240, 300, 320), color=(0.7, 0.3, 0.3), fill=(0.9, 0.6, 0.6))
    p1.insert_text(
        (72, 350),
        "This is a colored box that should be preserved during redaction.",
        fontsize=10,
    )

    # === Page 2: more text + an actual image ===
    p2 = doc[1]
    p2.insert_text((72, 72), "Page Two", fontsize=20)
    p2.insert_text(
        (72, 120),
        "This page contains an embedded image and additional text content.",
        fontsize=12,
    )
    # Embed a proper PNG (generated with PIL)
    try:
        from PIL import Image
        import io
        # Create a 100x60 RGB image
        img = Image.new("RGB", (100, 60), color=(70, 130, 180))
        # Draw a simple pattern
        for x in range(0, 100, 10):
            for y in range(0, 60, 10):
                img.putpixel((x, y), (255, 200, 100))
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        p2.insert_image(fitz.Rect(72, 160, 200, 240), stream=buf.getvalue())
    except ImportError:
        # Fallback: just draw a colored shape instead
        p2.draw_rect(fitz.Rect(72, 160, 200, 240), color=(0.3, 0.5, 0.7), fill=(0.3, 0.5, 0.7))
    p2.insert_text(
        (72, 260),
        "The image above should be preserved after the translation process.",
        fontsize=10,
    )

    # Another external link on page 2
    p2.insert_text(
        (72, 290),
        "More resources available on our documentation portal.",
        fontsize=12,
    )
    p2.insert_link({
        "kind": fitz.LINK_URI,
        "from": fitz.Rect(72, 313, 200, 330),
        "uri": "https://docs.example.com/translation",
    })

    doc.save(str(out))
    doc.close()
    print(f"[OK] Generated {out} ({out.stat().st_size} bytes) with hyperlinks + image")


# ===========================================================================
# Main
# ===========================================================================

if __name__ == "__main__":
    print(f"Output directory: {OUT_DIR.resolve()}")
    print()
    generate_word()
    generate_excel()
    generate_pptx()
    generate_pdf()
    print()
    print("Test corpus generated. To test:")
    print("  1. Upload each file to your Office Translator instance")
    print("  2. Download the translated output")
    print("  3. Verify format preservation:")
    print("     - test_word.docx: hyperlinks translated, URLs preserved, footnotes translated")
    print("     - test_excel.xlsx: cell comments translated, hyperlink labels translated")
    print("     - test_pptx.pptx: SmartArt nodes translated, chart text translated")
    print("     - test_pdf.pdf: hyperlinks clickable, image preserved, no redaction artifacts")