#!/usr/bin/env python3
"""Restore the published MP source from adjacent verified source-offer files."""
import argparse
import hashlib
import json
from pathlib import Path
import shutil
import subprocess
import tarfile
import tempfile


def digest(path):
    h = hashlib.sha256()
    with path.open('rb') as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b''):
            h.update(chunk)
    return h.hexdigest()


def restore(offer, target):
    manifest = json.loads((offer / 'manifest.json').read_text())
    if target.exists():
        raise ValueError('Destination must not exist; no existing checkout is modified.')
    target.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix='.bo1-source-check-', dir=target.parent) as temporary:
        probe = Path(temporary)
        (probe / 'CaseProbe').write_text('case')
        if (probe / 'caseprobe').exists():
            raise ValueError('Destination filesystem must be case-sensitive. Use Linux or a case-sensitive APFS volume.')
    def safe_relative(name):
        path = Path(name)
        return not path.is_absolute() and '..' not in path.parts and '.git' not in path.parts
    artifacts = {entry['name']: entry for entry in manifest['artifacts']}
    if not all(name in artifacts and Path(name).name == name for name in manifest['base']['parts']):
        raise ValueError('Unsafe Git bundle part name')
    if not all(safe_relative(name) for name in [*manifest['overlayFiles'], *manifest['frozenNativeFiles']]):
        raise ValueError('Unsafe source path')
    for name, entry in artifacts.items():
        if Path(name).name != name:
            raise ValueError('Unsafe source artifact name')
        path = offer / name
        if path.stat().st_size != entry['size'] or digest(path) != entry['sha256']:
            raise ValueError('Source offer is incomplete or corrupt: ' + name)
    with tempfile.TemporaryDirectory(prefix='bo1-source-restore-') as temporary:
        temporary = Path(temporary)
        bundle = temporary / 'base.bundle'
        with bundle.open('wb') as destination:
            for name in manifest['base']['parts']:
                with (offer / name).open('rb') as source:
                    shutil.copyfileobj(source, destination, 1024 * 1024)
        if digest(bundle) != manifest['base']['bundleSha256']:
            raise ValueError('Joined Git bundle checksum mismatch')
        subprocess.run(['git', '-c', 'pack.threads=2', '-c', 'index.threads=2', 'clone', str(bundle), str(target)], check=True)
        subprocess.run(['git', '-C', str(target), 'checkout', '--detach', manifest['base']['revision']], check=True)
        subprocess.run(['git', '-C', str(target), 'remote', 'set-url', 'origin', manifest['base']['repository']], check=True)
        with tarfile.open(offer / 'mp-source-overlay.tar.gz', 'r:gz') as archive:
            for member in archive.getmembers():
                path = Path(member.name)
                if member.isdir():
                    continue
                if not member.isfile() or path.is_absolute() or '..' in path.parts or '.git' in path.parts:
                    raise ValueError('Unsafe overlay member')
                destination = temporary / path
                destination.parent.mkdir(parents=True, exist_ok=True)
                with archive.extractfile(member) as source, destination.open('wb') as output:
                    shutil.copyfileobj(source, output, 1024 * 1024)
                destination.chmod(member.mode & 0o777)
        patch = temporary / 'mp.patch'
        if patch.stat().st_size:
            subprocess.run(['git', '-C', str(target), 'apply', '--binary', str(patch)], check=True)
        for name, expected in manifest['overlayFiles'].items():
            source = temporary / 'files' / name
            if digest(source) != expected:
                raise ValueError('Overlay source checksum mismatch: ' + name)
            destination = target / name
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(source, destination)
        for name, expected in manifest['frozenNativeFiles'].items():
            if digest(target / name) != expected:
                raise ValueError('Restored native source differs from frozen MP input: ' + name)
    print('Restored exact frozen MP source and preserved upstream Git history:', target)
    print('Build: ./scripts/build-engine.sh release')


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