#!/usr/bin/env python3
"""Restore the folder structure of this CKAN dataset and verify every file.

Usage:
    python3 restore.py [TARGET_DIR] [--only TEXT] [--manifest FILE] [--dataset ID] [--ckan URL]

How the dataset is organised:
  * Every data resource is named after its original path, with "__" in place of "/".
    Example: B_dot__180mm_antenna__1_5-Pa__IQ_Demo.csv -> B_dot/180mm_antenna/1_5-Pa/IQ_Demo.csv
  * A resource whose name ends in ".tar" is an archive. It holds paths relative to the
    dataset root and is unpacked into TARGET_DIR.
  * MANIFEST.sha256 lists the SHA-256 checksum and relative path of every original file.
    The script verifies the restored tree against it.

Only the Python standard library is needed (Python >= 3.8). Downloads are anonymous.
A file that already exists with the expected checksum is not downloaded again.
"""
import argparse
import hashlib
import json
import os
import sys
import tarfile
import urllib.request

CKAN = 'https://ckan.fdm.uni-greifswald.de'
DATASET = '46ef6872-09be-4ac4-87ae-bc2ca5c54f79'
CHUNK = 8 * 1024 * 1024
META = {'MANIFEST.sha256', 'README.md', 'restore.py'}


def api(ckan, action, **params):
    query = '&'.join(f'{k}={v}' for k, v in params.items())
    with urllib.request.urlopen(f'{ckan}/api/3/action/{action}?{query}', timeout=60) as r:
        return json.load(r)['result']


def sha256_of(path):
    h = hashlib.sha256()
    with open(path, 'rb') as f:
        for block in iter(lambda: f.read(CHUNK), b''):
            h.update(block)
    return h.hexdigest()


def download(url, dest):
    """Stream url to dest (via a .part file); return the SHA-256 of the data."""
    os.makedirs(os.path.dirname(dest) or '.', exist_ok=True)
    h = hashlib.sha256()
    tmp = dest + '.part'
    with urllib.request.urlopen(url, timeout=300) as r, open(tmp, 'wb') as f:
        for block in iter(lambda: r.read(CHUNK), b''):
            f.write(block)
            h.update(block)
    os.replace(tmp, dest)
    return h.hexdigest()


def read_manifest(path):
    entries = {}
    with open(path, encoding='utf-8') as f:
        for line in f:
            line = line.rstrip('\n')
            if line:
                digest, rel = line.split(None, 1)
                entries[rel.lstrip('*')] = digest
    return entries


def safe_extract(tar_path, target):
    root = os.path.realpath(target)
    with tarfile.open(tar_path) as tar:
        for m in tar.getmembers():
            dest = os.path.realpath(os.path.join(root, m.name))
            if not (dest == root or dest.startswith(root + os.sep)) or m.issym() or m.islnk():
                raise SystemExit(f'unsafe member in {tar_path}: {m.name}')
        if sys.version_info >= (3, 12):
            tar.extractall(root, filter='data')
        else:
            tar.extractall(root)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('target', nargs='?', default='for_publication')
    ap.add_argument('--only', help='only resources whose name contains this text')
    ap.add_argument('--manifest', help='local MANIFEST.sha256 (default: the one in the dataset)')
    ap.add_argument('--dataset', default=DATASET)
    ap.add_argument('--ckan', default=CKAN)
    a = ap.parse_args()

    resources = api(a.ckan, 'package_show', id=a.dataset)['resources']
    by_name = {r['name']: r for r in resources}
    manifest_path = a.manifest
    if not manifest_path and 'MANIFEST.sha256' in by_name:
        manifest_path = os.path.join(a.target, 'MANIFEST.sha256')
        download(by_name['MANIFEST.sha256']['url'], manifest_path)
    manifest = read_manifest(manifest_path) if manifest_path else {}

    for r in resources:
        name = r['name']
        if a.only and a.only not in name:
            continue
        if name in META:
            if name != 'MANIFEST.sha256':
                download(r['url'], os.path.join(a.target, name))
            continue
        if name.endswith('.tar'):
            tar_path = os.path.join(a.target, '.archives', name)
            print(f'archive {name}')
            download(r['url'], tar_path)
            safe_extract(tar_path, a.target)
            os.remove(tar_path)
            continue
        rel = name.replace('__', '/')
        dest = os.path.join(a.target, rel)
        if os.path.exists(dest) and manifest.get(rel) == sha256_of(dest):
            print(f'ok (exists) {rel}')
            continue
        digest = download(r['url'], dest)
        state = 'ok' if manifest.get(rel) == digest else ('NOT IN MANIFEST' if rel not in manifest else 'CHECKSUM MISMATCH')
        print(f'{state:17} {rel}')

    if manifest:
        wanted = {p: d for p, d in manifest.items() if not a.only or a.only in p.replace('/', '__')}
        bad = [p for p, d in sorted(wanted.items())
               if not os.path.exists(os.path.join(a.target, p)) or sha256_of(os.path.join(a.target, p)) != d]
        print(f'\nverified {len(wanted) - len(bad)} of {len(wanted)} files against the manifest')
        for p in bad:
            print(f'  MISSING OR WRONG: {p}')
        sys.exit(1 if bad else 0)


if __name__ == '__main__':
    main()