"""Проверка весов MuScriptor: заголовок safetensors, конфиг, число параметров, sha256.

Запуск:
    python tools/verify_model.py --model-dir muscriptor-large [--hash]

Ничего не загружает в память целиком: читает только 8-байтовую длину заголовка
и сам JSON-заголовок, поэтому работает и на 5 ГБ, и на 50 ГБ.
"""
import argparse
import hashlib
import json
import os
import struct
import sys

DTYPE_BYTES = {"F64": 8, "F32": 4, "F16": 2, "BF16": 2, "I64": 8, "I32": 4,
               "I16": 2, "I8": 1, "U8": 1, "BOOL": 1}


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


def read_header(path):
    with open(path, "rb") as handle:
        raw = handle.read(8)
        if len(raw) != 8:
            raise SystemExit("файл короче 8 байт: " + path)
        length = struct.unpack("<Q", raw)[0]
        if length <= 0 or length > 200 * 1024 * 1024:
            raise SystemExit("подозрительная длина заголовка: %d" % length)
        header = json.loads(handle.read(length))
    meta = header.pop("__metadata__", None)
    return header, meta, 8 + length


def sha256(path, chunk=8 * 1024 * 1024, label="sha256"):
    digest = hashlib.sha256()
    total = os.path.getsize(path)
    done = 0
    with open(path, "rb") as handle:
        while True:
            block = handle.read(chunk)
            if not block:
                break
            digest.update(block)
            done += len(block)
            if total:
                print("\r  %s: %5.1f%%" % (label, done * 100.0 / total), end="", flush=True)
    print("\r  %s: 100.0%%  %s" % (label, digest.hexdigest()))
    return digest.hexdigest()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", default="muscriptor-large")
    parser.add_argument("--hash", action="store_true", help="посчитать sha256 (долго)")
    args = parser.parse_args()

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

    print("[10%] читаю заголовок safetensors")
    tensors, meta, data_start = read_header(weights)
    file_size = os.path.getsize(weights)

    dtypes = {}
    total_bytes = 0
    for info in tensors.values():
        dtypes[info["dtype"]] = dtypes.get(info["dtype"], 0) + 1
        begin, end = info["data_offsets"]
        total_bytes += end - begin

    expected = data_start + total_bytes
    print("[30%%] файл: %s | данные: %s | заголовок: %s"
          % (human(file_size), human(total_bytes), human(data_start)))
    print("      тензоров: %d | типы: %s" % (len(tensors), dtypes))
    print("      параметров: %.1f млн" % (total_bytes / 4.0 / 1e6))

    if expected != file_size:
        print("[!] РАСХОЖДЕНИЕ: ожидалось %s, в файле %s" % (human(expected), human(file_size)))
        ok_size = False
    else:
        ok_size = True
        print("[50%] размеры сходятся байт в байт")

    config = {}
    if os.path.isfile(config_path):
        with open(config_path, encoding="utf-8") as handle:
            config = json.load(handle)
        print("[60%] config.json: " + json.dumps(config, ensure_ascii=False))
        dim = int(config.get("dim") or 0)
        layers = int(config.get("num_layers") or 0)
        if dim:
            probe = "linears.0.weight"
            if probe in tensors:
                out_features = tensors[probe]["shape"][0]
                print("      выходная голова: %s -> %d токенов" % (probe, out_features))
            hits = sum(1 for name in tensors if name.startswith("layers."))
            print("      тензоров в layers.*: %d (конфиг заявляет %d слоёв)" % (hits, layers))
    else:
        print("[!] нет config.json рядом с весами")

    if meta:
        print("[70%] метаданные: " + json.dumps(meta, ensure_ascii=False)[:400])

    exit_code = 0 if ok_size else 1
    if args.hash:
        print("[80%] считаю sha256 (это долго, ~1-2 минуты на 5 ГБ)")
        sha256(weights)
    else:
        print("[80%] sha256 пропущен (добавь --hash)")

    print("[100%] итог: " + ("ГОДНО" if exit_code == 0 else "ФАЙЛ ПОВРЕЖДЁН"))
    return exit_code


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