setup_counter_data.py

#!/usr/bin/env python3
"""Prepare counter lookup data: reorganize pictures and build counter_config.xlsx."""

from __future__ import annotations

import re
import shutil
import zipfile
from pathlib import Path

import openpyxl
from openpyxl.utils import get_column_letter

SCRIPT_DIR = Path(__file__).resolve().parent
TEMPLATES_DIR = SCRIPT_DIR / "templates"
PICTURES_DIR = TEMPLATES_DIR / "pictures"
OUTPUT_XLSX = SCRIPT_DIR / "counter_config.xlsx"

MONTAGE_TO_TIPORAZMER: dict[str, str] = {
    "СМТ-Смарт-К_G4-В": "4",
    "СМТ-Смарт-К_G6-В": "6",
    "СМТ-Смарт-К_G10-В": "10",
    "СМТ-Комплекс-К_G16": "16",
    "СМТ-Комплекс-К_G25": "25",
    "СМТ-Комплекс_-G40": "40",
    "СМТ-Комплекс_-G65": "65",
}

TIPORAZMER_ORDER = ["4", "6", "10", "16", "25", "40", "65"]

DIRECTIONS = [
    "слева направо",
    "сверху вниз",
    "справа налево",
]

COLUMNS = [
    "Типоразмер",
    "Направление",
    "Счетчик газа после замены",
    "Счетчик газа после замены полностью",
    "Чертеж",
]


def direction_to_filename(direction: str) -> str:
    return direction.replace(" ", "_")


def filename_to_direction(filename_stem: str) -> str:
    return filename_stem.replace("_", " ")


def extract_template_text(docx_path: Path) -> str:
    with zipfile.ZipFile(docx_path) as archive:
        xml = archive.read("word/document.xml").decode("utf-8")
    text = re.sub(r"</w:p>", "\n", xml)
    text = re.sub(r"<[^>]+>", "", text)
    return text


def extract_counter_values(template_path: Path) -> tuple[str, str]:
    text = extract_template_text(template_path)
    full_match = re.search(
        r"на интеллектуальный прибор учета(?: газа)?\s+(.+?)\.\s",
        text,
        flags=re.IGNORECASE | re.DOTALL,
    )
    short_match = re.search(
        r"установку счетчика газа\s+(.+?)\.\s",
        text,
        flags=re.IGNORECASE,
    )
    if not full_match or not short_match:
        raise ValueError(f"Could not extract counter values from {template_path.name}")
    return short_match.group(1).strip(), full_match.group(1).strip()


def load_counter_values_by_tiporazmer() -> dict[str, tuple[str, str]]:
    values: dict[str, tuple[str, str]] = {}
    for tiporazmer in TIPORAZMER_ORDER:
        template_path = TEMPLATES_DIR / f"template_{tiporazmer}.docx"
        values[tiporazmer] = extract_counter_values(template_path)
    return values


def delete_picture_pdfs() -> int:
    removed = 0
    for path in PICTURES_DIR.rglob("*.pdf"):
        path.unlink()
        removed += 1
    return removed


def reorganize_pictures() -> dict[tuple[str, str], Path]:
    picture_map: dict[tuple[str, str], Path] = {}

    for old_folder in sorted(PICTURES_DIR.iterdir()):
        if not old_folder.is_dir():
            continue

        tiporazmer = MONTAGE_TO_TIPORAZMER.get(old_folder.name)
        if tiporazmer is None:
            raise ValueError(f"Unknown montage folder: {old_folder.name}")

        new_folder = PICTURES_DIR / tiporazmer
        new_folder.mkdir(parents=True, exist_ok=True)

        for file_path in sorted(old_folder.iterdir()):
            if file_path.suffix.lower() == ".pdf":
                continue
            if file_path.suffix.lower() not in {".jpeg", ".jpg"}:
                continue

            direction_part = file_path.stem.split("__", 1)[-1]
            dest_name = f"{direction_part}{file_path.suffix.lower()}"
            dest_path = new_folder / dest_name
            shutil.move(str(file_path), str(dest_path))

            direction = filename_to_direction(direction_part)
            picture_map[(tiporazmer, direction)] = dest_path

        old_folder.rmdir()

    return picture_map


def build_picture_url(path: Path | None) -> str:
    if path is None or not path.exists():
        return ""
    return path.relative_to(SCRIPT_DIR).as_posix()


def write_counter_config(
    counter_values: dict[str, tuple[str, str]],
    picture_map: dict[tuple[str, str], Path],
) -> None:
    workbook = openpyxl.Workbook()
    sheet = workbook.active
    sheet.title = "counter_config"

    for column_index, header in enumerate(COLUMNS, start=1):
        sheet.cell(row=1, column=column_index, value=header)

    row_index = 2
    for tiporazmer in TIPORAZMER_ORDER:
        short_value, full_value = counter_values[tiporazmer]
        for direction in DIRECTIONS:
            picture_path = picture_map.get((tiporazmer, direction))
            sheet.cell(row=row_index, column=1, value=int(tiporazmer))
            sheet.cell(row=row_index, column=2, value=direction)
            sheet.cell(row=row_index, column=3, value=short_value)
            sheet.cell(row=row_index, column=4, value=full_value)
            sheet.cell(row=row_index, column=5, value=build_picture_url(picture_path))
            row_index += 1

    sheet.auto_filter.ref = f"A1:{get_column_letter(len(COLUMNS))}{row_index - 1}"
    sheet.freeze_panes = "A2"
    workbook.save(OUTPUT_XLSX)


def main() -> None:
    if not PICTURES_DIR.exists():
        raise FileNotFoundError(f"Pictures folder not found: {PICTURES_DIR}")

    removed_pdfs = delete_picture_pdfs()
    picture_map = reorganize_pictures()
    counter_values = load_counter_values_by_tiporazmer()
    write_counter_config(counter_values, picture_map)

    print(f"Removed {removed_pdfs} PDF file(s) from {PICTURES_DIR.name}/")
    print(f"Reorganized JPEG files into tiporazmer folders under {PICTURES_DIR.name}/")
    print(f"Created {OUTPUT_XLSX.name} with {len(TIPORAZMER_ORDER) * len(DIRECTIONS)} rows")


if __name__ == "__main__":
    main()