# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb", "huggingface_hub"]
# ///
"""Weekly updater for the 'Who's Shipping Open-Source AI?' Space.

Runs on HF Jobs ~12h after hfmlsoc/hub_weekly_snapshots adds a new week
(snapshots land Saturday ~07:30 UTC). Appends the new week(s) to data_ui.js.
State lives entirely in the Space repo: data_ui.js (payload) + org_table.json.
"""
import duckdb, json, os, sys, time, collections
import urllib.request, urllib.parse
from huggingface_hub import HfApi

SPACE = os.environ.get("TRACKER_SPACE", "huggingface/whos-shipping-open-source-ai")
DS = "hfmlsoc/hub_weekly_snapshots"
BASE = f"https://huggingface.co/datasets/{DS}/resolve/main"
TOKEN = os.environ.get("HF_TOKEN")
api = HfApi(token=TOKEN)

# ---------- load state ----------
pj = api.hf_hub_download(SPACE, "data_ui.js", repo_type="space")
payload = json.loads(open(pj).read().split("=", 1)[1].rstrip().rstrip(";"))
orgs = json.load(open(api.hf_hub_download(SPACE, "org_table.json", repo_type="space")))
org2comp = {e["org"]: e["company"] for e in orgs}
inlist = ",".join(f"'{o}'" for o in org2comp)
weeks = payload["weeks"]
companies = payload["companies"]

files = api.list_repo_files(DS, repo_type="dataset")
have = {t: sorted(f.split("/")[1] for f in files if f.startswith(t + "/") and f.endswith(".parquet"))
        for t in ["models", "datasets", "spaces"]}
new_dates = [d for d in have["models"] if d > weeks[-1]
             and d in have["datasets"] and d in have["spaces"]]
if not new_dates:
    print("no new snapshot weeks; latest =", weeks[-1]); sys.exit(0)
print("new weeks:", new_dates)
latest = new_dates[-1]
con = duckdb.connect()

def cols_of(url):
    return {r[0] for r in con.sql(f"DESCRIBE SELECT * FROM '{url}'").fetchall()}

# ---------- link tables from the latest snapshot ----------
murl = f"{BASE}/models/{latest}/models.parquet"
rows = con.sql(f"""SELECT m.id, m.author, m.baseModels.relation, [x.id for x in m.baseModels.models]
    FROM '{murl}' m WHERE m.baseModels.relation IS NOT NULL""").fetchall()
deriv_links, ext_ids = [], []
for mid, author, rel, bases in rows:
    basecomps = {org2comp.get(b.split("/")[0], b.split("/")[0]) for b in (bases or []) if "/" in b}
    if author in org2comp:
        ext_ids.append((mid, "ext" if (basecomps - {org2comp[author]}) else "self"))
    for b in (bases or []):
        if "/" in b and b.split("/")[0] in org2comp and author not in org2comp:
            deriv_links.append((mid, rel, b))
import csv
with open("deriv_links.csv", "w", newline="") as f:
    w = csv.writer(f); w.writerow(["deriv", "rel", "base"]); w.writerows(sorted(set(deriv_links)))
with open("ext_ids.csv", "w", newline="") as f:
    w = csv.writer(f); w.writerow(["id", "cls"]); w.writerows(sorted(set(ext_ids)))
print(f"{len(set(deriv_links))} derivative links, {len(set(ext_ids))} origin-classified")

# ---------- space links via Hub API ----------
model_ids = [r[0] for r in con.sql(f"SELECT id FROM '{murl}' WHERE author IN ({inlist})").fetchall()]
HDRS = {"Authorization": f"Bearer {TOKEN}"} if TOKEN else {}
space2comps, linked = {}, set()
from concurrent.futures import ThreadPoolExecutor
def fetch(model):
    out, url = [], f"https://huggingface.co/api/spaces?models={urllib.parse.quote(model, safe='')}&limit=100"
    while url:
        for attempt in range(6):
            try:
                req = urllib.request.Request(url, headers=HDRS)
                with urllib.request.urlopen(req, timeout=30) as r:
                    data = json.loads(r.read()); link = r.headers.get("Link", "")
                break
            except Exception:
                if attempt == 5: return model, []
                time.sleep(5 * (attempt + 1))
        out += [d["id"] for d in data]
        url = link.split("<")[1].split(">")[0] if 'rel="next"' in link else None
    return model, out
with ThreadPoolExecutor(10) as ex:
    for i, (model, sids) in enumerate(ex.map(fetch, model_ids)):
        c = org2comp.get(model.split("/")[0])
        for sid in sids:
            linked.add(sid)
            if c: space2comps.setdefault(sid, set()).add(c)
        if (i + 1) % 2000 == 0: print(f"space links {i+1}/{len(model_ids)}")
with open("linked_spaces.csv", "w", newline="") as f:
    w = csv.writer(f); w.writerow(["sid"]); w.writerows([[s] for s in sorted(linked)])
print(f"{len(linked)} linked spaces")

# ---------- per-week aggregation ----------
def org_rows(t, d):
    url = f"{BASE}/{t}/{d}/{t}.parquet"
    c = cols_of(url)
    author = "author" if "author" in c else "split_part(id,'/',1)"
    likes = "SUM(likes)" if "likes" in c else "NULL"
    dl = "SUM(downloads)" if ("downloads" in c and t != "spaces") else "NULL"
    dat = "SUM(downloadsAllTime)" if ("downloadsAllTime" in c and t != "spaces") else "NULL"
    return {r[0]: r[1:] for r in con.sql(f"""SELECT {author} a, COUNT(*),
        COUNT(*) FILTER (CAST(createdAt AS TIMESTAMP) >= TIMESTAMP '2025-01-01'),
        {likes}, {dl}, {dat} FROM '{url}' WHERE {author} IN ({inlist}) GROUP BY a""").fetchall()}

def row5(r):
    if r is None: return [0, 0, 0, 0, 0]
    n, n25, likes, d30, dat = r
    return [n, n25, likes if likes is not None else None,
            int(float(d30)) if d30 is not None else None,
            int(float(dat)) if dat is not None else None]

def totals(rows, payload_key, wi):
    n = l = 0
    for o in payload["orgs"]:
        v = o[payload_key][wi]
        if v and v[0] is not None: n += v[0]; l += (v[2] or 0)
    return n, l

for d in new_dates:
    print("processing", d)
    org_cache = {t: org_rows(t, d) for t in ["models", "datasets", "spaces"]}
    murl_d = f"{BASE}/models/{d}/models.parquet"
    c = cols_of(murl_d)
    likes_e = "SUM(m.likes)" if "likes" in c else "NULL"
    dl_e = "SUM(m.downloads)" if "downloads" in c else "NULL"
    dat_e = "SUM(m.downloadsAllTime)" if "downloadsAllTime" in c else "NULL"
    ext_rows = {(r[0], r[1]): r[2:] for r in con.sql(f"""SELECT split_part(m.id,'/',1) a, l.cls,
        COUNT(*), COUNT(*) FILTER (CAST(m.createdAt AS TIMESTAMP) >= TIMESTAMP '2025-01-01'),
        {likes_e}, {dl_e}, {dat_e}
        FROM '{murl_d}' m JOIN read_csv('ext_ids.csv') l ON m.id = l.id GROUP BY a, l.cls""").fetchall()}
    deriv_rows = con.sql(f"""SELECT l.base, COUNT(*), {likes_e}, {dl_e}, {dat_e}
        FROM '{murl_d}' m JOIN read_csv('deriv_links.csv') l ON m.id = l.deriv GROUP BY l.base""").fetchall()
    sp_rows = dict(con.sql(f"""SELECT m.id, m.likes FROM '{BASE}/spaces/{d}/spaces.parquet' m
        JOIN read_csv('linked_spaces.csv') l ON m.id = l.sid""").fetchall())

    # validity vs previous stored week (stale duplicate or big drop = partial)
    valid = {}
    for t, key in [("models", "m"), ("datasets", "d"), ("spaces", "s")]:
        n = sum(r[0] for r in org_cache[t].values())
        l = sum(r[2] or 0 for r in org_cache[t].values())
        pn, pl = totals(None, key, len(weeks) - 1) if False else (None, None)
        # previous totals from stored arrays
        pn = pl = 0
        for o in payload["orgs"]:
            v = o[key][-1]
            if v and v[0] is not None: pn += v[0]; pl += (v[2] or 0)
        valid[t] = not (pn and ((n, l) == (pn, pl) or n < pn * 0.90 or (pl and l < pl * 0.98)))
        if not valid[t]: print(f"  {t} snapshot looks stale/partial -> nulled")
    sp_n, sp_l = len(sp_rows), sum(v or 0 for v in sp_rows.values())
    prev_sp = [e["V"][-1] for e in payload["demos"]]
    valid_sp = valid["spaces"]

    weeks.append(d)
    for o in payload["orgs"]:
        for t, key in [("models", "m"), ("datasets", "d"), ("spaces", "s")]:
            o[key].append(row5(org_cache[t].get(o["o"])) if valid[t] else [None] * 5)
        for cls, key in [("ext", "me"), ("self", "ms")]:
            o[key].append(row5(ext_rows.get((o["o"], cls))) if valid["models"] else [None] * 5)
    for cname in companies:
        row = [0, 0, 0, 0, 0, 0] if valid["models"] else [None] * 4 + [0, 0]
        if not valid_sp: row[4] = row[5] = None
        payload["eco"][cname].append(row)
    if valid["models"]:
        for base, n, likes, d30, dat in deriv_rows:
            cname = org2comp.get(base.split("/")[0])
            if cname is None: continue
            row = payload["eco"][cname][-1]
            row[0] += n; row[1] += likes or 0
            row[2] += int(float(d30)) if d30 is not None else 0
            if dat is not None: row[3] = (row[3] or 0) + int(float(dat))
        for cname in companies:
            if payload["eco"][cname][-1][3] == 0: payload["eco"][cname][-1][3] = None
    if valid_sp:
        for sid, likes in sp_rows.items():
            for cname in space2comps.get(sid, ()):
                row = payload["eco"][cname][-1]
                row[4] += 1; row[5] += likes or 0
    known = {e["id"] for e in payload["demos"]}
    for e in payload["demos"]:
        e["V"].append(sp_rows.get(e["id"]) if valid_sp else None)
        e["cs"] = sorted(space2comps.get(e["id"], set())) or e["cs"]
    if valid_sp:
        for sid, likes in sorted(sp_rows.items(), key=lambda kv: -(kv[1] or 0))[:500]:
            if sid not in known and (likes or 0) > 100:
                payload["demos"].append({"id": sid, "cs": sorted(space2comps.get(sid, ())),
                                         "V": [None] * (len(weeks) - 1) + [likes]})

with open("data_ui.js", "w") as f:
    f.write("const UI_DATA = "); json.dump(payload, f, separators=(",", ":")); f.write(";\n")
api.upload_file(path_or_fileobj="data_ui.js", path_in_repo="data_ui.js",
                repo_id=SPACE, repo_type="space",
                commit_message=f"data: add week(s) {', '.join(new_dates)}")
print("uploaded; now at", weeks[-1])