test_sample_run.py

#!/usr/bin/env python3
"""Generate one sample PDF per template from the first matching input row.

For each template (by Типоразмер), picks the first row in input_data.xlsx
that matches and exports a PDF.

Outputs go to test_sample_run/{generated,pdf,errors}/.
PDF files are named from column A (№ по договору) as 4-digit numbers, e.g. 0001.pdf.
"""

from __future__ import annotations

import argparse
import shutil
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from edit_template import (
    CLOUD_SAMPLES_DIR,
    INPUT_XLSX,
    TEMPLATES_DIR,
    WordSession,
    copy_pdfs_to_cloud,
    count_pdf_pages,
    convert_docx_to_pdf,
    discover_tiporazmer_sizes,
    ensure_docx_template,
    enrich_context,
    format_duration,
    get_general_template_path,
    get_template_pdf_page_count,
    load_counter_config,
    render_row,
)
from test_templates_run import load_rows_by_tiporazmer

SCRIPT_DIR = Path(__file__).resolve().parent
OUTPUT_ROOT = SCRIPT_DIR / "test_sample_run"
GENERATED_DIR = OUTPUT_ROOT / "generated"
PDF_DIR = OUTPUT_ROOT / "pdf"
ERRORS_DIR = OUTPUT_ROOT / "errors"
INDEX_PATH = OUTPUT_ROOT / "index.txt"


@dataclass(frozen=True)
class SampleRow:
    tiporazmer: str
    excel_row: int
    output_stem: str
    context: dict[str, Any]


def select_first_rows(
    template_sizes: list[str],
    grouped_rows: dict[str, list[tuple[int, dict[str, Any], str]]],
) -> list[SampleRow]:
    selected: list[SampleRow] = []

    for tiporazmer in template_sizes:
        rows = grouped_rows.get(tiporazmer, [])
        if not rows:
            print(f"Warning: no input rows for template_{tiporazmer}", file=sys.stderr)
            continue

        excel_row, context, output_stem = rows[0]
        selected.append(
            SampleRow(
                tiporazmer=tiporazmer,
                excel_row=excel_row,
                output_stem=output_stem,
                context=context,
            )
        )

    return selected


def clean_output_dirs() -> None:
    for directory in (GENERATED_DIR, PDF_DIR, ERRORS_DIR):
        directory.mkdir(parents=True, exist_ok=True)
        for path in directory.iterdir():
            if path.is_dir():
                shutil.rmtree(path)
            else:
                path.unlink()


def build_output_stem(sample: SampleRow) -> str:
    return sample.output_stem


def move_to_errors(pdf_path: Path, docx_path: Path) -> None:
    ERRORS_DIR.mkdir(parents=True, exist_ok=True)
    shutil.move(str(pdf_path), str(ERRORS_DIR / pdf_path.name))
    if docx_path.exists():
        shutil.move(str(docx_path), str(ERRORS_DIR / docx_path.name))


def write_index(samples: list[SampleRow], page_errors: list[str]) -> None:
    lines = [
        "tiporazmer\tdogovor\texcel_row\toutput_pdf\tплощадка\tпокупатель\tсчетчик_газа_до_замены\n"
    ]
    for sample in samples:
        output_name = f"{build_output_stem(sample)}.pdf"
        if output_name in page_errors:
            output_name = f"errors/{output_name}"
        context = sample.context
        lines.append(
            f"{sample.tiporazmer}\t{sample.output_stem}\t{sample.excel_row}\t{output_name}\t"
            f"{str(context.get('Площадка', '')).replace(chr(9), ' ')}\t"
            f"{str(context.get('Покупатель', '')).replace(chr(9), ' ')}\t"
            f"{str(context.get('Счетчик газа до замены', '')).replace(chr(9), ' ')}\n"
        )
    INDEX_PATH.write_text("".join(lines), encoding="utf-8")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--input",
        type=Path,
        default=INPUT_XLSX,
        help=f"Excel input file (default: {INPUT_XLSX.name})",
    )
    parser.add_argument(
        "--templates-dir",
        type=Path,
        default=TEMPLATES_DIR,
        help=f"Templates folder (default: {TEMPLATES_DIR.name}/)",
    )
    return parser.parse_args()


def run(input_path: Path, templates_dir: Path) -> int:
    if not input_path.exists():
        print(f"Error: input file not found: {input_path}", file=sys.stderr)
        return 1

    if not templates_dir.exists():
        print(f"Error: templates folder not found: {templates_dir}", file=sys.stderr)
        return 1

    template_sizes = discover_tiporazmer_sizes()
    if not template_sizes:
        print(f"Error: no templates found in {templates_dir}", file=sys.stderr)
        return 1

    grouped_rows = load_rows_by_tiporazmer(input_path)
    samples = select_first_rows(template_sizes, grouped_rows)
    if not samples:
        print("No sample documents to generate.", file=sys.stderr)
        return 1

    clean_output_dirs()

    counter_lookup = load_counter_config()
    template_path = ensure_docx_template(get_general_template_path())

    started_at = time.perf_counter()
    page_errors: list[str] = []

    with WordSession() as word_session:
        word_app = word_session.app if word_session.active else None
        converter = "Microsoft Word" if word_session.active else "LibreOffice"
        print(f"Using {converter} for PDF export.")
        print(f"Generating {len(samples)} sample document(s) for {len(template_sizes)} template(s).")

        template_pages = get_template_pdf_page_count(template_path, word_app=word_app)

        for sample in samples:
            file_stem = build_output_stem(sample)
            output_docx = GENERATED_DIR / f"{file_stem}.docx"
            output_pdf = PDF_DIR / f"{file_stem}.pdf"

            enriched_context = enrich_context(sample.context, counter_lookup)
            render_row(template_path, enriched_context, output_docx)
            convert_docx_to_pdf(output_docx, output_pdf, word_app=word_app)

            generated_pages = count_pdf_pages(output_pdf)
            if generated_pages > template_pages:
                move_to_errors(output_pdf, output_docx)
                page_errors.append(f"{file_stem}.pdf")
                print(
                    f"Page overflow: {file_stem}.pdf has {generated_pages} page(s), "
                    f"template has {template_pages} — moved to errors/",
                    file=sys.stderr,
                )

    write_index(samples, page_errors)

    copied_to_cloud = copy_pdfs_to_cloud(PDF_DIR, CLOUD_SAMPLES_DIR)

    elapsed = time.perf_counter() - started_at
    ok_count = len(samples) - len(page_errors)
    print(f"Done. Generated {ok_count} PDF(s), {len(page_errors)} moved to errors/.")
    print(f"Copied {len(copied_to_cloud)} PDF(s) to {CLOUD_SAMPLES_DIR}")
    print(f"Index: {INDEX_PATH}")
    print(f"Total time: {format_duration(elapsed)}")
    return 0


def main() -> int:
    args = parse_args()
    return run(args.input, args.templates_dir)


if __name__ == "__main__":
    raise SystemExit(main())