"""Validate archived analysis before allowing deletion of downloaded weights."""
import argparse
import csv
import hashlib
import json
import pathlib
import shutil
import time


def validate(root):
    analysis = root / 'analysis'
    summary = json.loads((analysis / 'summary.json').read_text())
    assert summary['completed']
    manifest = json.loads((root / 'manifest.json').read_text())
    assert summary['revision'] == manifest['revision']
    for mode in ('tensor', 'block'):
        objects = json.loads((analysis / f'{mode}_objects.json').read_text())
        with (analysis / f'{mode}_pairs.csv').open() as f:
            pairs = list(csv.DictReader(f))
        coverage = json.loads((analysis / f'{mode}_coverage.json').read_text())
        assert len(pairs) == coverage['pairs_verified']
        edges = {}
        for p in pairs:
            a, b = int(p['a']), int(p['b'])
            assert objects[a]['shape'] == objects[b]['shape']
            assert objects[a]['role'] == objects[b]['role']
            edges[(b, a)] = float(p['error_member_to_rep'])
            edges[(a, b)] = float(p['error_rep_to_member'])
        for file in analysis.glob(f'groups_{mode}_*.json'):
            threshold = float(file.stem.rsplit('_', 1)[1])
            groups = json.loads(file.read_text())
            members = [i for g in groups for i in g['members']]
            assert len(set(members)) == len(members), file
            for g in groups:
                assert g['representative'] in g['members']
                for member in g['members']:
                    if member != g['representative']:
                        assert edges[(g['representative'], member)] <= threshold + 1e-12
                assert g['max_error'] <= threshold + 1e-12
    assert (root / 'REPORT.md').is_file()
    images = list((root / 'figures').glob('*.png'))
    assert len(images) >= 5
    for file in images:
        assert file.read_bytes()[:8] == b'\x89PNG\r\n\x1a\n'
    checksums = {}
    for file in sorted(root.rglob('*')):
        if (file.is_file() and 'weights' not in file.relative_to(root).parts
                and file.name not in ('checksums.json', 'validation.json', 'cleanup.json', 'status.json')
                and not file.name.endswith('.log')):
            checksums[str(file.relative_to(root))] = hashlib.sha256(file.read_bytes()).hexdigest()
    (root / 'checksums.json').write_text(json.dumps(checksums, indent=2))
    result = {'validated_unix': time.time(), 'files_checksummed': len(checksums),
              'group_edges_checked': True, 'passed': True}
    (root / 'validation.json').write_text(json.dumps(result, indent=2))
    return result


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--root', required=True, type=pathlib.Path)
    p.add_argument('--cleanup-weights', action='store_true')
    args = p.parse_args()
    root = args.root.resolve()
    result = validate(root)
    if args.cleanup_weights:
        weights = root / 'weights'
        manifest = json.loads((root / 'manifest.json').read_text())
        assert not weights.is_symlink()
        assert (weights / '.pilot-owned').read_text() == manifest['repo']
        expected_names = {entry['name'] for entry in manifest['files']} | {'.pilot-owned'}
        assert {file.name for file in weights.iterdir()} == expected_names
        assert all(file.is_file() and not file.is_symlink() for file in weights.iterdir())
        deleted_bytes = sum(file.stat().st_size for file in weights.iterdir())
        shutil.rmtree(weights)
        (root / 'cleanup.json').write_text(json.dumps({'deleted_unix': time.time(),
                                                      'deleted_bytes': deleted_bytes,
                                                      'deleted_directory': 'weights',
                                                      'reason': 'User-authorized delete-after-analysis; archived results validated'}, indent=2))
        result['weights_removed'] = not weights.exists()
    print(json.dumps(result))


if __name__ == '__main__':
    main()