#!/usr/bin/env python3
"""Remove obsolete FDS workspaces, retaining published outputs and saved inputs."""
import argparse
import contextlib
import fcntl
import os
from pathlib import Path
import re
import shutil
import stat
import subprocess
import sys


# Only names produced by repository builders/tests belong here. Never sweep out/*:
# it also holds signed releases, frozen inputs, user cartridges and DATA overlays.
PREFIXES = (
    'rootfs-build', 'kernel-build', 'system-build', 'initramfs-build',
    'boot-build', 'recovery-build', 'internal-build',
    'eeprom-production', 'eeprom-development',
    'cartridged-source', 'init-source', 'dasung-services', 'dasung-s6',
    'rust-notices', 'verify-package', 'dasung-check', 'dasung-s6-run',
    'm1-checks', 'm2-vm', 'm4-vm', 'm5-vm', 'm6-vm', 'm7-vm', 'm8-vm',
    'm9-images', 'm9-vm', 'm9-system', 'm10-vm', 'm11-vm', 'm11-faults',
    'm12-clock', 'm12-development', 'm12-eeprom', 'm12-internal',
    'm12-recovery', 'm12-release-contracts', 'm12-signing',
    'rust-profiles', 'workstation-images', 'emu-test', 'arch-package-test', 'version-image-test',
)
GENERATED = re.compile(r'(?:' + '|'.join(PREFIXES) + r')\.[A-Za-z0-9_-]{6,8}')
POINTERS = (
    'out/arch-package-current.txt',
    'out/version-image-current.txt',
    'out/workstation-images-current.txt',
    'out/workstation-emulator-current.txt',
    'out/manifests/dasung-s6-database.txt',
)
LOCKS = ('.clean.lock', '.rootfs.lock', '.base-packages.lock', '.image-tools.lock', '.workstation-package.lock')


def within(path, parent):
    return path.is_relative_to(parent)


def workspace_for(project, path):
    """Find the only possible candidate ancestor without scanning every run."""
    try:
        parts = path.relative_to(project).parts
    except ValueError:
        return None
    if parts and parts[0] == 'target':
        return project / 'target'
    if len(parts) >= 2 and parts[0] == 'out':
        return project / 'out' / parts[1]
    return None


def candidates(project):
    out = project / 'out'
    if out.is_symlink():
        raise ValueError('out is a symlink; refusing to clean an external output tree')
    if not out.exists():
        entries = []
    else:
        entries = sorted(out.iterdir())
    found = {p for p in entries if GENERATED.fullmatch(p.name)
             and p.is_dir() and not p.is_symlink()}
    target = project / 'target'
    if target.is_dir() and not target.is_symlink():
        found.add(target)
    protected = {}

    def retain(path, reason):
        # Retain both a lexical target and its resolved target, including a link
        # reached through a directory alias. Out-of-tree links never add candidates.
        lexical = Path(os.path.abspath(path))
        for value in (lexical, lexical.resolve()):
            candidate = workspace_for(project, value)
            if candidate in found:
                protected.setdefault(candidate, reason)

    # Git-tracked content is never build trash, even if named like a workspace.
    tracked = subprocess.check_output(
        ['git', '-C', str(project), 'ls-files', '-z', '--', 'out', 'target'])
    for name in os.fsdecode(tracked).split('\0'):
        if name:
            retain(project / name, 'contains tracked content')
    for path in found:
        if os.path.lexists(path / '.fds-keep'):
            protected[path] = '.fds-keep marker'

    manifests = out / 'manifests'
    links = entries + (list(manifests.iterdir())
                       if manifests.is_dir() and not manifests.is_symlink() else [])
    for link in links:
        if link.is_symlink():
            retain(link.parent / link.readlink(), f'published link {link.relative_to(project)}')
    for name in POINTERS:
        pointer = project / name
        if pointer.is_file():
            value = pointer.read_text().strip()
            if value:
                retain(project / value, f'current pointer {name}')

    # These suites consume the newest image fixture by mtime, without a link.
    image_runs = [p for p in found if p.name.startswith('m9-images.')]
    if image_runs:
        protected[max(image_runs, key=lambda p: p.stat().st_mtime_ns)] = 'latest media-image fixture'

    # A retained workspace can contain links to another generated workspace.
    scanned = set()
    while set(protected) - scanned:
        for path in set(protected) - scanned:
            scanned.add(path)
            for directory, dirs, files in os.walk(path, followlinks=False):
                for name in dirs + files:
                    link = Path(directory) / name
                    if link.is_symlink():
                        retain(link.parent / link.readlink(), f'dependency of {path.name}')
    return sorted(found - set(protected)), protected


def check_mounts(paths, mountinfo=None):
    if mountinfo is None:
        mountinfo = Path('/proc/self/mountinfo').read_text()
    for line in mountinfo.splitlines():
        value = re.sub(r'\\([0-7]{3})', lambda m: chr(int(m[1], 8)), line.split()[4])
        mount = Path(value)
        if any(within(mount, path) for path in paths):
            raise ValueError(f'mounted path {mount}; unmount it before cleaning')


def check_processes(project, paths):
    selected = set(paths)
    ancestors = set()
    pid = os.getpid()
    while pid > 0:
        ancestors.add(pid)
        try:
            status = Path(f'/proc/{pid}/status').read_text()
            pid = int(re.search(r'^PPid:\s+(\d+)', status, re.M)[1])
        except FileNotFoundError:
            break
    for proc in Path('/proc').iterdir():
        if not proc.name.isdigit() or int(proc.name) in ancestors:
            continue
        command = ''
        try:
            if proc.stat().st_uid != os.getuid():
                continue
            command = (proc / 'cmdline').read_bytes().replace(b'\0', b' ').decode(errors='replace')
            cwd = (proc / 'cwd').resolve(strict=True)
            # Catch a build before it has opened any candidate output files.
            name = (proc / 'comm').read_text().strip()
            building = (name in ('make', 'gmake', 'cargo', 'rustc', 'rustup', 'bwrap')
                        or name.startswith(('qemu-', 'xbps-'))
                        or re.search(r'(?:tools/|tests/integration/|image/build-)', command))
            if building and (within(cwd, project) or str(project) in command):
                raise ValueError(f'active build/test process {proc.name} ({name}); stop it before cleaning')
            for handle in [proc / 'cwd', proc / 'exe', *(proc / 'fd').iterdir()]:
                try:
                    value = Path(os.readlink(handle))
                except FileNotFoundError:
                    continue
                if workspace_for(project, value) in selected:
                    raise ValueError(f'process {proc.name} is using {value}; stop it before cleaning')
        except (FileNotFoundError, ProcessLookupError):
            continue
        except PermissionError as error:
            # Desktop session helpers may be nondumpable even for the same UID.
            # Refuse an inaccessible process naming this checkout; unrelated
            # protected processes must not make cleanup permanently unusable.
            if str(project) in command:
                raise ValueError(f'cannot inspect project process {proc.name}: {error}') from error


@contextlib.contextmanager
def locks(project):
    with contextlib.ExitStack() as stack:
        out = project / 'out'
        if out.is_symlink():
            raise ValueError('out is a symlink; refusing cleanup')
        out.mkdir(exist_ok=True)
        for name in LOCKS:
            fd = os.open(out / name, os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW, 0o600)
            stream = stack.enter_context(os.fdopen(fd, 'w'))
            try:
                fcntl.flock(stream, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except BlockingIOError as error:
                raise ValueError(f'{name} is locked; another build or cleanup is running') from error
        yield


def footprint(paths):
    if not paths:
        return 0
    result = subprocess.check_output(['du', '-sx', '-B1', '--', *map(str, paths)], text=True)
    # One invocation counts hardlinks once. Reflinks and snapshots can still
    # share physical extents, so this is not a promise of filesystem free space.
    return sum(int(line.split('\t', 1)[0]) for line in result.splitlines())


def remove_tree(path):
    if path.is_symlink() or not path.is_dir():
        raise ValueError(f'workspace changed since preview: {path}')
    # Some guest directories are deliberately read-only. Change only directory
    # permissions; file modes may belong to hardlinks in retained rootfs trees.
    def walk_error(error):
        raise error
    os.chmod(path, path.stat().st_mode | 0o700, follow_symlinks=False)
    for directory, dirs, _ in os.walk(path, topdown=True, followlinks=False, onerror=walk_error):
        for item in [Path(directory), *(Path(directory) / name for name in dirs)]:
            mode = item.lstat().st_mode
            if stat.S_ISDIR(mode) and mode & 0o700 != 0o700:
                os.chmod(item, mode | 0o700, follow_symlinks=False)
    shutil.rmtree(path)


def clean(project, dry_run):
    paths, protected = candidates(project)
    identities = {path: (path.stat().st_dev, path.stat().st_ino) for path in paths}
    check_mounts(paths)
    if not dry_run:
        check_processes(project, paths)
    for path in paths:
        print(f'{"WOULD REMOVE" if dry_run else "REMOVE"} {path.relative_to(project)}', flush=True)
    amount = footprint(paths)
    print(f'{len(paths)} disposable directories; allocated footprint {amount / 1024**3:.2f} GiB.', flush=True)
    print('Physical space recovered can be smaller with hardlinks, reflinks or snapshots.', flush=True)
    for path, reason in sorted(protected.items()):
        print(f'KEEP {path.relative_to(project)} ({reason})', flush=True)
    print('Saved releases/inputs/rebuild trees, unknown outputs, caches, out/logs and out/manifests are kept.', flush=True)
    if dry_run:
        print('Preview only. Run make clean to apply.', flush=True)
        return
    check_processes(project, paths)
    check_mounts(paths)
    if candidates(project)[0] != paths or any(
            path.is_symlink() or (path.stat().st_dev, path.stat().st_ino) != identities[path]
            for path in paths):
        raise ValueError('output selection changed during inspection; rerun cleanup after stopping builds')
    before = shutil.disk_usage(project).free
    for path in paths:
        remove_tree(path)
    after = shutil.disk_usage(project).free
    print(f'PASS: removed {len(paths)} disposable directories; '
          f'filesystem free-space change {(after - before) / 1024**3:+.2f} GiB.', flush=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--dry-run', action='store_true', help='list the exact cleanup selection without deleting anything')
    args = parser.parse_args()
    project = Path(__file__).resolve().parents[1]
    try:
        if args.dry_run:
            clean(project, True)
        else:
            with locks(project):
                clean(project, False)
    except (OSError, ValueError, subprocess.CalledProcessError) as error:
        parser.exit(1, f'ERROR: {error}\n')


if __name__ == '__main__':
    main()
