"""Obtem modelos Google e runtime npm oficiais, mantendo tudo local e versionado."""
import base64
import hashlib
import io
import json
from pathlib import Path
import tarfile
import urllib.request

ROOT = Path(__file__).resolve().parent
VERSION = '0.10.32'

def download(url):
    with urllib.request.urlopen(url, timeout=60) as response:
        return response.read()

def main():
    dest = ROOT / 'vendor/mediapipe'
    dest.mkdir(parents=True, exist_ok=True)
    metadata = json.loads(download('https://registry.npmjs.org/@mediapipe/tasks-vision/' + VERSION))
    archive = download(metadata['dist']['tarball'])
    integrity = metadata['dist']['integrity']
    assert integrity == 'sha512-' + base64.b64encode(hashlib.sha512(archive).digest()).decode()
    files = {}
    with tarfile.open(fileobj=io.BytesIO(archive), mode='r:gz') as tar:
        for member in tar.getmembers():
            name = member.name.removeprefix('package/')
            if name in ('vision_bundle.mjs', 'vision_bundle.js', 'LICENSE') or (
                name.startswith('wasm/') and '/' not in name[5:] and name.endswith(('.js', '.wasm'))):
                data = tar.extractfile(member).read()
                path = dest / name
                path.parent.mkdir(parents=True, exist_ok=True)
                path.write_bytes(data)
                files[name] = hashlib.sha256(data).hexdigest()
    for name, url in {
        'LICENSE':'https://raw.githubusercontent.com/google-ai-edge/mediapipe/master/LICENSE',
        'face.tflite':'https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/1/blaze_face_short_range.tflite',
        'hands.task':'https://storage.googleapis.com/mediapipe-models/hand_landmarker/hand_landmarker/float16/1/hand_landmarker.task',
    }.items():
        data = download(url); (dest/name).write_bytes(data)
        files[name] = {'url':url, 'sha256':hashlib.sha256(data).hexdigest()}
    (dest/'PROVENIENCIA.json').write_text(json.dumps({'pacote':'@mediapipe/tasks-vision','versao':VERSION,
        'licenca':metadata.get('license'), 'tarball':metadata['dist']['tarball'],'integrity':integrity,
        'arquivos':files},indent=2),encoding='utf-8')
    print('MediaPipe local: runtime e modelos baixados e registrados.')

if __name__ == '__main__':
    main()