#!/usr/bin/env python3
"""Minimal authenticated GCS access for the Waymo bucket using the existing gcloud
legacy ADC refresh token (no gcloud/gsutil/google-cloud libs needed; requests only)."""
import json, os, urllib.parse, glob, requests

BUCKET = "waymo_open_dataset_v_1_4_3"
PREFIX = "individual_files/training/"


def _adc_path():
    cands = glob.glob(os.path.expanduser("~/.config/gcloud/legacy_credentials/*/adc.json"))
    cands += [os.path.expanduser("~/.config/gcloud/application_default_credentials.json")]
    for c in cands:
        if os.path.exists(c):
            return c
    raise FileNotFoundError("no ADC json found")


def access_token():
    d = json.load(open(_adc_path()))
    r = requests.post("https://oauth2.googleapis.com/token", data={
        "client_id": d["client_id"], "client_secret": d["client_secret"],
        "refresh_token": d["refresh_token"], "grant_type": "refresh_token"}, timeout=30)
    r.raise_for_status()
    return r.json()["access_token"]


def list_training_tfrecords(tok):
    names, page = [], None
    while True:
        params = {"prefix": PREFIX, "maxResults": 1000}
        if page:
            params["pageToken"] = page
        r = requests.get(f"https://storage.googleapis.com/storage/v1/b/{BUCKET}/o",
                         headers={"Authorization": f"Bearer {tok}"}, params=params, timeout=60)
        r.raise_for_status()
        j = r.json()
        names += [it["name"] for it in j.get("items", []) if it["name"].endswith(".tfrecord")]
        page = j.get("nextPageToken")
        if not page:
            break
    return sorted(names)


def download(tok, name, dst):
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    url = (f"https://storage.googleapis.com/storage/v1/b/{BUCKET}/o/"
           f"{urllib.parse.quote(name, safe='')}?alt=media")
    with requests.get(url, headers={"Authorization": f"Bearer {tok}"}, stream=True, timeout=600) as r:
        r.raise_for_status()
        with open(dst, "wb") as f:
            for chunk in r.iter_content(chunk_size=1 << 22):
                f.write(chunk)
    return dst


if __name__ == "__main__":
    tok = access_token()
    print("got access token:", tok[:12], "...")
    segs = list_training_tfrecords(tok)
    print("training tfrecords in bucket:", len(segs))
    for s in segs[:3]:
        print("  ", s)