"""Fetch and hash-check the documented Menagerie G1 model, including its licence."""
import argparse
import hashlib
import json
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--output', type=Path, default=Path('menagerie/unitree_g1'))
args = parser.parse_args()
manifest = json.loads(Path(__file__).with_name('g1_model_files.json').read_text(encoding='utf-8'))
base = 'https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/'

def fetch(record):
    name = record['path']
    relative = Path(name)
    if relative.is_absolute() or '..' in relative.parts:
        raise ValueError('Invalid model path')
    target = args.output / relative
    if target.is_file() and hashlib.sha256(target.read_bytes()).hexdigest() == record['sha256']:
        return
    request = urllib.request.Request(base + manifest['commit'] + '/unitree_g1/' + name,
                                     headers={'User-Agent': 'TechniaHQRobot-model-example/1.0'})
    with urllib.request.urlopen(request, timeout=60) as response:
        content = response.read()
    if hashlib.sha256(content).hexdigest() != record['sha256']:
        raise RuntimeError('Model file hash mismatch for ' + name)
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_bytes(content)

with ThreadPoolExecutor(max_workers=4) as pool:
    list(pool.map(fetch, manifest['files']))
print('Verified', len(manifest['files']), 'model files at commit', manifest['commit'])