"""Упаковка моделей для Docker: models/best + models/segmenter -> models-bundle.tar.gz.

    python tools/pack_models.py                       # models-bundle.tar.gz в корне проекта
    python tools/pack_models.py --out dist/m.tar.gz   # другой путь
    python tools/pack_models.py --stage docker-models # только каталог docker-models/models (для docker build)

В архив попадает только нужное сервису: веса, токенизатор, label_maps, конфиги. Графики, eval.json,
архивные копии и запуски обучения не берутся. У старых чекпоинтов без config.json конфиг основы
берётся из кэша Hugging Face (нужен один раз на машине, где лежит модель) — после этого сервис
работает полностью офлайн.

Распаковка: tar -xzf models-bundle.tar.gz   (даёт каталог models/)
"""

from __future__ import annotations

import argparse
import shutil
import sys
import tarfile
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

MAIN_FILES = ("model.pt", "dub_config.json", "label_maps.json", "tokenizer.json", "tokenizer_config.json",
              "config.json", "special_tokens_map.json", "vocab.txt")
SEG_FILES = ("config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json",
             "special_tokens_map.json", "vocab.txt")


def stage(dest: Path, main: Path, seg: Path) -> list[str]:
    """Собирает dest/models/{best,segmenter}; возвращает список предупреждений."""
    notes: list[str] = []
    if not (main / "model.pt").exists():
        raise SystemExit(f"нет {main / 'model.pt'}: сначала обучите и выберите модель (tools/collect.py, пункт 7)")
    out_main, out_seg = dest / "models" / "best", dest / "models" / "segmenter"
    out_main.mkdir(parents=True)
    for name in MAIN_FILES:
        if (main / name).exists():
            shutil.copy2(main / name, out_main / name)
    if not (out_main / "config.json").exists():
        import json

        from transformers import AutoConfig
        backbone = json.loads((main / "dub_config.json").read_text(encoding="utf-8"))["backbone"]
        AutoConfig.from_pretrained(backbone).save_pretrained(out_main)
        notes.append(f"config.json основы {backbone} взят из кэша Hugging Face")
    if seg.exists():
        out_seg.mkdir(parents=True)
        for name in SEG_FILES:
            if (seg / name).exists():
                shutil.copy2(seg / name, out_seg / name)
    else:
        notes.append(f"нет {seg}: сервис будет сегментировать страницы правилами, а не моделью")
    return notes


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", default=str(ROOT / "models" / "best"))
    ap.add_argument("--segmenter", default=str(ROOT / "models" / "segmenter"))
    ap.add_argument("--out", default=str(ROOT / "models-bundle.tar.gz"))
    ap.add_argument("--stage", help="только подготовить каталог <stage>/models и не архивировать")
    args = ap.parse_args()

    if args.stage:
        dest = Path(args.stage)
        shutil.rmtree(dest / "models", ignore_errors=True)
        dest.mkdir(parents=True, exist_ok=True)
        for n in stage(dest, Path(args.model), Path(args.segmenter)):
            print("•", n)
        (dest / "models" / ".gitkeep").touch()           # каталог остаётся в git пустым
        print(f"готово: {dest / 'models'}")
        return
    with tempfile.TemporaryDirectory() as tmp:
        for n in stage(Path(tmp), Path(args.model), Path(args.segmenter)):
            print("•", n)
        out = Path(args.out)
        out.parent.mkdir(parents=True, exist_ok=True)
        with tarfile.open(out, "w:gz") as tar:
            tar.add(Path(tmp) / "models", arcname="models")
    print(f"готово: {out} ({out.stat().st_size / 1e6:.0f} МБ)")


if __name__ == "__main__":
    main()