#!/usr/bin/env python3
"""Freeze and restore explicit project build inputs; never copy user credentials."""
import argparse
import fcntl
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import stat
import subprocess
import sys
import tomllib

project = Path(__file__).resolve().parents[1]
OWNED = ('.cargo', '.gitignore', '.gitmodules', 'AGENTS.md', 'Cargo.lock',
         'Cargo.toml', 'LICENSE', 'Makefile', 'README.md', 'VOID_PACKAGES_COMMIT',
         'config', 'docs', 'image', 'packages', 'profiles', 'rust-toolchain.toml',
         'rust', 's6', 'tests', 'tools')


def digest(path):
    with path.open('rb') as stream:
        return hashlib.file_digest(stream, 'sha256').hexdigest()


def copy(source, destination):
    destination.parent.mkdir(parents=True, exist_ok=True)
    if source.is_dir() and not source.is_symlink():
        shutil.copytree(source, destination, symlinks=True,
                        ignore=shutil.ignore_patterns('__pycache__', '*.pyc'))
    else:
        shutil.copy2(source, destination, follow_symlinks=False)


def inventory(root):
    result = []
    for directory, dirs, files in os.walk(root, followlinks=False):
        for name in sorted(dirs + files):
            path = Path(directory) / name
            relative = path.relative_to(root).as_posix()
            if relative == 'lock.json':
                continue
            metadata = path.lstat()
            item = {'path': relative, 'mode': stat.S_IMODE(metadata.st_mode)}
            if stat.S_ISLNK(metadata.st_mode):
                item.update(type='symlink', target=os.readlink(path))
            elif stat.S_ISREG(metadata.st_mode):
                item.update(type='file', bytes=metadata.st_size, sha256=digest(path))
            elif stat.S_ISDIR(metadata.st_mode):
                item.update(type='directory')
            else:
                raise ValueError(f'Unexpected special build input: {relative}')
            result.append(item)
    return sorted(result, key=lambda item: item['path'])


def source_digest(items):
    source = [dict(item, path=item['path'].removeprefix('project/'))
              for item in items if item['path'].startswith('project/')]
    return hashlib.sha256(json.dumps(source, sort_keys=True, separators=(',', ':')).encode()).hexdigest()


def new_directory(path):
    if path.exists() or path.is_symlink():
        raise ValueError(f'Output already exists: {path}')
    path.mkdir(mode=0o700, parents=False)


def run(*args, **kwargs):
    return subprocess.run(list(map(str, args)), check=True, **kwargs)


def check_pin():
    pin = (project / 'VOID_PACKAGES_COMMIT').read_text().strip()
    void = project / 'vendor/void-packages'
    actual = subprocess.check_output(['git', '-C', str(void), 'rev-parse', 'HEAD'], text=True).strip()
    if pin != actual:
        raise ValueError('Void checkout differs from its pin')
    run('git', '-C', void, 'diff', '--exit-code', 'HEAD', '--')
    build = project / '.host/void-packages'
    if subprocess.check_output(['git', '-C', str(build), 'rev-parse', 'HEAD'], text=True).strip() != pin:
        raise ValueError('Void build checkout differs from its pin')
    run('git', '-C', build, 'diff', '--exit-code', 'HEAD', '--')
    if subprocess.check_output(['git', '-C', str(void), 'status', '--porcelain', '--untracked-files=all'], text=True).strip():
        raise ValueError('Void upstream submodule is dirty')
    return build, pin


def add_package(source, destination):
    target = destination / source.name
    if target.exists():
        if digest(target) != digest(source):
            raise ValueError(f'Conflicting same-version package inputs: {source.name}')
        return
    copy(source, target)


def index(directory, arch):
    packages = sorted([*directory.glob(f'*.{arch}.xbps'), *directory.glob('*.noarch.xbps')])
    if not packages:
        raise ValueError(f'No {arch} packages in {directory}')
    run(project / '.host/xbps/usr/bin/xbps-rindex', '-fa', *packages,
        env=dict(os.environ, XBPS_ARCH=arch))


def create(destination):
    void, pin = check_pin()
    epoch = int(subprocess.check_output(['git', '-C', str(void), 'show', '-s', '--format=%ct', 'HEAD']))
    # All profiles must exist before a freeze. Their exact selected packages
    # become independent repositories, excluding locally rebuilt FDS packages.
    roots = {}
    for profile in ('cli', 'development', 'recovery'):
        archive = (project / f'out/rootfs-{profile}.tar').resolve(strict=True)
        roots[profile] = archive.parent
        if not (archive.parent / 'packages.json').is_file():
            raise ValueError(f'Missing selected-package inventory for {profile}')
    new_directory(destination)
    for name in OWNED:
        copy(project / name, destination / 'project' / name)
    run('git', '-C', void, 'bundle', 'create', destination / 'void.bundle', 'HEAD')
    # Bundles carry commit objects, but do not retain a shallow checkout's
    # boundary file. Without it, history-aware commands follow missing parents.
    shallow = Path(subprocess.check_output(
        ['git', '-C', str(void), 'rev-parse', '--git-path', 'shallow'], text=True).strip())
    if not shallow.is_absolute():
        shallow = void / shallow
    if shallow.is_file():
        copy(shallow, destination / 'void-shallow')

    print('Freezing locked Cargo sources and the selected Rust toolchain', flush=True)
    # Fetching is a separate, explicit preparation step. An incomplete cache
    # fails here rather than extending the lock during the snapshot.
    run('cargo', 'vendor', '--locked', '--offline', '--versioned-dirs',
        destination / 'cargo-vendor', cwd=project, stdout=subprocess.DEVNULL)
    toolchain = Path(subprocess.check_output(['rustup', 'which', 'rustc'], cwd=project, text=True).strip()).parent.parent
    copy(toolchain, destination / 'rust-toolchain')
    toolchain_version = tomllib.loads((project / 'rust-toolchain.toml').read_text())['toolchain']['channel']
    copy(project / '.host/xbps', destination / 'xbps')
    copy(project / '.host/image-tools', destination / 'image-tools')
    # This is the explicit, project-local binary build environment. Exclude
    # temporary build trees, mounted paths, home directories and download caches.
    master = void / 'masterdir-x86_64'
    for name in ('usr', 'etc', 'bin', 'sbin', 'lib', 'lib32', 'lib64',
                 '.xbps_chroot_init', '.xbps_chroot_configured', '.xbps-aarch64-done'):
        path = master / name
        if path.exists() or path.is_symlink():
            copy(path, destination / 'masterdir' / name)
    copy(master / 'var/db/xbps', destination / 'masterdir/var/db/xbps')
    for name in ('dev', 'sys', 'tmp', 'proc', 'host', 'boot', 'void-packages',
                 'builddir', 'destdir', 'home', 'var/cache/xbps', 'var/tmp', 'var/log'):
        (destination / 'masterdir' / name).mkdir(parents=True, exist_ok=True)
    (destination / 'masterdir/tmp').chmod(0o1777)
    (destination / 'masterdir/var/tmp').chmod(0o1777)

    print('Freezing selected target packages, build dependencies and source archives', flush=True)
    package_directories = [void / 'hostdir/repocache-x86_64', void / 'hostdir/repocache-aarch64',
                           master / 'var/cache/xbps', project / 'out/downloads',
                           project / 'out/cache/image-tools', project / 'out/cache/rootfs']
    build = destination / 'repositories/build'
    build.mkdir(parents=True)
    for directory in package_directories:
        for package in sorted(directory.glob('*.xbps')):
            if not package.name.startswith('fds-'):
                add_package(package, build)
    for profile, root in roots.items():
        packages = json.loads((root / 'packages.json').read_text())
        repo = destination / 'repositories' / profile
        repo.mkdir()
        for package in packages:
            source = root / 'packages' / package['filename']
            if digest(source) != package['sha256']:
                raise ValueError(f'Selected {profile} package changed: {source.name}')
            if package['name'].startswith('fds-'):
                continue
            add_package(source, repo)
            add_package(source, build)
        copy(root / 'packages.json', destination / 'profiles' / f'{profile}.json')
        index(repo, 'aarch64')
    for arch in ('x86_64', 'aarch64'):
        index(build, arch)
    source_cache = void / 'hostdir/sources'
    for path in sorted(source_cache.rglob('*')):
        if path.is_file() and path.name.endswith(('.tar.gz', '.tar.xz', '.tar.bz2', '.tar.zst', '.tgz', '.zip')):
            copy(path, destination / 'sources' / path.relative_to(source_cache))
    for category in ('boot', 'vm', 'image-tools'):
        for path in sorted((project / 'out/cache' / category).glob('*.xbps')):
            copy(path, destination / 'cache' / category / path.name)
    host_config = dict(line.split('=', 1) for line in (project / 'config/host-tools.conf').read_text().splitlines()
                       if line and not line.startswith('#'))
    host_archive = project / 'out/downloads' / host_config['XBPS_STATIC_URL'].rsplit('/', 1)[1]
    if digest(host_archive) != host_config['XBPS_STATIC_SHA256']:
        raise ValueError('Static XBPS archive differs from its configured digest')
    copy(host_archive, destination / 'downloads' / host_archive.name)
    eeprom = json.loads((project / 'config/eeprom/inputs.json').read_text())
    for item in eeprom['files']:
        path = project / '.host/eeprom' / eeprom['commit'] / item['name']
        if digest(path) != item['sha256']:
            raise ValueError(f'EEPROM cache differs from pin: {item["name"]}')
        copy(path, destination / 'eeprom' / eeprom['commit'] / item['name'])
    host = {}
    for name, command in [('kernel', ['uname', '-srvmo']), ('python', ['python3', '--version']),
                          ('bubblewrap', ['bwrap', '--version']), ('git', ['git', '--version']),
                          ('tar', ['tar', '--version']), ('rustup', ['rustup', '--version'])]:
        host[name] = subprocess.check_output(command, text=True, stderr=subprocess.DEVNULL).splitlines()[0]
    files = inventory(destination)
    lock = {'format': 1, 'version': '0.1.0', 'void_commit': pin, 'source_epoch': epoch,
            'source_sha256': source_digest(files), 'rust_toolchain': toolchain_version,
            'host_prerequisites': host, 'files': files}
    (destination / 'lock.json').write_text(json.dumps(lock, indent=2) + '\n')
    print(f'PASS: frozen {len(files)} entries; source SHA-256 {lock["source_sha256"]}: {destination}')


def verify(directory):
    if directory.is_symlink() or not directory.is_dir():
        raise ValueError('Frozen input directory must be a real directory')
    lockfile = directory / 'lock.json'
    if lockfile.is_symlink() or not lockfile.is_file() or lockfile.stat().st_size > 64 * 1024 * 1024:
        raise ValueError('Missing or invalid frozen input lock')
    lock = json.loads(lockfile.read_text())
    if lock.get('format') != 1 or lock.get('version') != '0.1.0':
        raise ValueError('Unsupported frozen input format')
    if (set(lock) != {'format', 'version', 'void_commit', 'source_epoch', 'source_sha256',
                     'rust_toolchain', 'host_prerequisites', 'files'}
            or not re.fullmatch(r'[0-9a-f]{40}', lock['void_commit'])
            or not re.fullmatch(r'[0-9a-f]{64}', lock['source_sha256'])
            or not re.fullmatch(r'[0-9]+\.[0-9]+\.[0-9]+', lock['rust_toolchain'])
            or not isinstance(lock['source_epoch'], int) or not 0 < lock['source_epoch'] <= 4102444800):
        raise ValueError('Invalid frozen input lock metadata')
    files = inventory(directory)
    if lock['files'] != files or source_digest(files) != lock['source_sha256']:
        raise ValueError('Frozen input content, modes, paths or source identity changed')
    print(f'PASS: frozen input lock verified: {directory}', flush=True)
    return lock


def restore(source, destination):
    if any(character.isspace() for character in str(destination)):
        raise ValueError('The restored checkout path must not contain whitespace')
    lock = verify(source)
    new_directory(destination)
    for path in sorted((source / 'project').iterdir()):
        copy(path, destination / path.name)
    void = destination / 'vendor/void-packages'
    void.parent.mkdir()
    run('git', 'clone', '--quiet', source / 'void.bundle', void)
    if (source / 'void-shallow').is_file():
        copy(source / 'void-shallow', void / '.git/shallow')
    run('git', '-C', void, 'checkout', '--quiet', '--detach', lock['void_commit'])
    epoch = int(subprocess.check_output(['git', '-C', str(void), 'show', '-s', '--format=%ct', 'HEAD']))
    if epoch != lock['source_epoch']:
        raise ValueError('Restored Void commit timestamp differs from the snapshot')
    # Older frozen projects expect caches under vendor; their saved sources and
    # release inputs remain unchanged. New projects use the separate workspace.
    prepare = destination / 'tools/prepare-void-workspace'
    if prepare.is_file():
        run(sys.executable, prepare)
        void = destination / '.host/void-packages'
    copy(source / 'masterdir', void / 'masterdir-x86_64')
    copy(source / 'sources', void / 'hostdir/sources')
    copy(source / 'repositories/build', void / 'hostdir/frozen-repository')
    (void / 'hostdir/binpkgs').mkdir()
    # xbps-src copies custom local repositories to both its host and cross
    # package configurations while -N excludes every upstream remote.
    custom = void / 'etc/xbps.d/custom'
    custom.mkdir(parents=True, exist_ok=True)
    # The pinned xbps-src regenerates the cross configuration from *local*.conf
    # after initially copying custom files, so the filename must contain local.
    (custom / '05-fds-frozen-local.conf').write_text('repository=/host/frozen-repository\n')
    conf = void / 'masterdir-x86_64/etc/xbps.d'
    shutil.rmtree(conf)
    conf.mkdir()
    (conf / '00-repository-main.conf').symlink_to('/dev/null')
    (conf / '05-fds-frozen-local.conf').write_text('repository=/host/frozen-repository\n')
    for name in ('xbps', 'image-tools', 'eeprom'):
        copy(source / name, destination / '.host' / name)
    for profile in ('cli', 'development', 'recovery'):
        copy(source / 'repositories' / profile, destination / '.host/frozen' / profile)
    copy(source / 'cache', destination / 'out/cache')
    copy(source / 'downloads', destination / 'out/downloads')
    copy(source / 'lock.json', destination / '.host/frozen/lock.json')
    copy(source / 'cargo-vendor', destination / '.host/repro-cargo/registry/src/fds-frozen')
    cargo = destination / '.host/repro-cargo'
    vendor = cargo / 'registry/src/fds-frozen'
    (cargo / 'config.toml').write_text('[source.crates-io]\nreplace-with="fds-frozen"\n'
                                      f'[source.fds-frozen]\ndirectory={json.dumps(str(vendor))}\n')
    rustup = destination / '.host/repro-rustup'
    toolchain = lock['rust_toolchain'] + '-x86_64-unknown-linux-gnu'
    copy(source / 'rust-toolchain', rustup / 'toolchains' / toolchain)
    (rustup / 'settings.toml').write_text(f'version = "12"\ndefault_toolchain = "{toolchain}"\nprofile = "minimal"\n')
    (destination / 'out/logs').mkdir(exist_ok=True)
    print(f'PASS: fresh isolated build tree restored: {destination}')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    subs = parser.add_subparsers(dest='command', required=True)
    for name in ('create', 'verify'):
        sub = subs.add_parser(name)
        sub.add_argument('directory', type=Path)
    sub = subs.add_parser('restore')
    sub.add_argument('directory', type=Path)
    sub.add_argument('destination', type=Path)
    args = parser.parse_args()
    # Resolve only parent components, so a symlink at the final path remains
    # visible to the new-output and verification guards.
    directory = args.directory.absolute()
    try:
        if args.command == 'create':
            # Hold the same locks as the composers for the entire snapshot.
            # Standalone VM/kernel commands must still be run sequentially.
            with (project / 'out/.rootfs.lock').open('a') as root_lock, (project / 'out/.base-packages.lock').open('a') as package_lock:
                fcntl.flock(root_lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
                fcntl.flock(package_lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
                create(directory)
        elif args.command == 'verify':
            verify(directory)
        else:
            restore(directory, args.destination.absolute())
    except BlockingIOError:
        sys.exit('ERROR: a rootfs or base-package build is active; finish it before freezing inputs')
    except (OSError, ValueError, KeyError, subprocess.CalledProcessError) as error:
        sys.exit(f'ERROR: {error}')


if __name__ == '__main__':
    main()
