#!/usr/bin/env python3
"""Пересохраняет чекпоинт FANTA с разбиением весов на шарды.

Примеры:
  python3 scripts/reshard_checkpoint.py runs/checkpoint-5000 runs/checkpoint-5000-sharded
  python3 scripts/reshard_checkpoint.py runs/checkpoint-5000 out --max_shard_size 1GB

Результат:
  out/
    config.json
    model-00001-of-00008.safetensors
    ...
    model-00008-of-00008.safetensors
    model.safetensors.index.json

HF (`from_pretrained`) автоматически читает index.json и собирает шарды — склеивать вручную не нужно.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

# Позволяем запускать из корня проекта без PYTHONPATH: добавляем родительский каталог.
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

from fanta.modeling_fanta import FantaForPredictionAndAnomaly  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("src", help="Папка исходного чекпоинта (с model.safetensors + config.json)")
    ap.add_argument("dst", help="Папка куда сохранить шардированный чекпоинт")
    ap.add_argument(
        "--max_shard_size",
        default="1GB",
        help="Размер одного шарда (строка для HF). ~1GB → ~8 шардов для 8GB модели.",
    )
    args = ap.parse_args()

    src = Path(args.src)
    dst = Path(args.dst)
    if not (src / "config.json").is_file():
        print(f"Нет config.json в {src}", file=sys.stderr)
        sys.exit(1)

    print(f"Загружаю {src} …")
    model = FantaForPredictionAndAnomaly.from_pretrained(src)

    dst.mkdir(parents=True, exist_ok=True)
    print(f"Сохраняю шарды max_shard_size={args.max_shard_size} → {dst}")
    model.save_pretrained(
        dst,
        safe_serialization=True,
        max_shard_size=args.max_shard_size,
    )

    shards = sorted(dst.glob("model-*.safetensors"))
    if not shards:
        shards = sorted(dst.glob("*.safetensors"))
    total = sum(p.stat().st_size for p in shards)
    print(f"Готово. Файлов весов: {len(shards)}, суммарно {total / 1024**3:.2f} GiB")
    for p in shards:
        print(f"  {p.name}  {p.stat().st_size / 1024**3:.2f} GiB")


if __name__ == "__main__":
    main()