"""Restore the pinned source archives locally without executing upstream code."""
import argparse
import hashlib
import io
import json
from pathlib import Path, PurePosixPath
import sys
import tarfile


def digest(data):
    return hashlib.sha256(data).hexdigest()


def archive_bytes(directory, entry, downloads):
    filename = entry.get('file', f"{entry['id']}-{entry.get('commit', '')}.tar.gz")
    source = directory / filename
    if source.is_file():
        data = source.read_bytes()
    else:
        record = next((item for item in downloads.get('archives', []) if item['file'] == filename), None)
        if record is None:
            raise ValueError(f'Missing source archive: {filename}')
        chunks = []
        for part in record['parts']:
            name = PurePosixPath(part['file'])
            if len(name.parts) != 1 or name.is_absolute():
                raise ValueError('Unsafe part name')
            chunk = (directory / name.name).read_bytes()
            if len(chunk) != part['bytes'] or digest(chunk) != part['sha256']:
                raise ValueError(f'Invalid source part: {name}')
            chunks.append(chunk)
        data = b''.join(chunks)
    if digest(data) != entry['sha256']:
        raise ValueError(f'Invalid source archive: {filename}')
    return data


def restore(directory, output):
    if sys.version_info < (3, 12):
        raise RuntimeError('Python 3.12+ is required for safe archive extraction')
    if output.exists():
        raise ValueError('Output already exists; choose a new directory')
    lock = json.loads((directory / 'source-lock.json').read_text())
    downloads_path = directory / 'downloads.json'
    downloads = json.loads(downloads_path.read_text()) if downloads_path.exists() else {}
    # Verify all archives before creating the output, including a corrupt late part.
    for entry in lock['components']:
        archive_bytes(directory, entry, downloads)
    output.mkdir(parents=True)
    for entry in lock['components']:
        destination = PurePosixPath(entry['restorePath'])
        if destination.is_absolute() or '..' in destination.parts:
            raise ValueError('Unsafe restore path')
        with tarfile.open(fileobj=io.BytesIO(archive_bytes(directory, entry, downloads))) as archive:
            members = archive.getmembers()
            roots = {PurePosixPath(member.name).parts[0] for member in members}
            if len(roots) != 1:
                raise ValueError('Expected one archive root')
            for member in members:
                path = PurePosixPath(member.name)
                if path.is_absolute() or '..' in path.parts:
                    raise ValueError(f'Unsafe archive member: {member.name}')
                relative = PurePosixPath(*path.parts[1:])
                # Preserve the tree within a fixed target. data_filter validates
                # links and refuses devices or paths escaping that target.
                member.name = str(destination / relative)
            archive.extractall(output, members=members, filter='data')
    print(json.dumps({'archives': len(lock['components']), 'output': str(output), 'upstreamCodeExecuted': False}))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--materials', type=Path, default=Path(__file__).resolve().parent)
    args = parser.parse_args()
    restore(args.materials.resolve(), args.output.resolve())