from __future__ import annotations

import io
import struct
import zipfile
from dataclasses import asdict

import pytest

from app.services.document_ingestion import (
    BlockKind,
    ChunkingConfig,
    CloudTextPolicy,
    DocumentFormat,
    DocumentIngestionError,
    IngestionErrorCode,
    IngestionStatus,
    ingest_document,
)

DOCX_MIME = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
CONTENT_TYPES = b"""<?xml version="1.0" encoding="UTF-8"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
  <Override PartName="/word/document.xml"
    ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
</Types>
"""


def _docx_bytes(document_xml: bytes, extra: dict[str, bytes] | None = None) -> bytes:
    output = io.BytesIO()
    with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as package:
        package.writestr("[Content_Types].xml", CONTENT_TYPES)
        package.writestr("word/document.xml", document_xml)
        for name, value in (extra or {}).items():
            package.writestr(name, value)
    return output.getvalue()


def _word_document(body: str) -> bytes:
    return f"""<?xml version="1.0" encoding="UTF-8"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
  <w:body>{body}<w:sectPr/></w:body>
</w:document>
""".encode()


def _mark_zip_encrypted(value: bytes) -> bytes:
    mutated = bytearray(value)
    local = mutated.find(b"PK\x03\x04")
    central = mutated.find(b"PK\x01\x02")
    assert local >= 0 and central >= 0
    local_flags = struct.unpack_from("<H", mutated, local + 6)[0]
    central_flags = struct.unpack_from("<H", mutated, central + 8)[0]
    struct.pack_into("<H", mutated, local + 6, local_flags | 0x1)
    struct.pack_into("<H", mutated, central + 8, central_flags | 0x1)
    return bytes(mutated)


def _assert_error(
    code: IngestionErrorCode,
    *,
    filename: str,
    mime: str,
    content: bytes,
    max_upload_bytes: int = 1024 * 1024,
) -> DocumentIngestionError:
    with pytest.raises(DocumentIngestionError) as captured:
        ingest_document(
            filename=filename,
            declared_mime_type=mime,
            content=content,
            max_upload_bytes=max_upload_bytes,
        )
    assert captured.value.code is code
    return captured.value


def test_utf8_markdown_preserves_heading_page_line_and_chunk_anchors() -> None:
    content = (
        "# 区域地质\n\n第一段包含铜矿化描述。\n\n"
        "## 蚀变特征\n\n钾化与绢英岩化可作为演示找矿标志。\f"
        "## 第二页\n\n第二页保留逻辑页号。"
    ).encode()

    artifact = ingest_document(
        filename="synthetic.md",
        declared_mime_type="text/markdown; charset=utf-8",
        content=content,
        chunking=ChunkingConfig(target_tokens=24, max_tokens=40, overlap_tokens=4),
    )

    assert artifact.status is IngestionStatus.READY_FOR_LOCAL_REVIEW
    assert artifact.document_format is DocumentFormat.MARKDOWN
    assert [page.page_number for page in artifact.pages] == [1, 2]
    headings = [block for block in artifact.blocks if block.kind is BlockKind.HEADING]
    assert [block.section_path for block in headings] == [
        ("区域地质",),
        ("区域地质", "蚀变特征"),
        ("区域地质", "第二页"),
    ]
    assert artifact.chunks
    assert all(chunk.anchor.block_ids for chunk in artifact.chunks)
    assert all(chunk.anchor.line_start <= chunk.anchor.line_end for chunk in artifact.chunks)
    assert artifact.chunks[-1].anchor.page_end == 2
    assert artifact.manifest is not None
    assert len(artifact.manifest.items) == len(artifact.chunks)


def test_utf16_text_is_supported_and_form_feed_creates_logical_pages() -> None:
    content = "第一页面。\f第二页面。".encode("utf-16")

    artifact = ingest_document(
        filename="synthetic.txt",
        declared_mime_type="text/plain",
        content=content,
    )

    assert [page.page_number for page in artifact.pages] == [1, 2]
    assert "第一页面" in artifact.pages[0].text
    assert "第二页面" in artifact.pages[1].text


def test_markdown_heading_inside_fence_is_not_promoted() -> None:
    artifact = ingest_document(
        filename="synthetic.md",
        declared_mime_type="text/markdown",
        content=b"# Heading\n\n```text\n# not-a-heading\n```\n",
    )

    headings = [block.text for block in artifact.blocks if block.kind is BlockKind.HEADING]
    assert headings == ["Heading"]
    assert "# not-a-heading" in artifact.blocks[-1].text


def test_docx_parses_heading_paragraph_and_table_without_third_party_library() -> None:
    document = _word_document(
        """
        <w:p><w:pPr><w:pStyle w:val="Heading1"/></w:pPr><w:r><w:t>矿床概况</w:t></w:r></w:p>
        <w:p><w:r><w:t>这是虚构的铜矿化描述。</w:t></w:r></w:p>
        <w:tbl><w:tr>
          <w:tc><w:p><w:r><w:t>样品号</w:t></w:r></w:p></w:tc>
          <w:tc><w:p><w:r><w:t>Cu_%</w:t></w:r></w:p></w:tc>
        </w:tr></w:tbl>
        """
    )

    artifact = ingest_document(
        filename="synthetic.docx",
        declared_mime_type=DOCX_MIME,
        content=_docx_bytes(document),
    )

    assert artifact.document_format is DocumentFormat.DOCX
    assert [block.kind for block in artifact.blocks] == [
        BlockKind.HEADING,
        BlockKind.PARAGRAPH,
        BlockKind.TABLE_ROW,
    ]
    assert artifact.blocks[1].section_path == ("矿床概况",)
    assert artifact.blocks[2].text == "样品号\tCu_%"
    assert artifact.pages[0].page_number is None


@pytest.mark.parametrize(
    ("filename", "mime", "content", "code"),
    [
        ("empty.txt", "text/plain", b"", IngestionErrorCode.EMPTY_FILE),
        (
            "spoof.txt",
            "application/pdf",
            b"plain text",
            IngestionErrorCode.MIME_EXTENSION_MISMATCH,
        ),
        (
            "spoof.txt",
            "text/plain",
            b"%PDF-1.7\n",
            IngestionErrorCode.MIME_CONTENT_MISMATCH,
        ),
        (
            "invalid.txt",
            "text/plain",
            b"\x81\x82\x83",
            IngestionErrorCode.INVALID_TEXT_ENCODING,
        ),
        (
            "archive.zip",
            "application/zip",
            b"PK\x03\x04",
            IngestionErrorCode.UNSUPPORTED_MEDIA_TYPE,
        ),
    ],
)
def test_upload_envelope_and_encoding_fail_closed(
    filename: str,
    mime: str,
    content: bytes,
    code: IngestionErrorCode,
) -> None:
    _assert_error(code, filename=filename, mime=mime, content=content)


def test_upload_size_is_checked_before_parsing() -> None:
    _assert_error(
        IngestionErrorCode.FILE_TOO_LARGE,
        filename="large.txt",
        mime="text/plain",
        content=b"12345",
        max_upload_bytes=4,
    )


def test_docx_rejects_path_traversal_and_active_content() -> None:
    document = _word_document("<w:p><w:r><w:t>safe</w:t></w:r></w:p>")
    traversal = _docx_bytes(document, {"../outside.txt": b"bad"})
    active = _docx_bytes(document, {"word/vbaProject.bin": b"macro"})

    _assert_error(
        IngestionErrorCode.DOCX_PATH_TRAVERSAL,
        filename="unsafe.docx",
        mime=DOCX_MIME,
        content=traversal,
    )
    _assert_error(
        IngestionErrorCode.DOCX_ACTIVE_CONTENT,
        filename="active.docx",
        mime=DOCX_MIME,
        content=active,
    )


def test_docx_rejects_encryption_flag_and_compression_bomb() -> None:
    document = _word_document("<w:p><w:r><w:t>safe</w:t></w:r></w:p>")
    encrypted = _mark_zip_encrypted(_docx_bytes(document))
    bomb = _docx_bytes(document, {"word/media/repeated.bin": b"A" * 2_000_000})

    _assert_error(
        IngestionErrorCode.DOCX_ENCRYPTED,
        filename="encrypted.docx",
        mime=DOCX_MIME,
        content=encrypted,
    )
    _assert_error(
        IngestionErrorCode.DOCX_PACKAGE_LIMIT,
        filename="bomb.docx",
        mime=DOCX_MIME,
        content=bomb,
        max_upload_bytes=3_000_000,
    )


def test_docx_rejects_doctype_and_missing_required_parts() -> None:
    unsafe_xml = b"""<!DOCTYPE x [<!ENTITY secret "unsafe">]>
    <w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
      <w:body><w:p><w:r><w:t>&secret;</w:t></w:r></w:p></w:body>
    </w:document>"""
    incomplete = io.BytesIO()
    with zipfile.ZipFile(incomplete, "w") as package:
        package.writestr("[Content_Types].xml", CONTENT_TYPES)

    _assert_error(
        IngestionErrorCode.DOCX_UNSAFE_XML,
        filename="unsafe.docx",
        mime=DOCX_MIME,
        content=_docx_bytes(unsafe_xml),
    )
    _assert_error(
        IngestionErrorCode.INVALID_DOCX_PACKAGE,
        filename="incomplete.docx",
        mime=DOCX_MIME,
        content=incomplete.getvalue(),
    )


def test_pdf_is_routed_fail_closed_without_text_or_spatial_claims() -> None:
    artifact = ingest_document(
        filename="synthetic.pdf",
        declared_mime_type="application/pdf",
        content=b"%PDF-1.7\nsynthetic bytes only",
    )

    assert artifact.status is IngestionStatus.OCR_REQUIRED
    assert artifact.pages == ()
    assert artifact.blocks == ()
    assert artifact.chunks == ()
    assert artifact.manifest is None
    assert "NO_MAP_OR_SPATIAL_UNDERSTANDING" in artifact.limitations


def test_chunking_is_bounded_overlapping_and_deterministic() -> None:
    content = (" ".join(f"term{index}" for index in range(900))).encode()
    config = ChunkingConfig(target_tokens=512, max_tokens=800, overlap_tokens=64)

    first = ingest_document(
        filename="long.txt",
        declared_mime_type="text/plain",
        content=content,
        chunking=config,
    )
    second = ingest_document(
        filename="renamed.txt",
        declared_mime_type="text/plain",
        content=content,
        chunking=config,
    )

    assert [chunk.token_count for chunk in first.chunks] == [512, 452]
    assert all(chunk.token_count <= config.max_tokens for chunk in first.chunks)
    assert first == second
    assert first.manifest is not None and second.manifest is not None
    assert first.manifest.manifest_sha256 == second.manifest.manifest_sha256
    first_tail = first.chunks[0].display_text.split()[-64:]
    second_head = first.chunks[1].display_text.split()[:64]
    assert first_tail == second_head
    reconstructed = (
        first.chunks[0].display_text.split()
        + first.chunks[1].display_text.split()[config.overlap_tokens :]
    )
    assert reconstructed == content.decode().split()


def test_cloud_and_embedding_text_are_separate_and_hash_bound() -> None:
    artifact = ingest_document(
        filename="redacted.md",
        declared_mime_type="text/markdown",
        content="# 钻孔\n\n项目代号 DEMO-42 位于虚构地区。".encode(),
        cloud_policy=CloudTextPolicy(
            policy_id="synthetic-redaction-v1",
            redact_literals=("DEMO-42",),
        ),
    )
    chunk = artifact.chunks[0]

    assert "DEMO-42" in chunk.display_text
    assert "DEMO-42" not in chunk.cloud_text
    assert "[REDACTED]" in chunk.cloud_text
    assert chunk.embedding_text == chunk.embedding_prefix + chunk.cloud_text
    assert chunk.display_text_sha256 != chunk.cloud_text_sha256
    assert chunk.cloud_text_sha256 != chunk.embedding_text_sha256
    assert artifact.manifest is not None
    assert artifact.manifest.items[0].cloud_text_sha256 == chunk.cloud_text_sha256


def test_credential_shapes_never_enter_errors_or_artifacts() -> None:
    secret = "sk-" + "A" * 24
    error = _assert_error(
        IngestionErrorCode.SENSITIVE_CONTENT_DETECTED,
        filename="secret.txt",
        mime="text/plain",
        content=f"credential={secret}".encode(),
    )

    assert secret not in str(error)
    assert secret not in repr(error)


def test_manifest_and_citation_anchor_change_when_source_changes() -> None:
    first = ingest_document(
        filename="anchor.md",
        declared_mime_type="text/markdown",
        content="# 标题\n\n第一版证据。".encode(),
    )
    second = ingest_document(
        filename="anchor.md",
        declared_mime_type="text/markdown",
        content="# 标题\n\n第二版证据。".encode(),
    )

    assert first.manifest is not None and second.manifest is not None
    assert first.raw_sha256 != second.raw_sha256
    assert first.chunks[0].anchor.anchor_id != second.chunks[0].anchor.anchor_id
    assert first.manifest.manifest_sha256 != second.manifest.manifest_sha256
    assert first.chunks[0].anchor.normalized_text_sha256 == first.normalized_text_sha256
    assert first.chunks[0].anchor.char_start < first.chunks[0].anchor.char_end
    assert asdict(first)["manifest"]["manifest_sha256"] == first.manifest.manifest_sha256