#!/usr/bin/env python3
"""Verify and reassemble this source collection; optionally extract complete trees."""
import argparse
import hashlib
import json
import pathlib
import shutil
import sys
import tarfile
import zipfile

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('destination', type=pathlib.Path, help='New output directory outside the downloaded collection')
parser.add_argument('--extract', action='store_true', help='Also extract source trees and populate pinned submodules (Python 3.12+)')
args = parser.parse_args()
bundle = pathlib.Path(__file__).resolve().parent
destination = args.destination.resolve()
if destination == bundle or bundle in destination.parents:
    parser.error('Choose a destination outside this downloaded collection.')
if destination.exists() and any(destination.iterdir()):
    parser.error('Destination already contains files. Choose a new/empty directory.')
if args.extract and sys.version_info < (3, 12):
    parser.error('--extract requires Python 3.12+ for safe tar extraction filtering.')
destination.mkdir(parents=True, exist_ok=True)
archives = destination / 'archives'
archives.mkdir(exist_ok=True)
manifest = json.loads((bundle / 'manifest.json').read_text())
components = manifest['components']
archive_paths = {}
for item in components:
    if item.get('error'):
        raise RuntimeError('Source collection contains a failed component: ' + item['component'])
    target = archives / item['archive_file']
    digest = hashlib.sha256()
    size = 0
    with target.open('wb') as stream:
        for part in item['files']:
            path = bundle / part['file']
            data = path.read_bytes()
            if len(data) != part['bytes'] or hashlib.sha256(data).hexdigest() != part['sha256']:
                raise RuntimeError('Source part checksum mismatch: ' + part['file'])
            stream.write(data)
            digest.update(data)
            size += len(data)
    if size != item['archive_bytes'] or digest.hexdigest() != item['archive_sha256']:
        raise RuntimeError('Reassembled archive checksum mismatch: ' + item['archive_file'])
    archive_paths[item['component']] = target
    print('Verified archive:', item['component'], size, 'bytes')

if args.extract:
    trees = destination / 'trees'
    trees.mkdir(exist_ok=True)
    staged = destination / '_extract'
    staged.mkdir(exist_ok=True)
    restored = {}
    for item in sorted(components, key=lambda x: bool(x.get('parent_component'))):
        component = item['component']
        archive = archive_paths[component]
        scratch = staged / component
        scratch.mkdir(exist_ok=True)
        if archive.suffix == '.zip':
            with zipfile.ZipFile(archive) as z:
                for info in z.infolist():
                    path = pathlib.PurePosixPath(info.filename)
                    if path.is_absolute() or '..' in path.parts:
                        raise RuntimeError('Unsafe ZIP member: ' + info.filename)
                z.extractall(scratch)
        else:
            with tarfile.open(archive, 'r:gz') as t:
                t.extractall(scratch, filter='data')
        roots = item['archive_roots']
        if len(roots) != 1:
            raise RuntimeError('Unexpected archive root count: ' + component)
        extracted = scratch / roots[0]
        if item.get('parent_component'):
            target = restored[item['parent_component']] / item['submodule_path']
        else:
            target = trees / component
        if target.exists():
            if any(target.iterdir()):
                raise RuntimeError('Submodule destination already contains files: ' + str(target))
            target.rmdir()
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.move(str(extracted), str(target))
        restored[component] = target
        print('Restored tree:', component, '->', target)
    shutil.rmtree(staged)
print('Complete:', destination)