"""Check a downloaded release without training or evaluating a model."""
from pathlib import Path
import argparse
import ast
import hashlib
import importlib
import json
import sys
import numpy as np


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


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--package', type=Path, default=Path(__file__).resolve().parent)
    p.add_argument('--models', action='store_true', help='Also load every checkpoint on CPU and check its interface and tensors')
    p.add_argument('--output', type=Path, help='Optional external JSON receipt')
    a = p.parse_args()
    root = a.package.resolve()
    manifest = json.loads((root / 'MANIFEST.json').read_text())
    errors = []
    for row in manifest['files']:
        path = root / row['path']
        if root not in path.resolve().parents:
            errors.append('Path escapes release: ' + row['path']); continue
        if not path.is_file():
            errors.append('Missing: ' + row['path']); continue
        if path.stat().st_size != row['bytes'] or sha256(path) != row['sha256']:
            errors.append('Release content differs: ' + row['path'])
        if path.suffix == '.py':
            try:
                ast.parse(path.read_text())
            except SyntaxError as error:
                errors.append(f'{row["path"]}: {error}')
    data = json.loads((root / 'data/manifest.json').read_text())
    statuses = {}
    for row in data['datasets']:
        statuses[row['status']] = statuses.get(row['status'], 0) + 1
        with np.load(root / row['path'], allow_pickle=False) as arrays:
            schema = {key: {'shape': list(arrays[key].shape), 'dtype': str(arrays[key].dtype)} for key in arrays.files if key != 'metadata'}
            if schema != row['arrays']:
                errors.append('Dataset schema differs: ' + row['path'])
            for key in arrays.files:
                value = arrays[key]
                if np.issubdtype(value.dtype, np.number) and not np.isfinite(value).all():
                    errors.append('Non-finite dataset array: ' + row['path'] + ':' + key)
    model_count = 0
    if a.models:
        import torch
        torch.set_num_threads(1)
        sys.path.insert(0, str(root))
        for row in manifest['files']:
            path = root / row['path']
            if path.suffix != '.pt':
                continue
            method = Path(row['path']).parts[1]
            config = importlib.import_module(f'experiments.{method}.src.config').default_config()
            model = importlib.import_module(f'experiments.{method}.src.model').build_model(config)
            state = torch.load(path, map_location='cpu', weights_only=True)
            try:
                model.load_state_dict(state, strict=True)
            except RuntimeError as error:
                errors.append(f'Checkpoint interface: {row["path"]}: {error}')
            if not all(torch.isfinite(t).all() for t in state.values()):
                errors.append('Non-finite checkpoint: ' + row['path'])
            model_count += 1
    result = {'passed': not errors, 'files_checked': len(manifest['files']),
              'datasets_checked': len(data['datasets']), 'dataset_provenance_counts': statuses,
              'checkpoints_loaded_on_cpu': model_count, 'models_executed': False,
              'release_integrity_only': True, 'errors': errors}
    if a.output:
        output = a.output.resolve()
        supplied = {str((root / r['path']).resolve()) for r in manifest['files']}
        if str(output) in supplied or output == root / 'MANIFEST.json':
            raise ValueError('Receipt must not overwrite a supplied file')
        output.parent.mkdir(parents=True, exist_ok=True)
        output.write_text(json.dumps(result, indent=2) + '\n')
    print(json.dumps(result, indent=2))
    if errors:
        raise SystemExit(1)


if __name__ == '__main__':
    main()