#!/usr/bin/env python3
"""Match local category/scene IDs to Hub HTMLs and capture their initial cameras.

Optional build tooling: npm install --no-save playwright && npx playwright install chromium.
The running demo needs neither Playwright nor a browser installation.
"""
import argparse
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
from pathlib import Path
import shutil
import subprocess
from urllib.parse import quote
from urllib.request import urlopen

ROOT = Path(__file__).resolve().parents[1]


def read_json(url):
    with urlopen(url, timeout=60) as response:
        return json.load(response)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--dataset', default='ritianyu/SplatAnything_Demo')
    parser.add_argument('--examples_dir', type=Path, default=ROOT / 'examples')
    parser.add_argument('--work_dir', type=Path, default=ROOT / '.cache/preview-build')
    args = parser.parse_args()
    api = f'https://huggingface.co/api/datasets/{args.dataset}'
    revision = read_json(api)['sha']
    remote = read_json(f'{api}/tree/{revision}?recursive=true&limit=1000')
    scenes = {}
    for folder in ('object', 'indoor', 'street', 'aerial', 'rig', 'pano'):
        for scene in sorted((args.examples_dir / folder).iterdir()):
            if not scene.is_dir() or not (scene / 'images').is_dir():
                continue
            prefix = f'{folder}/{scene.name}/'
            candidates = [x for x in remote if x['type'] == 'file' and x['path'].startswith(prefix)
                          and x['path'].endswith('.html') and x['size'] > 0]
            if len(candidates) != 1:
                raise ValueError(f'Expected exactly one nonempty HTML for {prefix}: found {len(candidates)}')
            item = candidates[0]
            key = f'{folder}-{scene.name}'
            scenes[key] = dict(path=item['path'], bytes=item['size'], sha256=item['lfs']['oid'], poster=f'{key}.webp')
    catalog = dict(dataset=args.dataset, revision=revision, capture=dict(time=0, width=1440, height=900), scenes=scenes)
    args.work_dir.mkdir(parents=True, exist_ok=True)
    source = args.work_dir / 'source'
    source.mkdir(exist_ok=True)

    def fetch(item):
        key, record = item
        target = source / f'{key}.html'
        if target.is_file() and target.stat().st_size == record['bytes']:
            if hashlib.sha256(target.read_bytes()).hexdigest() == record['sha256']:
                return
        temporary = target.with_suffix('.download')
        try:
            url = f'https://huggingface.co/datasets/{args.dataset}/resolve/{revision}/' + quote(record['path'], safe='/')
            digest = hashlib.sha256()
            with urlopen(url, timeout=90) as response, temporary.open('wb') as output:
                while chunk := response.read(1024 * 1024):
                    output.write(chunk)
                    digest.update(chunk)
            if temporary.stat().st_size != record['bytes'] or digest.hexdigest() != record['sha256']:
                raise ValueError(f'Download verification failed: {key}')
            temporary.replace(target)
        finally:
            temporary.unlink(missing_ok=True)
        print(f'Downloaded {key}', flush=True)

    with ThreadPoolExecutor(max_workers=4) as pool:
        list(pool.map(fetch, scenes.items()))
    build_catalog = args.work_dir / 'catalog.json'
    build_catalog.write_text(json.dumps(catalog, indent=2) + '\n')
    posters = args.work_dir / 'posters' / revision
    subprocess.run(['node', str(ROOT / 'scripts/capture_scene_previews.cjs'),
                    str(build_catalog.resolve()), str(source.resolve()), str(posters.resolve())], check=True)
    assets = ROOT / 'assets/scene-previews'
    assets.mkdir(parents=True, exist_ok=True)
    for record in scenes.values():
        shutil.copyfile(posters / record['poster'], assets / record['poster'])
    staged = assets / 'catalog.json.tmp'
    staged.write_text(json.dumps(catalog, indent=2) + '\n')
    staged.replace(assets / 'catalog.json')
    print(f'Prepared {len(scenes)} initial-frame previews at {revision}. Restart the demo to use the new catalog.')


if __name__ == '__main__':
    main()