from pathlib import Path, PurePosixPath
import argparse
import collections
import hashlib
import json
import re
import tarfile

parser = argparse.ArgumentParser(description="Safely extract an experiment archive and create audit path manifests (Python 3.11+).")
parser.add_argument("--archive", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
SOURCE = args.archive.resolve()
DEST = args.output.resolve()
if DEST.exists() and any(DEST.iterdir()):
    raise SystemExit(f"Output directory must be empty: {DEST}")
DEST.mkdir(parents=True, exist_ok=True)
start_size = SOURCE.stat().st_size

def safe_component(value):
    value = value.replace('%', '%25')
    value = re.sub(r'[<>:"\\|?*\x00-\x1f]', lambda m: f'%{ord(m.group()):02X}', value)
    while value.endswith((' ', '.')):
        value = value[:-1] + f'%{ord(value[-1]):02X}'
    if value.split('.')[0].upper() in {'CON', 'PRN', 'AUX', 'NUL', *(f'COM{i}' for i in range(1, 10)), *(f'LPT{i}' for i in range(1, 10))}:
        value = '%5F' + value
    return value

manifest = []
seen = set()
with tarfile.open(SOURCE, 'r:gz') as archive:
    for member in archive:
        original = PurePosixPath(member.name)
        if original.is_absolute() or '..' in original.parts or '\\' in member.name:
            raise RuntimeError(f'Unsafe archive path: {member.name!r}')
        relative = Path(*(safe_component(x) for x in original.parts if x not in ('', '.')))
        target = DEST / relative
        if not target.resolve().is_relative_to(DEST.resolve()):
            raise RuntimeError(f'Path escapes extraction root: {member.name!r}')
        row = {'original': member.name, 'local': relative.as_posix(), 'size': member.size, 'type': 'other'}
        if member.isdir():
            target.mkdir(parents=True, exist_ok=True)
            row['type'] = 'directory'
        elif member.isfile():
            key = relative.as_posix().casefold()
            if key in seen:
                raise RuntimeError(f'Archive filename collision: {member.name!r}')
            seen.add(key)
            target.parent.mkdir(parents=True, exist_ok=True)
            extracted = archive.extractfile(member)
            assert extracted is not None
            with extracted, target.open('wb') as handle:
                while chunk := extracted.read(1024 * 1024):
                    handle.write(chunk)
            if target.stat().st_size != member.size:
                raise RuntimeError(f'Incomplete extraction: {member.name!r}')
            row['type'] = 'file'
        else:
            row['type'] = 'skipped_link_or_special'
        manifest.append(row)

if SOURCE.stat().st_size != start_size:
    raise RuntimeError('Archive changed during extraction')
with SOURCE.open('rb') as handle:
    digest = hashlib.file_digest(handle, 'sha256').hexdigest()
(DEST / 'archive_manifest.json').write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding='utf-8')
files = [x for x in manifest if x['type'] == 'file']
summary = {'source': str(SOURCE), 'compressed_bytes': start_size, 'sha256': digest,
           'files': len(files), 'uncompressed_bytes': sum(x['size'] for x in files),
           'skipped_special': [x for x in manifest if x['type'] == 'skipped_link_or_special'],
           'top_level_files': dict(collections.Counter(PurePosixPath(x['original']).parts[0] for x in files)),
           'largest_files': sorted(files, key=lambda x: x['size'], reverse=True)[:12]}
(DEST / 'archive_summary.json').write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding='utf-8')
print(json.dumps(summary, ensure_ascii=False, indent=2))