"""Fetch reproducible Poly Haven glTF sources and PBR ground materials (CC0)."""
import concurrent.futures, hashlib, json, pathlib, urllib.request, time

ROOT = pathlib.Path(__file__).resolve().parents[2]
DEST = ROOT / 'Content' / 'Vegitations'
ASSETS = ['pine_tree_01', 'fir_tree_01', 'tree_small_02', 'fir_sapling_medium',
          'shrub_01', 'shrub_04', 'fern_02', 'grass_medium_01', 'grass_medium_02', 'weed_plant_02']
HEADERS = {'User-Agent': 'UniverseVegetationImporter/1.0 (Poly Haven CC0 asset preparation)'}

def request(url):
    return urllib.request.urlopen(urllib.request.Request(url, headers=HEADERS), timeout=120)

def md5(path):
    h = hashlib.md5()
    with path.open('rb') as f:
        for block in iter(lambda: f.read(4*1024*1024), b''): h.update(block)
    return h.hexdigest()

def fetch(item):
    path = DEST / item['path']
    if not path.resolve().is_relative_to(DEST.resolve()): raise RuntimeError('Asset path escapes content folder')
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.exists() and path.stat().st_size == item['size'] and md5(path) == item['md5']:
        return
    partial = path.with_suffix(path.suffix + '.partial')
    for attempt in range(4):
        try:
            with request(item['url']) as response, partial.open('wb') as output:
                received = 0
                while True:
                    block = response.read(1024*1024)
                    if not block: break
                    received += len(block)
                    if received > item['size']: raise RuntimeError('Download exceeds declared size: '+str(path))
                    output.write(block)
            if partial.stat().st_size != item['size'] or md5(partial) != item['md5']:
                raise RuntimeError('Checksum/length mismatch: ' + str(path))
            partial.replace(path)
            print('Downloaded', item['path'], round(item['size']/1048576, 1), 'MiB', flush=True)
            return
        except Exception:
            if attempt == 3: raise
            time.sleep(attempt + 1)

def main():
    DEST.mkdir(parents=True, exist_ok=True)
    manifest_path = DEST / 'sources.json'
    if manifest_path.exists():
        manifest = json.loads(manifest_path.read_text())
    else:
        files = []
        for asset in ASSETS:
            with request('https://api.polyhaven.com/files/' + asset) as f: data = json.load(f)
            model = data['gltf']['2k']['gltf']
            entries = {'model.gltf': model, **model.get('include', {})}
            for relative, entry in entries.items():
                files.append({**{k: entry[k] for k in ('url','size','md5')}, 'path': asset + '/source/' + relative})
        # Ground materials share the same biome field as the plants.
        for asset in ['forest_ground_04', 'withered_grass', 'snow_02']:
            with request('https://api.polyhaven.com/files/' + asset) as f: data = json.load(f)
            for kind in ['diff', 'nor_gl', 'rough']:
                key = next(k for k in data if k.lower() in ({'diff': ['diff','diffuse'], 'rough':['rough','roughness'], 'nor_gl':['nor_gl']}[kind]))
                entry = data[key]['2k']['jpg']
                files.append({**{k:entry[k] for k in ('url','size','md5')}, 'path':'ground/'+asset+'/'+kind+'.jpg'})
        manifest = {'provider':'Poly Haven', 'license':'CC0', 'source':'https://polyhaven.com', 'files':files}
        size = sum(f['size'] for f in files)
        if size > 2_000_000_000: raise RuntimeError('Original asset budget exceeded: '+str(size))
        manifest_path.write_text(json.dumps(manifest, indent=2)+'\n')
    if not manifest.get('foliageMasksIncluded'):
        for asset in ASSETS:
            with request('https://api.polyhaven.com/files/' + asset) as f: data = json.load(f)
            for kind, variants in data.items():
                if 'alpha' not in kind.lower() and 'dry_diff' not in kind.lower(): continue
                formats=variants.get('2k',{})
                extension='png' if 'png' in formats else 'jpg'
                if extension not in formats: continue
                entry=formats[extension]
                manifest['files'].append({**{k:entry[k] for k in ('url','size','md5')},
                    'path':asset+'/source/masks/'+kind+'.'+extension})
        manifest['foliageMasksIncluded']=True
        if sum(f['size'] for f in manifest['files'])>2_000_000_000: raise RuntimeError('Asset budget exceeded')
        manifest_path.write_text(json.dumps(manifest,indent=2)+'\n')
    print('Poly Haven CC0 assets:', len(manifest['files']), 'files,', round(sum(f['size'] for f in manifest['files'])/1048576), 'MiB', flush=True)
    if any(f['size']<=0 for f in manifest['files']) or sum(f['size'] for f in manifest['files'])>2_000_000_000:
        raise RuntimeError('Source download budget exceeded')
    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
        list(executor.map(fetch, manifest['files']))
    print('All source checksums verified.', flush=True)

if __name__ == '__main__': main()