#!/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 content sha256. LFS and Xet-backed files redirect the resolve endpoint to
a CDN (302), whose own headers carry the truth: X-Linked-Size is the byte
count and X-Linked-ETag is the sha256 of the file's bytes. The redirect must
NOT be followed: on Xet-backed repos the bridge answers with ETag = the Xet
CAS chunk-dag hash, which is NOT the content sha256 (seen on
Qwen/Qwen3.8-Flash-Next-FP8, 2026-09-26: every shard "failed" verification
against a hash no download could ever match). Non-LFS files (config.json,
tokenizer, ...) are answered directly (200) and carry no 64-hex etag, 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 http.client
import json
import os
import re
import sys
import urllib.request
from urllib.parse import urljoin, urlsplit

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):
    """HEAD via http.client, which does not follow redirects.

    urllib's default opener follows the resolve endpoint's 302, and the Xet
    bridge's ETag is the Xet CAS hash, not the content sha256 — so the 302's
    own X-Linked-* headers must be read instead. Returns an http.client
    response; the caller reads r.status and r.headers.
    """
    parts = urlsplit(url)
    conn = http.client.HTTPSConnection(parts.netloc, timeout=30)
    path = parts.path + (f"?{parts.query}" if parts.query else "")
    h = {"User-Agent": "hf-manifest/1.0"}
    if token:
        h["Authorization"] = f"Bearer {token}"
    conn.request("HEAD", path, headers=h)
    return conn.getresponse()


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 content sha256s
    # come from the resolve endpoint's HEAD (not followed — see head());
    # 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 sha256 from the resolve endpoint's HEAD. 200 =
        # direct answer (Content-Length + ETag). 3xx WITH X-Linked-Size =
        # LFS/Xet redirect: X-Linked-Size + X-Linked-ETag are the byte
        # count and the sha256 of the file's bytes, and the redirect must
        # NOT be followed — the Xet bridge's own ETag is the Xet CAS hash,
        # not the content sha256. 3xx WITHOUT X-Linked-Size = non-LFS
        # redirect (307 to HF's own /api/resolve-cache): the stub's
        # Content-Length is NOT the file size, so follow ONCE; the final
        # ETag there is the git blob sha (40 hex), so the 64-hex gate
        # leaves the hash blank (hf-fetch.sh then checks size only).
        size = 0
        sha = ""
        resolve = f"https://huggingface.co/{repo}/resolve/main/{name}"
        try:
            with head(resolve, token) as r:
                if r.status == 404:
                    raise FileNotFoundError(name)
                if r.status == 200:
                    size = int(r.headers.get("Content-Length") or 0)
                    etag = (r.headers.get("ETag") or "").strip('"')
                else:
                    linked = r.headers.get("X-Linked-Size")
                    if linked is not None:
                        size = int(linked)
                        etag = (r.headers.get("X-Linked-ETag") or "").strip('"')
                    else:
                        loc = r.headers.get("Location", "")
                        with head(urljoin(resolve, loc), token) as r2:
                            if r2.status == 200:
                                size = int(r2.headers.get("Content-Length") or 0)
                                etag = (r2.headers.get("ETag") or "").strip('"')
                            else:
                                etag = ""
                if re.fullmatch(r"[0-9a-f]{64}", etag):
                    sha = etag
        except FileNotFoundError:
            pass
        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()