#!/usr/bin/env python3
"""Сборка публикационного дерева: код + доки + веса в fp16.

    python3 src/publish.py --out tmp/publish          # только собрать
    python3 src/publish.py --out tmp/publish --push   # собрать и залить (REPO_ID)

Локальные fp32-веса остаются на месте: они нужны обучению, а на Hub едет fp16.
"""
from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file

ROOT = Path(__file__).resolve().parent.parent
REPO_ID = "AiArtLab/Ra"

COPY = ["README.md", ".gitattributes", ".gitignore", "docs", "src",
        "vae/config.json", "vae/LICENSE", "vae/README.md", "vae/.gitattributes", "vae/example.py",
        "vae/ae", "text_encoder/config.json", "text_encoder/tokenizer.json",
        "text_encoder/tokenizer_config.json", "text_encoder/special_tokens_map.json",
        "text_encoder/README.md", "text_encoder/.gitattributes",
        "transformer/config.json"]


def half_file(src: Path, dst: Path, key_prefix: str = ""):
    tensors = load_file(str(src))
    out, overflow = {}, 0
    for k, v in tensors.items():
        name = f"{key_prefix}{k}" if key_prefix else k
        if v.is_floating_point():
            if v.numel() and float(v.abs().max()) > 65504:
                overflow += 1
            v = v.half()
        out[name] = v.contiguous()
        assert torch.isfinite(out[name]).all(), f"{name}: не-конечные значения после fp16"
    dst.parent.mkdir(parents=True, exist_ok=True)
    save_file(out, str(dst))
    print(f"[fp16] {dst.relative_to(dst.parents[1])}: {len(out)} тензоров, "
          f"{src.stat().st_size / 2**20:.0f} → {dst.stat().st_size / 2**20:.0f} МиБ"
          + (f", ВНИМАНИЕ: {overflow} тензоров вне диапазона fp16" if overflow else ""), flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=str(ROOT / "tmp" / "publish"))
    ap.add_argument("--repo", default=REPO_ID)
    ap.add_argument("--push", action="store_true", help="залить собранное на Hub (публичный репозиторий)")
    args = ap.parse_args()
    out = Path(args.out)
    if out.exists():
        shutil.rmtree(out)
    out.mkdir(parents=True)

    for rel in COPY:
        src = ROOT / rel
        dst = out / rel
        dst.parent.mkdir(parents=True, exist_ok=True)
        if src.is_dir():
            shutil.copytree(src, dst, ignore=shutil.ignore_patterns("__pycache__", ".ipynb_checkpoints"))
        else:
            shutil.copy2(src, dst)

    half_file(ROOT / "transformer" / "diffusion_pytorch_model.safetensors",
              out / "transformer" / "diffusion_pytorch_model.safetensors")
    half_file(ROOT / "vae" / "model.safetensors", out / "vae" / "model.safetensors")

    sd = torch.load(ROOT / "text_encoder" / "pytorch_model.bin", map_location="cpu", weights_only=True)
    sd = {k: v for k, v in sd.items() if not k.startswith(("head.", "decoder."))}   # MLM-голова не нужна
    save_file({k: v.half().contiguous() for k, v in sd.items()},
              str(out / "text_encoder" / "model.safetensors"))
    print(f"[fp16] text_encoder/model.safetensors: {len(sd)} тензоров "
          f"({sum(v.numel() for v in sd.values()) / 1e6:.1f}M), "
          f"{(out / 'text_encoder' / 'model.safetensors').stat().st_size / 2**20:.0f} МиБ", flush=True)

    cfg_path = out / "transformer" / "config.json"
    cfg = json.loads(cfg_path.read_text())
    cfg["storage_dtype"] = "float16"
    cfg_path.write_text(json.dumps(cfg, indent=2, ensure_ascii=False))

    total = sum(f.stat().st_size for f in out.rglob("*") if f.is_file())
    print(f"[готово] {out}: {total / 2**30:.2f} ГиБ", flush=True)

    if args.push:
        from huggingface_hub import HfApi
        api = HfApi()
        api.create_repo(args.repo, repo_type="model", private=False, exist_ok=True)
        api.upload_folder(folder_path=str(out), repo_id=args.repo, repo_type="model",
                          commit_message=f"Ra fp16 (step {cfg['step']})")
        print(f"[push] https://huggingface.co/{args.repo}", flush=True)


if __name__ == "__main__":
    main()