big-rag-rus

"""
Генерация тестовых файлов (PNG + PDF) на английском и русском языках.
Запуск: python test-fixtures/generate_fixtures.py
"""

import os
from PIL import Image, ImageDraw, ImageFont
from fpdf import FPDF

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))


def get_font(size: int = 40) -> ImageFont.truetype:
    """Получить шрифт, поддерживающий кириллицу (Windows)."""
    candidates = [
        "C:/Windows/Fonts/arial.ttf",
        "C:/Windows/Fonts/calibri.ttf",
        "C:/Windows/Fonts/times.ttf",
        "C:/Windows/Fonts/verdana.ttf",
    ]
    for path in candidates:
        if os.path.exists(path):
            return ImageFont.truetype(path, size)
    return ImageFont.load_default()


def create_text_image(text: str, output_path: str) -> None:
    """Создать PNG-изображение с текстом на белом фоне."""
    font = get_font(48)
    tmp_img = Image.new("RGB", (100, 100), "white")
    tmp_draw = ImageDraw.Draw(tmp_img)
    bbox = tmp_draw.textbbox((0, 0), text, font=font)
    text_w = bbox[2] - bbox[0]
    text_h = bbox[3] - bbox[1]
    padding = 60
    img_w = text_w + padding * 2
    img_h = text_h + padding * 2
    img = Image.new("RGB", (img_w, img_h), "white")
    draw = ImageDraw.Draw(img)
    draw.text((padding, padding), text, fill="black", font=font)
    img.save(output_path, "PNG")
    print(f"  Created: {output_path} ({img_w}x{img_h})")


def create_text_pdf(text: str, output_path: str) -> None:
    """Создать PDF с текстом. Используем встроенный шрифт + Unicode."""
    pdf = FPDF()
    pdf.add_page()
    # Используем шрифт, поддерживающий кириллицу
    font_path = None
    for candidate in [
        "C:/Windows/Fonts/arial.ttf",
        "C:/Windows/Fonts/calibri.ttf",
        "C:/Windows/Fonts/times.ttf",
    ]:
        if os.path.exists(candidate):
            font_path = candidate
            break
    if font_path:
        pdf.add_font("TestFont", "", font_path, uni=True)
        pdf.set_font("TestFont", size=14)
    else:
        pdf.set_font("Helvetica", size=14)
    pdf.multi_cell(0, 10, text)
    pdf.output(output_path)
    print(f"  Created: {output_path}")


def create_image_pdf(text: str, output_path: str) -> None:
    """Создать PDF без текстового слоя — текст рендерится как изображение."""
    # Сначала создаём картинку с текстом
    font = get_font(36)
    tmp_img = Image.new("RGB", (100, 100), "white")
    tmp_draw = ImageDraw.Draw(tmp_img)
    bbox = tmp_draw.textbbox((0, 0), text, font=font)
    text_w = bbox[2] - bbox[0]
    text_h = bbox[3] - bbox[1]
    padding = 40
    img_w = text_w + padding * 2
    img_h = text_h + padding * 2
    img = Image.new("RGB", (img_w, img_h), "white")
    draw = ImageDraw.Draw(img)
    draw.text((padding, padding), text, fill="black", font=font)

    # Сохраняем как временный PNG и встраиваем в PDF как изображение
    tmp_png = output_path + ".tmp.png"
    img.save(tmp_png, "PNG")

    pdf = FPDF()
    pdf.add_page()
    # Встраиваем изображение — PDF не будет содержать текстового слоя
    pdf.image(tmp_png, x=10, y=10, w=img_w * 0.5)
    pdf.output(output_path)

    # Удаляем временный файл
    os.remove(tmp_png)
    print(f"  Created: {output_path} (image-only PDF, no text layer)")


def main():
    print("Generating test fixtures...\n")

    # === PNG images ===
    print("[PNG] Creating test images...")

    create_text_image(
        "Hello World OCR Test",
        os.path.join(SCRIPT_DIR, "sample_eng.png"),
    )

    create_text_image(
        "Привет мир тест распознавания",
        os.path.join(SCRIPT_DIR, "sample_rus.png"),
    )

    # === PDF files ===
    print("\n[PDF] Creating test PDFs...")

    create_text_pdf(
        "This is an English PDF document.\n"
        "It contains sample text for testing PDF parsing.\n"
        "The quick brown fox jumps over the lazy dog.",
        os.path.join(SCRIPT_DIR, "sample_eng.pdf"),
    )

    create_text_pdf(
        "Это русский PDF документ.\n"
        "Он содержит пример текста для тестирования парсинга PDF.\n"
        "Съешь ещё этих мягких французских булок, да выпей чаю.",
        os.path.join(SCRIPT_DIR, "sample_rus.pdf"),
    )

    # === Mixed language (EN+RU) files ===
    print("\n[MIX] Creating mixed-language test files...")

    # 1. Mixed image (EN + RU text)
    create_text_image(
        "Hello Привet World Мир OCR Test",
        os.path.join(SCRIPT_DIR, "sample_mix.png"),
    )

    # 2. Mixed PDF with text layer (selectable text)
    create_text_pdf(
        "This is a mixed language document.\n"
        "Это документ на двух языках.\n"
        "The quick brown fox jumps over the lazy dog.\n"
        "Съешь ещё этих мягких французских булок, да выпей чаю.",
        os.path.join(SCRIPT_DIR, "sample_mix_text.pdf"),
    )

    # 3. Mixed PDF without text layer (image-only — rendered as image in PDF)
    create_image_pdf(
        "Hello Привet World Мир Mixed OCR Test for bilingual document recognition",
        os.path.join(SCRIPT_DIR, "sample_mix_ocr.pdf"),
    )

    print("\nDone! All test fixtures generated.")


if __name__ == "__main__":
    main()