"""Import upload guards: decompression-bomb docx must be rejected before parsing (H2 DoS)."""

from __future__ import annotations

import io
import zipfile

import pytest
from docx import Document

from app.document_extractors import ImportExtractionError, extract_text

CT = ('<?xml version="1.0"?><Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">'
      '<Default Extension="xml" ContentType="application/xml"/>'
      '<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>'
      '<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/></Types>')
RELS = ('<?xml version="1.0"?><Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
        '<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/></Relationships>')


def _custom_docx(document_xml: str) -> bytes:
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("[Content_Types].xml", CT)
        zf.writestr("_rels/.rels", RELS)
        zf.writestr("word/document.xml", document_xml)
    return buffer.getvalue()


def test_decompression_bomb_docx_is_rejected_before_parsing() -> None:
    """~20MB decompressed XML must fail fast instead of burning CPU in the parser."""
    para = "<w:p><w:r><w:t>放大攻击</w:t></w:r></w:p>"
    body = para * (20 * 1024 * 1024 // len(para.encode()))
    bomb = _custom_docx('<?xml version="1.0" encoding="UTF-8"?>'
                        '<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
                        f"<w:body>{body}</w:body></w:document>")
    assert len(bomb) < 1024 * 1024  # small compressed payload is the point of the attack

    with pytest.raises(ImportExtractionError):
        extract_text(extension=".docx", content=bomb)


def test_normal_docx_still_parses() -> None:
    document = Document()
    document.add_paragraph("张三 后端工程师")
    buffer = io.BytesIO()
    document.save(buffer)

    assert "张三" in extract_text(extension=".docx", content=buffer.getvalue())