YG_FT/backend/tests/test_data_process_algorithms.py

from __future__ import annotations

import io
import json
import xml.etree.ElementTree as ET
import zipfile
from datetime import datetime

import pytest
from docx import Document
from openpyxl import Workbook
from pptx import Presentation
from pptx.util import Inches
from pypdf import PdfWriter

from app.modules.data_process.algorithms import (
    PdfPageText,
    chunk_unstructured,
    content_quality_flags,
    desensitize_pii,
    desensitize_structured_record,
    detect_document_structure,
    detect_pdf_document_noise,
    detect_text_format,
    estimate_token_count,
    extract_pdf_page_texts,
    extract_structured_records,
    generate_standard_records,
    is_near_duplicate,
    merge_short_blocks,
    normalize_text,
    parse_text_content,
    preprocess_structured_records,
    record_fingerprint,
    remove_document_noise,
    score_quality,
    stable_split,
)


def _pdf_page_texts(*texts: str) -> tuple[PdfPageText, ...]:
    pages: list[PdfPageText] = []
    offset = 0
    for page_number, text in enumerate(texts, start=1):
        normalized = normalize_text(text)
        if pages:
            offset += 2
        start = offset
        offset += len(normalized)
        pages.append(
            PdfPageText(
                page_number=page_number,
                text=normalized,
                source_start=start,
                source_end=offset,
            )
        )
    return tuple(pages)


def _minimal_pdf(text: str = "Hello PDF") -> bytes:
    stream = f"BT /F1 12 Tf 72 720 Td ({text}) Tj ET".encode("ascii")
    objects = [
        b"<< /Type /Catalog /Pages 2 0 R >>",
        b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
        (
            b"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] "
            b"/Resources << /Font << /F1 5 0 R >> >> /Contents 4 0 R >>"
        ),
        b"<< /Length " + str(len(stream)).encode("ascii") + b" >>\nstream\n"
        + stream
        + b"\nendstream",
        b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
    ]
    result = bytearray(b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n")
    offsets = [0]
    for object_number, value in enumerate(objects, start=1):
        offsets.append(len(result))
        result.extend(f"{object_number} 0 obj\n".encode("ascii"))
        result.extend(value)
        result.extend(b"\nendobj\n")
    xref_offset = len(result)
    result.extend(f"xref\n0 {len(objects) + 1}\n".encode("ascii"))
    result.extend(b"0000000000 65535 f \n")
    for offset in offsets[1:]:
        result.extend(f"{offset:010d} 00000 n \n".encode("ascii"))
    result.extend(
        (
            f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n"
            f"startxref\n{xref_offset}\n%%EOF\n"
        ).encode("ascii")
    )
    return bytes(result)


def _aes_encrypted_pdf(*, user_password: str) -> bytes:
    writer = PdfWriter(clone_from=io.BytesIO(_minimal_pdf()))
    writer.encrypt(
        user_password=user_password,
        owner_password="owner-secret",
        algorithm="AES-256",
    )
    output = io.BytesIO()
    writer.write(output)
    return output.getvalue()


def _docx_bytes() -> bytes:
    document = Document()
    document.add_heading("服务说明", level=1)
    document.add_paragraph("这是 DOCX 正文。")
    table = document.add_table(rows=1, cols=2)
    table.cell(0, 0).text = "字段"
    table.cell(0, 1).text = "内容"
    output = io.BytesIO()
    document.save(output)
    return output.getvalue()


def _xlsx_bytes() -> bytes:
    workbook = Workbook()
    worksheet = workbook.active
    worksheet.title = "数据"
    worksheet.append(["name", "score", "created_at"])
    worksheet.append(["Alice", 95, datetime(2026, 7, 23, 10, 30)])
    worksheet.append(["Bob", 88, datetime(2026, 7, 24, 9, 0)])
    output = io.BytesIO()
    workbook.save(output)
    workbook.close()
    return output.getvalue()


def _xlsx_with_worksheet_relationship(
    raw: bytes,
    target: str,
    *,
    target_mode: str | None = None,
) -> bytes:
    member_name = "xl/_rels/workbook.xml.rels"
    source = io.BytesIO(raw)
    output = io.BytesIO()
    with zipfile.ZipFile(source) as original, zipfile.ZipFile(output, "w") as rewritten:
        for member in original.infolist():
            content = original.read(member.filename)
            if member.filename == member_name:
                root = ET.fromstring(content)
                worksheet_relationship = next(
                    element
                    for element in root
                    if element.attrib.get("Type", "").endswith("/worksheet")
                )
                worksheet_relationship.set("Target", target)
                if target_mode is None:
                    worksheet_relationship.attrib.pop("TargetMode", None)
                else:
                    worksheet_relationship.set("TargetMode", target_mode)
                content = ET.tostring(root, encoding="utf-8", xml_declaration=True)
            rewritten.writestr(member, content)
    return output.getvalue()


def _pptx_bytes() -> bytes:
    presentation = Presentation()
    slide = presentation.slides.add_slide(presentation.slide_layouts[6])
    text_box = slide.shapes.add_textbox(Inches(1), Inches(1), Inches(6), Inches(1))
    text_box.text = "PPTX 页面正文"
    output = io.BytesIO()
    presentation.save(output)
    return output.getvalue()


def test_parse_utf8_json_jsonl_csv_markdown_and_txt() -> None:
    parsed_json = parse_text_content(
        b'\xef\xbb\xbf{"data":[{"name":"\xe5\xbc\xa0\xe4\xb8\x89"}]}',
        filename="records.json",
    )
    assert parsed_json.format == "json"
    assert parsed_json.records == ({"name": "张三"},)

    parsed_jsonl = parse_text_content('{"id":1}\n\n{"id":2}\n', filename="records.jsonl")
    assert parsed_jsonl.format == "jsonl"
    assert parsed_jsonl.records == ({"id": 1}, {"id": 2})

    parsed_csv = parse_text_content("name,answer\r\nAlice,yes\r\nBob,no", filename="records.csv")
    assert parsed_csv.format == "csv"
    assert parsed_csv.text == "name,answer\nAlice,yes\nBob,no"
    assert parsed_csv.records[1] == {"name": "Bob", "answer": "no"}

    parsed_markdown = parse_text_content("# 标题\n\n正文", filename="README.md")
    assert parsed_markdown.format == "markdown"
    assert parsed_markdown.records == ()

    parsed_txt = parse_text_content("普通文本", filename="note.txt")
    assert parsed_txt.format == "txt"
    assert parsed_txt.text == "普通文本"


def test_parse_pdf_docx_xlsx_and_pptx() -> None:
    parsed_pdf = parse_text_content(_minimal_pdf(), filename="manual.pdf")
    assert parsed_pdf.format == "pdf"
    assert "Hello PDF" in parsed_pdf.text
    assert parsed_pdf.records == ()

    pdf_pages = extract_pdf_page_texts(_minimal_pdf())
    assert len(pdf_pages) == 1
    assert pdf_pages[0].page_number == 1
    assert pdf_pages[0].text == "Hello PDF"
    assert pdf_pages[0].source_start == 0
    assert pdf_pages[0].source_end == len(parsed_pdf.text)

    parsed_docx = parse_text_content(_docx_bytes(), filename="manual.docx")
    assert parsed_docx.format == "docx"
    assert "服务说明" in parsed_docx.text
    assert "这是 DOCX 正文。" in parsed_docx.text
    assert "字段\t内容" in parsed_docx.text
    assert parsed_docx.records == ()

    parsed_xlsx = parse_text_content(_xlsx_bytes(), filename="records.xlsx")
    assert parsed_xlsx.format == "xlsx"
    assert parsed_xlsx.records == (
        {"name": "Alice", "score": 95, "created_at": "2026-07-23T10:30:00"},
        {"name": "Bob", "score": 88, "created_at": "2026-07-24T09:00:00"},
    )
    assert json.loads(parsed_xlsx.text.splitlines()[0]) == parsed_xlsx.records[0]

    parsed_pptx = parse_text_content(_pptx_bytes(), filename="slides.pptx")
    assert parsed_pptx.format == "pptx"
    assert parsed_pptx.text == "PPTX 页面正文"
    assert parsed_pptx.records == ()


def test_pdf_document_noise_removes_headers_page_numbers_and_toc_safely() -> None:
    pages = _pdf_page_texts(
        """
        远光制度文件 文件编码 2024
        秘密等级 商密【中】
        第 1 页 共 5 页
        正文第一页，关于适用范围的说明。
        业务提示保留
        第一页补充说明甲
        第一页补充说明乙
        第一页补充说明丙
        """,
        """
        远光制度文件 文件编码 2024
        秘密等级 商密【中】
        第 2 页 共 5 页
        目 录
        第一章 总则........3
        第二章 报销申请........4
        第三章 附则........5
        """,
        """
        远光制度文件 文件编码 2024
        秘密等级 商密【中】
        第 3 页 共 5 页
        1.1 管理要求........6
        1.2 审批职责 7
        1.3 费用标准........8
        1.4 例外处理........9
        """,
        """
        远光制度文件 文件编码 2024
        秘密等级 商密【中】
        第 4 页 共 5 页
        正文中可以说“请参见第 3 页说明”，不应误删。
        第 99 页 共 100 页
        系统可用率........99.9%
        业务提示保留
        第四页补充说明甲
        第四页补充说明乙
        第四页补充说明丙
        """,
        """
        远光制度文件 文件编码 2024
        秘密等级 商密【中】
        第 5 页 共 5 页
        本办法自发布之日起施行。
        业务提示保留
        第五页补充说明甲
        第五页补充说明乙
        第五页补充说明丙
        """,
    )
    source = "\n\n".join(page.text for page in pages)

    spans = detect_pdf_document_noise(pages)
    cleaned = remove_document_noise(source, spans)

    assert {span.kind for span in spans} == {
        "page_number",
        "repeated_margin",
        "table_of_contents",
    }
    assert "远光制度文件" not in cleaned
    assert "商密【中】" not in cleaned
    assert "第 1 页 共 5 页" not in cleaned
    assert "第一章 总则" not in cleaned
    assert "1.2 审批职责 7" not in cleaned
    assert "请参见第 3 页说明" in cleaned
    assert "第 99 页 共 100 页" in cleaned
    assert "系统可用率........99.9%" in cleaned
    assert cleaned.count("业务提示保留") == 3


def test_pdf_document_noise_does_not_infer_repeated_margins_for_short_documents() -> None:
    pages = _pdf_page_texts(
        "公司内部文件\n正文 A",
        "公司内部文件\n正文 B",
    )

    spans = detect_pdf_document_noise(pages)

    assert not any(span.kind == "repeated_margin" for span in spans)


def test_xlsx_merged_multilevel_headers_are_flattened_without_losing_columns() -> None:
    workbook = Workbook()
    worksheet = workbook.active
    worksheet.merge_cells("A1:A2")
    worksheet.merge_cells("B1:C1")
    worksheet["A1"] = "地区"
    worksheet["B1"] = "销售"
    worksheet["B2"] = "Q1"
    worksheet["C2"] = "Q2"
    worksheet.append(["华东", 100, 120])
    output = io.BytesIO()
    workbook.save(output)
    workbook.close()

    parsed = parse_text_content(output.getvalue(), filename="sales.xlsx")
    assert parsed.records == ({"地区": "华东", "销售.Q1": 100, "销售.Q2": 120},)


def test_xlsx_header_inference_skips_more_than_eight_merged_report_titles() -> None:
    workbook = Workbook()
    worksheet = workbook.active
    for row_number in range(1, 13):
        worksheet.merge_cells(
            start_row=row_number,
            start_column=1,
            end_row=row_number,
            end_column=4,
        )
        worksheet.cell(row_number, 1, f"报表说明 {row_number}")
    worksheet.append(["姓名", "部门", "得分", "日期"])
    worksheet.append(["张三", "研发", 95, "2026-07-23"])
    output = io.BytesIO()
    workbook.save(output)
    workbook.close()

    parsed = parse_text_content(output.getvalue(), filename="report.xlsx")
    assert parsed.records == (
        {"姓名": "张三", "部门": "研发", "得分": 95, "日期": "2026-07-23"},
    )


def test_xlsx_header_inference_ignores_continuous_body_merges() -> None:
    workbook = Workbook()
    worksheet = workbook.active
    worksheet.append(["类别", "名称", "数量"])
    worksheet.append(["水果", "苹果", 10])
    worksheet.append([None, "香蕉", 12])
    worksheet.append(["蔬菜", "白菜", 8])
    worksheet.append([None, "萝卜", 9])
    worksheet.merge_cells("A2:A3")
    worksheet.merge_cells("A4:A5")
    output = io.BytesIO()
    workbook.save(output)
    workbook.close()

    parsed = parse_text_content(output.getvalue(), filename="inventory.xlsx")
    assert parsed.records == (
        {"类别": "水果", "名称": "苹果", "数量": 10},
        {"类别": "", "名称": "香蕉", "数量": 12},
        {"类别": "蔬菜", "名称": "白菜", "数量": 8},
        {"类别": "", "名称": "萝卜", "数量": 9},
    )


def test_xlsx_header_inference_supports_title_and_two_header_levels() -> None:
    workbook = Workbook()
    worksheet = workbook.active
    worksheet.merge_cells("A1:C1")
    worksheet["A1"] = "区域销售报表"
    worksheet["A2"] = "统计日期"
    worksheet["B2"] = "2026-07-23"
    worksheet.merge_cells("A4:A5")
    worksheet.merge_cells("B4:C4")
    worksheet["A4"] = "地区"
    worksheet["B4"] = "销售"
    worksheet["B5"] = "Q1"
    worksheet["C5"] = "Q2"
    worksheet.append(["华南", 88, 92])
    output = io.BytesIO()
    workbook.save(output)
    workbook.close()

    parsed = parse_text_content(output.getvalue(), filename="two-level.xlsx")
    assert parsed.records == (
        {"地区": "华南", "销售.Q1": 88, "销售.Q2": 92},
    )


def test_xlsx_header_inference_supports_title_and_three_header_levels() -> None:
    workbook = Workbook()
    worksheet = workbook.active
    worksheet.merge_cells("A1:D1")
    worksheet["A1"] = "年度销售分析报告"
    worksheet["A2"] = "统计日期"
    worksheet["B2"] = "2026-07-23"
    worksheet.merge_cells("A4:A6")
    worksheet.merge_cells("B4:D4")
    worksheet.merge_cells("B5:C5")
    worksheet.merge_cells("D5:D6")
    worksheet["A4"] = "地区"
    worksheet["B4"] = "销售"
    worksheet["B5"] = "国内"
    worksheet["D5"] = "海外"
    worksheet["B6"] = "Q1"
    worksheet["C6"] = "Q2"
    worksheet.append(["华东", 100, 120, 80])
    output = io.BytesIO()
    workbook.save(output)
    workbook.close()

    parsed = parse_text_content(output.getvalue(), filename="three-level.xlsx")
    assert parsed.records == (
        {
            "地区": "华东",
            "销售.国内.Q1": 100,
            "销售.国内.Q2": 120,
            "销售.海外": 80,
        },
    )


def test_xlsx_header_inference_keeps_an_ordinary_single_header_row() -> None:
    parsed = parse_text_content(_xlsx_bytes(), filename="ordinary.xlsx")
    assert tuple(parsed.records[0]) == ("name", "score", "created_at")
    assert len(parsed.records) == 2


@pytest.mark.parametrize(
    "target",
    [
        "./worksheets/../worksheets/sheet1.xml",
        "./worksheets/%2e%2e/worksheets/sheet1.xml",
        "../xl/worksheets/sheet1.xml",
        "/xl/worksheets/./sheet1.xml",
    ],
)
def test_xlsx_worksheet_relationship_allows_safe_dot_segments(target: str) -> None:
    raw = _xlsx_with_worksheet_relationship(_xlsx_bytes(), target)
    parsed = parse_text_content(raw, filename="records.xlsx")
    assert parsed.records[0]["name"] == "Alice"


@pytest.mark.parametrize(
    "target",
    [
        "../../outside.xml",
        "worksheets\\sheet1.xml",
        "%2e%2e/%2e%2e/outside.xml",
        "%252e%252e/%252e%252e/outside.xml",
        "https://example.com/sheet1.xml",
    ],
)
def test_xlsx_worksheet_relationship_rejects_path_traversal(target: str) -> None:
    raw = _xlsx_with_worksheet_relationship(_xlsx_bytes(), target)
    with pytest.raises(ValueError, match="unsafe worksheet path"):
        parse_text_content(raw, filename="unsafe.xlsx")


def test_xlsx_worksheet_relationship_rejects_external_and_missing_targets() -> None:
    external = _xlsx_with_worksheet_relationship(
        _xlsx_bytes(),
        "https://example.com/sheet1.xml",
        target_mode="External",
    )
    with pytest.raises(ValueError, match="external relationship"):
        parse_text_content(external, filename="external.xlsx")

    missing = _xlsx_with_worksheet_relationship(
        _xlsx_bytes(),
        "worksheets/missing.xml",
    )
    with pytest.raises(ValueError, match="target does not exist"):
        parse_text_content(missing, filename="missing.xlsx")


@pytest.mark.parametrize(
    ("filename", "replacement"),
    [
        ("legacy.doc", ".docx"),
        ("legacy.xls", ".xlsx"),
        ("legacy.ppt", ".pptx"),
    ],
)
def test_legacy_office_formats_require_conversion(filename: str, replacement: str) -> None:
    with pytest.raises(ValueError, match=rf"convert the file to \{replacement}"):
        parse_text_content(b"legacy", filename=filename)


def test_office_zip_bomb_and_invalid_pdf_are_rejected_before_parsing() -> None:
    archive = io.BytesIO()
    with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as package:
        package.writestr("[Content_Types].xml", "<Types/>")
        package.writestr("word/document.xml", b"A" * (2 * 1024 * 1024))
    with pytest.raises(ValueError, match="unsafe compression ratio"):
        parse_text_content(archive.getvalue(), filename="unsafe.docx")

    active_xml = io.BytesIO()
    with zipfile.ZipFile(active_xml, "w") as package:
        package.writestr("[Content_Types].xml", "<Types/>")
        package.writestr(
            "word/document.xml",
            '<!DOCTYPE document [<!ENTITY xxe SYSTEM "file:///etc/passwd">]><document/>',
        )
    with pytest.raises(ValueError, match="unsupported active XML"):
        parse_text_content(active_xml.getvalue(), filename="active.docx")

    with pytest.raises(ValueError, match="missing PDF header"):
        parse_text_content(b"not a pdf", filename="broken.pdf")

    blank_pdf = io.BytesIO()
    blank_writer = PdfWriter()
    blank_writer.add_blank_page(width=612, height=792)
    blank_writer.write(blank_pdf)
    with pytest.raises(ValueError, match="scanned PDF requires OCR"):
        parse_text_content(blank_pdf.getvalue(), filename="scanned.pdf")

    aes_pdf_without_open_password = parse_text_content(
        _aes_encrypted_pdf(user_password=""),
        filename="aes-no-password.pdf",
    )
    assert "Hello PDF" in aes_pdf_without_open_password.text

    with pytest.raises(ValueError, match="password-protected PDF files are not supported"):
        parse_text_content(
            _aes_encrypted_pdf(user_password="secret"),
            filename="aes-password.pdf",
        )


def test_invalid_utf8_and_malformed_structured_content_fail_loudly() -> None:
    with pytest.raises(ValueError, match="not valid UTF-8"):
        parse_text_content(b"\xff\xfe", filename="broken.txt")
    with pytest.raises(ValueError, match="invalid JSONL at line 2"):
        extract_structured_records('{"id":1}\nnot-json', "jsonl")
    with pytest.raises(ValueError, match="more fields"):
        extract_structured_records("a,b\n1,2,3", "csv")


def test_detect_format_from_content_and_normalize() -> None:
    assert detect_text_format(text='{"id":1}\n{"id":2}') == "jsonl"
    assert detect_text_format(text="# Heading\ntext") == "markdown"
    assert detect_text_format(text="a,b\n1,2") == "csv"
    assert normalize_text("\ufeffＡＢＣ  \r\n第二\x00行\u200b\t \r\n") == "ABC\n第二行"


def test_extract_json_scalar_and_nested_values_are_stable() -> None:
    assert extract_structured_records("[1, true, null]", "json") == [
        {"value": 1},
        {"value": True},
        {"value": None},
    ]
    result = extract_structured_records(
        json.dumps({"items": [{"text": "  内容  "}], "ignored": 1}, ensure_ascii=False),
        "json",
    )
    assert result == [{"text": "内容"}]


def test_desensitize_pii_returns_masked_text_and_counts() -> None:
    source = "邮箱 a.user+tag@example.com，手机 +86 13800138000，身份证 11010519491231002X。"
    masked, counts = desensitize_pii(source)
    assert masked == "邮箱 [EMAIL]，手机 [PHONE]，身份证 [ID_CARD]。"
    assert counts == {"email": 1, "phone": 1, "id_card": 1, "total": 3}


def test_every_structured_preprocess_option_has_independent_behavior() -> None:
    clean_source = [
        {"id": "1", "name": "有效", "empty_column": ""},
        {"id": "", "name": "缺少关键字段", "empty_column": ""},
        {"id": "2", "name": "有效", "empty_column": ""},
    ]
    assert preprocess_structured_records(clean_source, []) == clean_source
    assert preprocess_structured_records(clean_source, ["clean_invalid"]) == [
        {"id": "1", "name": "有效"},
        {"id": "2", "name": "有效"},
    ]

    nested = [{"id": 1, "profile": {"name": "张三", "level": 2}}]
    assert "profile" in preprocess_structured_records(nested, [])[0]
    assert preprocess_structured_records(nested, ["detect_structure"])[0] == {
        "id": 1,
        "profile.name": "张三",
        "profile.level": 2,
    }

    duplicates = [
        {"customer_id": "C-1", "value": "first"},
        {"customer_id": "C-1", "value": "updated"},
        {"customer_id": "", "value": "blank-one"},
        {"customer_id": "", "value": "blank-two"},
    ]
    assert len(preprocess_structured_records(duplicates, [])) == 4
    deduplicated = preprocess_structured_records(duplicates, ["deduplicate"])
    assert [record["value"] for record in deduplicated] == [
        "first",
        "blank-one",
        "blank-two",
    ]

    unnormalized = [{" User Name ": "ＡＢＣ\r\n第二行"}]
    assert preprocess_structured_records(unnormalized, []) == unnormalized
    assert preprocess_structured_records(unnormalized, ["normalize_format"]) == [
        {"user_name": "ABC\n第二行"}
    ]

    anomaly_source = [
        {"id": 10_000 + index, "amount": amount, "text": "正常内容"}
        for index, amount in enumerate((10, 10, 11, 11, 12, 12, 13, 1000))
    ]
    assert len(preprocess_structured_records(anomaly_source, [])) == 8
    filtered = preprocess_structured_records(anomaly_source, ["filter_anomaly"])
    assert len(filtered) == 7
    assert all(record["amount"] != 1000 for record in filtered)
    assert max(record["id"] for record in filtered) > 10_000

    sensitive = [{"姓名": "张三", "phone": "13800138000", "email": "a@b.com"}]
    assert preprocess_structured_records(sensitive, []) == sensitive
    masked = preprocess_structured_records(sensitive, ["desensitize"])[0]
    assert masked == {"姓名": "[NAME]", "phone": "[PHONE]", "email": "[EMAIL]"}


def test_structured_desensitization_counts_and_document_helpers() -> None:
    masked, counts = desensitize_structured_record(
        {"联系人姓名": "李四", "说明": "邮箱 user@example.com，手机 13900139000"}
    )
    assert masked == {
        "联系人姓名": "[NAME]",
        "说明": "邮箱 [EMAIL]，手机 [PHONE]",
    }
    assert counts == {"email": 1, "phone": 1, "id_card": 0, "name": 1, "total": 3}

    structure = detect_document_structure(
        "# 第一章\n正文\n\n## 细节\n- 项目一\n- 项目二\n\n```python\nprint(1)\n```"
    )
    assert [heading.title for heading in structure.headings] == ["第一章", "细节"]
    assert structure.list_block_count == 1
    assert structure.code_block_count == 1
    assert merge_short_blocks(["短一", "短二", "这是一段足够长的正文内容"], min_token_count=4)
    assert "mojibake" in content_quality_flags("正常文字锟斤拷内容", min_chars=0, min_tokens=0)
    assert is_near_duplicate(
        "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron",
        "alpha beta gamma, delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron",
        similarity_threshold=0.92,
        max_hamming_distance=2,
    )


@pytest.mark.parametrize("method", ["structure", "fixed", "custom"])
def test_chunk_methods_preserve_offsets_and_always_advance(method: str) -> None:
    text = "# 第一章\n" + "甲。" * 18 + "\n# 第二章\n" + "乙。" * 18
    kwargs = {"custom_delimiter": "\\n"} if method == "custom" else {}
    chunks = chunk_unstructured(
        text,
        method=method,  # type: ignore[arg-type]
        chunk_size=12,
        chunk_overlap=2,
        min_chunk_size=4,
        **kwargs,
    )
    assert len(chunks) > 1
    assert all(chunk.content == normalize_text(text)[chunk.start : chunk.end] for chunk in chunks)
    assert all(chunk.end > chunk.start for chunk in chunks)
    assert all(left.start < right.start for left, right in zip(chunks, chunks[1:]))
    assert all(chunk.start_line <= chunk.end_line for chunk in chunks)


def test_default_and_structure_chunking_split_headings_without_cross_section_overlap() -> None:
    text = (
        "# 第一章\n"
        + " ".join(f"alpha{i}" for i in range(18))
        + "\n# 第二章\n"
        + " ".join(f"beta{i}" for i in range(18))
    )
    normalized = normalize_text(text)
    second_chapter_start = normalized.index("# 第二章")
    kwargs = {"chunk_size": 10, "chunk_overlap": 3, "min_chunk_size": 4}

    default_chunks = chunk_unstructured(text, **kwargs)
    structure_chunks = chunk_unstructured(text, method="structure", **kwargs)

    assert default_chunks == structure_chunks
    assert len(structure_chunks) > 2
    assert all(
        chunk.content == normalized[chunk.start : chunk.end] for chunk in structure_chunks
    )
    assert all(
        not (chunk.start < second_chapter_start < chunk.end) for chunk in structure_chunks
    )
    second_chapter_chunks = [
        chunk for chunk in structure_chunks if chunk.start >= second_chapter_start
    ]
    assert second_chapter_chunks[0].start == second_chapter_start
    assert second_chapter_chunks[0].content.startswith("# 第二章")


def test_fixed_chunk_offsets_and_actual_token_overlap_are_exact() -> None:
    text = "  ".join(f"token{i}" for i in range(30))
    normalized = normalize_text(text)
    chunks = chunk_unstructured(
        text,
        method="fixed",
        chunk_size=10,
        chunk_overlap=3,
        min_chunk_size=4,
    )
    assert len(chunks) > 2
    assert all(chunk.content == normalized[chunk.start : chunk.end] for chunk in chunks)
    assert all(chunk.token_count == estimate_token_count(chunk.content) for chunk in chunks)
    assert all(chunk.token_count == 10 for chunk in chunks[:-1])
    for left, right in zip(chunks, chunks[1:]):
        overlap_text = normalized[right.start : left.end]
        assert right.start < left.end
        assert estimate_token_count(overlap_text) == 3
        assert left.content.endswith(overlap_text)
        assert right.content.startswith(overlap_text)


def test_chunk_line_numbers_treat_newline_as_previous_line_boundary() -> None:
    chunks = chunk_unstructured(
        "第一行。\n第二行。\n第三行。",
        method="custom",
        chunk_size=8,
        chunk_overlap=0,
        min_chunk_size=2,
        custom_delimiter="\\n",
    )
    assert chunks[0].content.endswith("\n")
    assert chunks[0].start_line == 1
    assert chunks[0].end_line == 1
    assert chunks[1].start_line == 2


def test_custom_delimiter_is_preserved_as_the_chunk_boundary() -> None:
    custom_chunks = chunk_unstructured(
        "a b c d <CUT> e f g h i j",
        method="custom",
        chunk_size=8,
        chunk_overlap=0,
        min_chunk_size=2,
        custom_delimiter="<CUT>",
    )
    assert custom_chunks[0].content.endswith("<CUT>")


@pytest.mark.parametrize(
    ("field", "block"),
    [
        (
            "preserve_code_blocks",
            "```python\n" + "\n".join(f"value_{i} = {i}" for i in range(30)) + "\n```",
        ),
        (
            "preserve_tables",
            "| 字段 | 说明 |\n| --- | --- |\n"
            + "\n".join(f"| field_{i} | value_{i} |" for i in range(30)),
        ),
        (
            "preserve_lists",
            "\n".join(f"- 第 {i} 项需要完整保留" for i in range(30)),
        ),
    ],
)
def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None:
    text = "前言。" * 15 + "\n" + block + "\n" + "结尾。" * 40
    unprotected = chunk_unstructured(
        text,
        method="fixed",
        chunk_size=40,
        chunk_overlap=0,
        min_chunk_size=10,
    )
    chunks = chunk_unstructured(
        text,
        method="fixed",
        chunk_size=40,
        chunk_overlap=0,
        min_chunk_size=10,
        **{field: True},
    )
    assert all(block not in chunk.content for chunk in unprotected)
    assert any(block in chunk.content for chunk in chunks)


@pytest.mark.parametrize(
    ("kwargs", "message"),
    [
        ({"chunk_size": 0}, "chunk_size"),
        ({"chunk_size": 10, "chunk_overlap": 10}, "chunk_overlap"),
        ({"chunk_size": 10, "chunk_overlap": 0, "min_chunk_size": 11}, "min_chunk_size"),
        (
            {"chunk_size": 10, "chunk_overlap": 5, "min_chunk_size": 6},
            "cannot exceed",
        ),
        ({"method": "custom", "custom_delimiter": ""}, "custom_delimiter"),
        ({"method": "semantic"}, "unsupported chunk method"),
        ({"method": "heading"}, "unsupported chunk method"),
    ],
)
def test_chunk_configuration_validation(kwargs: dict[str, object], message: str) -> None:
    with pytest.raises(ValueError, match=message):
        chunk_unstructured("some text", **kwargs)  # type: ignore[arg-type]


def test_quality_scoring_covers_all_dimensions_and_duplicates() -> None:
    valid = {
        "instruction": "如何修改收货地址？",
        "input": "订单尚未发货",
        "output": "可以在订单详情页申请修改收货地址。",
    }
    source = "订单尚未发货时，可以在订单详情页申请修改收货地址。"
    first_score = score_quality(valid, min_output_length=10, source_content=source)
    assert first_score.is_valid
    assert first_score.completeness == 100
    assert first_score.length == 100
    assert first_score.readability >= 90
    assert first_score.relevance >= 70
    assert first_score.duplicate == 100

    duplicate_score = score_quality(valid, known_fingerprints={first_score.fingerprint})
    assert duplicate_score.duplicate == 0
    assert "duplicate_record" in duplicate_score.flags

    unrelated_score = score_quality(
        valid,
        min_output_length=10,
        source_content="量子计算使用量子比特处理信息。",
    )
    assert unrelated_score.relevance < first_score.relevance
    assert "low_source_relevance" in unrelated_score.flags

    invalid_score = score_quality({"instruction": "", "output": "短"}, min_output_length=10)
    assert not invalid_score.is_valid
    assert {"missing_instruction", "output_too_short"}.issubset(invalid_score.flags)
    assert record_fingerprint(valid) == record_fingerprint(dict(reversed(list(valid.items()))))


def test_stable_split_is_reproducible_and_validates_ratios() -> None:
    first = stable_split("record-42", seed="task-1")
    assert stable_split("record-42", seed="task-1") == first
    assert first in {"train", "validation", "test"}
    assert stable_split("record-42", {"train": 100, "validation": 0, "test": 0}) == "train"
    with pytest.raises(ValueError, match="sum to 100"):
        stable_split("record", {"train": 80, "validation": 10, "test": 9})


def test_generate_standard_records_supports_json_qa_and_stable_variants() -> None:
    previews = [
        {
            "id": "preview-json",
            "edited_content": json.dumps(
                {"instruction": "问题", "input": "上下文", "output": "答案"},
                ensure_ascii=False,
            ),
        },
        {"id": "preview-qa", "editedContent": "问：如何操作？\n答：按步骤操作。"},
    ]
    records = generate_standard_records(
        previews,
        qa_pairs_per_item=2,
        semantic_enrichment=True,
        split={"train": 100, "validation": 0, "test": 0},
        split_seed="task-1",
    )
    assert len(records) == 4
    assert records[0]["instruction"] == "问题"
    assert records[0]["input"] == "上下文"
    assert records[0]["output"] == "答案"
    assert records[1]["instruction"].endswith("问题")
    assert records[2]["instruction"] == "如何操作?"
    assert records[2]["output"] == "按步骤操作。"
    assert all(record["status"] == "valid" for record in records)
    assert all(record["split"] == "train" for record in records)
    assert records == generate_standard_records(
        previews,
        qa_pairs_per_item=2,
        semantic_enrichment=True,
        split={"train": 100, "validation": 0, "test": 0},
        split_seed="task-1",
    )