# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Inspect and convert .forge weight repositories (see src/core/fmodel.h for
the format: git-style manifests + content-addressed page-aligned shards).

Usage:
  python3 tools/fmodel.py inspect  model.forge [--manifest manifest-000002.json]
  python3 tools/fmodel.py log      model.forge                 # manifest history
  python3 tools/fmodel.py to-safetensors model.forge out.safetensors
        [--manifest m.json]        # pure-numpy writer, no torch needed

Only numpy is required. The safetensors export makes .forge weights loadable
from PyTorch/HF (`safetensors.torch.load_file`) for interop.
"""
import argparse
import json
import os
import struct

import numpy as np

PAGE = 16384
NP_DTYPE = {"f32": np.float32, "f16": np.float16, "u16": np.uint16, "i32": np.int32}
ST_DTYPE = {"f32": "F32", "f16": "F16", "bf16": "BF16", "u16": "U16", "i32": "I32"}


def resolve_manifest(repo, manifest=None):
    if os.path.isfile(repo):  # given a manifest path directly
        return os.path.dirname(os.path.dirname(os.path.abspath(repo))) \
            if os.path.basename(os.path.dirname(repo)) == "manifests" \
            else os.path.dirname(os.path.abspath(repo)), repo
    path = os.path.join(repo, "manifests", manifest) if manifest else \
        os.path.join(repo, "manifest-latest.json")
    return repo, path


def load(repo, manifest=None):
    repo, mpath = resolve_manifest(repo, manifest)
    with open(mpath) as f:
        return repo, json.load(f)


def tensor_bytes(repo, entry):
    with open(os.path.join(repo, entry["shard"]), "rb") as f:
        f.seek(entry["offset"])
        return f.read(entry["nbytes"])


def tensor_array(repo, entry):
    raw = tensor_bytes(repo, entry)
    dt = entry["dtype"]
    if dt == "bf16":  # numpy has no bf16: widen via bit tricks
        u16 = np.frombuffer(raw, dtype=np.uint16).astype(np.uint32) << 16
        return u16.view(np.float32).reshape(entry["shape"])
    return np.frombuffer(raw, dtype=NP_DTYPE[dt]).reshape(entry["shape"])


def cmd_inspect(args):
    repo, m = load(args.repo, args.manifest)
    cfg = m.get("config", {}).get("model", {})
    print(f"{m.get('self', '?')}  step={m.get('step')}  tag='{m.get('tag', '')}'"
          f"  dtype={m.get('dtype')}  parent='{m.get('parent', '')}'")
    if cfg:
        print(f"model: {cfg.get('name')}  layers={cfg.get('n_layers')}"
              f"  d_model={cfg.get('d_model')}  vocab={cfg.get('vocab_size')}")
    total = 0
    shards = {}
    for name, e in sorted(m["tensors"].items()):
        total += e["nbytes"]
        shards.setdefault(e["shard"], 0)
        shards[e["shard"]] += e["nbytes"]
        print(f"  {name:<40} {e['dtype']:<5} {str(e['shape']):<20}"
              f" {e['nbytes'] / 1e6:8.2f} MB  {e['shard'].split('/')[-1][:8]}…")
    print(f"{len(m['tensors'])} tensors, {total / 1e6:.1f} MB across "
          f"{len(shards)} shard(s)")


def cmd_log(args):
    repo, _ = resolve_manifest(args.repo)
    mdir = os.path.join(repo, "manifests")
    for name in sorted(os.listdir(mdir)):
        with open(os.path.join(mdir, name)) as f:
            m = json.load(f)
        shards = {e["shard"] for e in m["tensors"].values()}
        print(f"{name}  step={m.get('step'):>8}  tag='{m.get('tag', '')}'"
              f"  dtype={m.get('dtype')}  shards={len(shards)}"
              f"  parent='{m.get('parent', '')}'")


def cmd_to_safetensors(args):
    repo, m = load(args.repo, args.manifest)
    header = {"__metadata__": {"format": "forge",
                               "step": str(m.get("step", -1)),
                               "config": json.dumps(m.get("config", {}))}}
    offset = 0
    order = sorted(m["tensors"].items())
    for name, e in order:
        header[name] = {"dtype": ST_DTYPE[e["dtype"]], "shape": e["shape"],
                        "data_offsets": [offset, offset + e["nbytes"]]}
        offset += e["nbytes"]
    hjson = json.dumps(header, separators=(",", ":")).encode()
    pad = (8 - len(hjson) % 8) % 8  # safetensors aligns the header to 8
    hjson += b" " * pad
    with open(args.out, "wb") as f:
        f.write(struct.pack("<Q", len(hjson)))
        f.write(hjson)
        for name, e in order:
            f.write(tensor_bytes(repo, e))
    print(f"wrote {args.out} ({(8 + len(hjson) + offset) / 1e6:.1f} MB, "
          f"{len(order)} tensors)")


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    for name, fn in [("inspect", cmd_inspect), ("log", cmd_log),
                     ("to-safetensors", cmd_to_safetensors)]:
        p = sub.add_parser(name)
        p.add_argument("repo", help=".forge directory or manifest json")
        p.add_argument("--manifest", default=None,
                       help="specific manifest inside manifests/")
        if name == "to-safetensors":
            p.add_argument("out")
        p.set_defaults(fn=fn)
    args = ap.parse_args()
    args.fn(args)


if __name__ == "__main__":
    main()