"""Упаковка MuScriptor в один самодостаточный архив.

Внутри: веса, конфиг, карточка модели, все инструменты и manifest с хешем.
Архив намеренно без сжатия: safetensors это уже плотные числа, gzip даёт ~2%
выигрыша за десятки минут CPU, а манифест и хеши считаются на лету.

Запуск:
    python tools/pack_bundle.py --model-dir muscriptor-large --out out/muscriptor-bundle.tar
"""
import argparse
import hashlib
import io
import json
import os
import struct
import sys
import tarfile
import time

BUNDLE_VERSION = "1.0"
CHUNK = 8 * 1024 * 1024


def human(size):
    for unit in ("Б", "КБ", "МБ", "ГБ", "ТБ"):
        if size < 1024 or unit == "ТБ":
            return "%.2f %s" % (size, unit)
        size /= 1024.0


def model_stats(path):
    """Размер, число тензоров и параметров — по заголовку, без чтения данных."""
    with open(path, "rb") as handle:
        length = struct.unpack("<Q", handle.read(8))[0]
        header = json.loads(handle.read(length))
    header.pop("__metadata__", None)
    params = 0
    dtypes = {}
    for info in header.values():
        begin, end = info["data_offsets"]
        dtypes[info["dtype"]] = dtypes.get(info["dtype"], 0) + 1
        params += (end - begin) // 4
    return {"tensors": len(header), "parameters": params, "dtypes": dtypes,
            "bytes": os.path.getsize(path)}


def add_with_progress(archive, source, arcname, digest, total, label):
    """Кладёт файл в tar потоком, попутно считая sha256 и печатая проценты."""
    size = os.path.getsize(source)
    info = tarfile.TarInfo(arcname)
    info.size = size
    info.mtime = int(time.time())
    with open(source, "rb") as handle:
        reader = _HashingReader(handle, digest, total, label)
        archive.addfile(info, reader)


class _HashingReader(io.RawIOBase):
    def __init__(self, handle, digest, total, label):
        self.handle, self.digest, self.total, self.label = handle, digest, total, label
        self.done = 0
        self.last = -1

    def readable(self):
        return True

    def read(self, size=-1):
        block = self.handle.read(size if size and size > 0 else CHUNK)
        if block:
            self.digest.update(block)
            self.done += len(block)
            if self.total:
                percent = int(self.done * 100 / self.total)
                if percent != self.last:
                    self.last = percent
                    print("\r  %s: %3d%%  %s / %s"
                          % (self.label, percent, human(self.done), human(self.total)),
                          end="", flush=True)
        elif self.total:
            print("\r  %s: 100%%  %s" % (self.label, human(self.total)), flush=True)
        return block


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", default="muscriptor-large")
    parser.add_argument("--out", default="out/muscriptor-large-bundle.tar")
    parser.add_argument("--tools", default=None, help="папка с инструментами (по умолчанию рядом с этим скриптом)")
    parser.add_argument("--license-note", default=None)
    args = parser.parse_args()

    weights = os.path.join(args.model_dir, "model.safetensors")
    config = os.path.join(args.model_dir, "config.json")
    if not os.path.isfile(weights):
        raise SystemExit("нет весов: " + weights)

    print("[5%] считаю параметры по заголовку")
    stats = model_stats(weights)
    print("      тензоров %d, параметров %.1f млн, типы %s, вес %s"
          % (stats["tensors"], stats["parameters"] / 1e6, stats["dtypes"], human(stats["bytes"])))

    tools_dir = args.tools or os.path.dirname(os.path.abspath(__file__))
    tools = sorted(name for name in os.listdir(tools_dir) if name.endswith(".py"))
    print("[10%] инструменты: " + ", ".join(tools) if tools else "[10%] инструменты не найдены")

    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    digest = hashlib.sha256()
    manifest = {
        "bundle_version": BUNDLE_VERSION,
        "built_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "model": {
            "name": "MuScriptor Large",
            "source": "https://huggingface.co/MuScriptor/muscriptor-large",
            "variant": "large",
            "tensors": stats["tensors"],
            "parameters": stats["parameters"],
            "dtypes": stats["dtypes"],
            "weights_bytes": stats["bytes"],
        },
        "license": {
            "weights": "CC-BY-NC-4.0",
            "code": "MIT",
            "note": "Веса некоммерческие: их нельзя использовать в платном продукте.",
        },
        "files": ["model/model.safetensors", "model/config.json", "model/README.md", "manifest.json"],
        "tools": tools,
    }

    print("[15%] пишу архив " + args.out)
    with tarfile.open(args.out, "w") as archive:
        for name in tools:
            path = os.path.join(tools_dir, name)
            info = tarfile.TarInfo("tools/" + name)
            info.size = os.path.getsize(path)
            info.mtime = int(time.time())
            with open(path, "rb") as handle:
                archive.addfile(info, handle)
            manifest["files"].append("tools/" + name)
        if args.license_note and os.path.isfile(args.license_note):
            data = open(args.license_note, "rb").read()
            info = tarfile.TarInfo("LICENSE_NOTICE.txt")
            info.size = len(data)
            info.mtime = int(time.time())
            archive.addfile(info, io.BytesIO(data))
            manifest["files"].append("LICENSE_NOTICE.txt")

        if os.path.isfile(config):
            data = open(config, "rb").read()
            info = tarfile.TarInfo("model/config.json")
            info.size = len(data)
            info.mtime = int(time.time())
            archive.addfile(info, io.BytesIO(data))
        card = os.path.join(args.model_dir, "README.md")
        if os.path.isfile(card):
            data = open(card, "rb").read()
            info = tarfile.TarInfo("model/README.md")
            info.size = len(data)
            info.mtime = int(time.time())
            archive.addfile(info, io.BytesIO(data))

        add_with_progress(archive, weights, "model/model.safetensors", digest,
                          stats["bytes"], "веса")

        manifest["model"]["sha256"] = digest.hexdigest()
        blob = json.dumps(manifest, ensure_ascii=False, indent=2).encode("utf-8")
        info = tarfile.TarInfo("manifest.json")
        info.size = len(blob)
        info.mtime = int(time.time())
        archive.addfile(info, io.BytesIO(blob))

    size = os.path.getsize(args.out)
    print("[100%%] готово: %s (%s)" % (args.out, human(size)))
    print("      sha256 весов: " + manifest["model"]["sha256"])
    print("      в истории: python tools/pack_bundle.py --model-dir ... --out ...")
    return 0


if __name__ == "__main__":
    sys.exit(main())