#!/usr/bin/env python3
"""hf-manifest.py — generate an hf-fetch.sh manifest from a HuggingFace repo.

Usage: python3 scripts/hf-manifest.py <owner/repo> [outfile]

Lists the repo tree via the HF API and, for each file, records the size and
the LFS sha256 (the X-Linked-ETag header on the resolve endpoint). Non-LFS
files (config.json, tokenizer, ...) have no LFS sha256, so they get a blank
hash field — hf-fetch.sh then checks size only for them.

Output: one line per file, three space-separated fields:
    <filename> <size_bytes> <sha256-or-blank>

Token: $HF_TOKEN, else ~/.cache/huggingface/token, else the llamaswap conf
(grep hf_...). Gated repos need a token that has been granted access.
"""
import json
import os
import re
import sys
import urllib.error
import urllib.request


def get_token():
    t = os.environ.get("HF_TOKEN", "")
    if not t:
        p = os.path.expanduser("~/.cache/huggingface/token")
        if os.path.exists(p):
            with open(p) as f:
                t = f.read().strip()
    if not t:
        p = os.path.expanduser("~/.config/llamaswap.conf.yaml")
        if os.path.exists(p):
            with open(p) as f:
                m = re.search(r"hf_[a-zA-Z0-9]+", f.read())
                if m:
                    t = m.group(0)
    return t


def head(url, token):
    h = {"User-Agent": "hf-manifest/1.0"}
    if token:
        h["Authorization"] = f"Bearer {token}"
    req = urllib.request.Request(url, method="HEAD", headers=h)
    return urllib.request.urlopen(req, timeout=30)


def main():
    if len(sys.argv) < 2:
        print(f"usage: {sys.argv[0]} <owner/repo> [outfile]", file=sys.stderr)
        sys.exit(2)
    repo = sys.argv[1]
    outfile = sys.argv[2] if len(sys.argv) > 2 else None
    token = get_token()
    # 1. List the repo tree (siblings[].rfilename). Sizes and LFS sha256s
    # come from the resolve endpoint's HEAD (Content-Length + X-Linked-ETag);
    # the API's siblings carry only rfilename.
    api = f"https://huggingface.co/api/models/{repo}"
    req = urllib.request.Request(
        api, headers={"User-Agent": "hf-manifest/1.0",
                      **({"Authorization": f"Bearer {token}"} if token else {})})
    with urllib.request.urlopen(req, timeout=30) as r:
        meta = json.load(r)
    siblings = meta.get("siblings", [])
    print(f"# {repo}: {len(siblings)} entries", file=sys.stderr)

    lines = []
    for s in siblings:
        name = s.get("rfilename", "")
        if not name or name.endswith("/"):
            continue
        # 2. Size (Content-Length) + LFS sha256 (X-Linked-ETag) from the
        # resolve endpoint's HEAD. Non-LFS files have no 64-hex etag, so the
        # hash field stays blank (hf-fetch.sh checks size only for them).
        size = 0
        sha = ""
        resolve = f"https://huggingface.co/{repo}/resolve/main/{name}"
        try:
            with head(resolve, token) as r:
                size = int(r.headers.get("Content-Length") or 0)
                etag = (r.headers.get("X-Linked-ETag")
                        or r.headers.get("ETag") or "").strip('"')
                if re.fullmatch(r"[0-9a-f]{64}", etag):
                    sha = etag
        except urllib.error.HTTPError as e:
            if e.code != 404:
                raise
        lines.append(f"{name} {size} {sha}")
        print(f"  {name}  {size:>12}  {sha or '(size only)'}", file=sys.stderr)

    out = "\n".join(lines) + "\n"
    if outfile:
        with open(outfile, "w") as f:
            f.write(out)
        print(f"wrote {outfile}", file=sys.stderr)
    else:
        sys.stdout.write(out)


if __name__ == "__main__":
    main()