#!/usr/bin/env python3
"""#58 local fleet repair. Default preflight; explicit pinned apply or rollback.

Only two source directories and the Resume-hotfix launcher may change. Settings,
credentials, state, sessions and enforcement extensions are never written.
Linux renameat2 exchanges keep originals at the private backup paths, including
when interrupted. A receipt can recover a partially applied transaction.
"""
import argparse
from contextlib import contextmanager
import ctypes
import fcntl
import hashlib
import json
import os
from pathlib import Path
import stat
import uuid

SHARED = '.pi/extensions/goal'
ALIASES = ('fleet/extensions/goal', 'fleet/agents/velma/.pi/extensions/goal')
LAUNCHER = 'fleet/bin/launch-seat.sh'
DEPENDENCIES = ('.pi/extensions/mosaic-core', 'fleet/extensions/mosaic-core',
                'fleet/agents/velma/.pi/extensions/mosaic-core',
                'fleet/extensions/wrapper-guard/index.ts', 'tools/unslop-hook/extension.ts')
HOTFIX_SHA = 'dd9e5ece5f1a86cc286027668560198bb7a44d2ab891f7b1ab4f89d7379ca0e6'
ORIGINAL_SHA = '9352feed0acf9d449c26c0556ba00aba1c65ded43d376930a39ff6b5cee21986'


def require(condition, message):
    if not condition:
        raise RuntimeError(message)


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


def snapshot(path):
    """Ordinary source files only. Do not call this on goal-state directories."""
    require(path.resolve() == path, f'non-canonical source path: {path}')
    result = {}
    paths = [path] + (sorted(path.rglob('*')) if path.is_dir() else [])
    for p in paths:
        s = p.lstat()
        kind = 'dir' if stat.S_ISDIR(s.st_mode) else 'file' if stat.S_ISREG(s.st_mode) else None
        require(kind is not None, f'symlink/special source file refused: {p}')
        require(s.st_uid == os.getuid(), f'foreign-owned source: {p}')
        result[str(p.relative_to(path))] = {'type': kind, 'mode': stat.S_IMODE(s.st_mode)}
        if kind == 'file':
            result[str(p.relative_to(path))]['sha256'] = sha(p.read_bytes())
    return result


def inventory(root):
    """Pin configuration without copying its contents or reading auth/state."""
    paths = []
    for group in ('agents', 'roles'):
        paths += list((root / 'fleet' / group).glob('*/.pi/agent/settings.json'))
    paths += [root / 'fleet/templates/pi-settings.json.template']
    result = {}
    for p in sorted(paths):
        data = p.read_bytes()
        entry = {'sha256': sha(data), 'resolved': str(p.resolve()),
                 'link': os.readlink(p) if p.is_symlink() else None}
        if p.name == 'settings.json':
            settings = json.loads(data)
            entry['extensions'] = settings.get('extensions', [])
            require(not settings.get('packages'), f'unreviewed package source: {p}')
        result[str(p.relative_to(root))] = entry
    # Environment files may carry credentials. Pin filesystem identity/change
    # metadata only; never read or copy their contents for this source repair.
    for p in sorted((root / 'fleet/agents').glob('*/launch.env')):
        s = p.stat()
        result[str(p.relative_to(root))] = {
            'resolved': str(p.resolve()), 'link': os.readlink(p) if p.is_symlink() else None,
            'device': s.st_dev, 'inode': s.st_ino, 'size': s.st_size,
            'mtime_ns': s.st_mtime_ns, 'ctime_ns': s.st_ctime_ns,
        }
    return result


def restored_launcher(data):
    require(sha(data) == HOTFIX_SHA, 'launcher is not the reviewed Resume hotfix')
    start = data.index(b'# Issue #57: Resume uses the operator-selected shared NG goal')
    end = data.index(b'for f in "${PROMPT_FILES[@]}"', start)
    result = data[:start] + data[end:]
    require(sha(result) == ORIGINAL_SHA, 'launcher restoration differs from original')
    return result


def prepare(root):
    require(root.is_dir() and root.resolve() == root, 'root must be a canonical directory')
    after = restored_launcher((root / LAUNCHER).read_bytes())
    before = {p: snapshot(root / p) for p in (*ALIASES, LAUNCHER)}
    launcher_after = {'.': {**before[LAUNCHER]['.'], 'sha256': sha(after)}}
    return {'version': 1, 'root': str(root), 'shared': snapshot(root / SHARED),
            'dependencies': {p: snapshot(root / p) for p in DEPENDENCIES},
            'before': before, 'launcherAfter': launcher_after, 'inventory': inventory(root)}


def validate_plan(plan):
    require(plan['version'] == 1, 'unsupported plan')
    root = Path(plan['root'])
    require(root.is_absolute() and root.resolve() == root, 'invalid root')
    require(set(plan['before']) == set((*ALIASES, LAUNCHER)), 'invalid mutation scope')
    require(plan['before'][LAUNCHER]['.']['sha256'] == HOTFIX_SHA, 'unreviewed launcher pin')
    require(plan['launcherAfter']['.']['sha256'] == ORIGINAL_SHA, 'unreviewed restoration pin')
    return root


def verify_common(plan):
    root = validate_plan(plan)
    require(snapshot(root / SHARED) == plan['shared'], 'shared source drift')
    require(set(plan['dependencies']) == set(DEPENDENCIES), 'invalid dependency scope')
    for p in DEPENDENCIES:
        require(snapshot(root / p) == plan['dependencies'][p], f'dependency drift: {p}')
    require(inventory(root) == plan['inventory'], 'configuration inventory drift')
    return root


def verify_before(plan):
    root = verify_common(plan)
    for p in (*ALIASES, LAUNCHER):
        require(snapshot(root / p) == plan['before'][p], f'before drift: {p}')


def is_alias(path, shared):
    return path.is_symlink() and os.readlink(path) == str(shared) and path.resolve() == shared


def verify_after(plan):
    root = verify_common(plan)
    for p in ALIASES:
        require(is_alias(root / p, root / SHARED), f'alias drift: {p}')
    require(snapshot(root / LAUNCHER) == plan['launcherAfter'], 'launcher after drift')


def exchange(a, b):
    fn = ctypes.CDLL(None, use_errno=True).renameat2
    fn.argtypes = [ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_char_p, ctypes.c_uint]
    fn.restype = ctypes.c_int
    if fn(-100, os.fsencode(a), -100, os.fsencode(b), 2):
        error = ctypes.get_errno()
        raise OSError(error, os.strerror(error))


def write_once(path, data):
    with path.open('xb') as f:
        f.write(data)
        f.flush()
        os.fsync(f.fileno())


def stage_path(backup, relative):
    return backup / ('launcher' if relative == LAUNCHER else 'legacy' if relative == ALIASES[0] else 'velma')


def recover(plan, backup):
    """Validate every pair before reverting any. Handles partial exchanges."""
    root = verify_common(plan)
    require(backup.parent == root / '.pi/goal-backups' and backup.resolve() == backup,
            'backup outside approved directory')
    swapped = []
    for p in (*ALIASES, LAUNCHER):
        target, stage = root / p, stage_path(backup, p)
        if p in ALIASES:
            target_after = is_alias(target, root / SHARED)
            stage_after = is_alias(stage, root / SHARED)
        else:
            target_after = snapshot(target) == plan['launcherAfter']
            stage_after = snapshot(stage) == plan['launcherAfter']
        if target_after:
            require(snapshot(stage) == plan['before'][p], f'rollback backup drift: {p}')
            swapped.append(p)
        else:
            require(snapshot(target) == plan['before'][p] and stage_after, f'rollback pair drift: {p}')
    for p in reversed(swapped):
        exchange(root / p, stage_path(backup, p))
    verify_before(plan)
    write_once(backup / ('rollback-' + uuid.uuid4().hex + '.json'),
               json.dumps({'status': 'rolled back', 'restored': swapped}).encode())


def apply(plan, raw):
    verify_before(plan)
    root = Path(plan['root'])
    backup_root = root / '.pi/goal-backups'
    backup_root.mkdir(mode=0o700, parents=True, exist_ok=True)
    require(backup_root.resolve() == backup_root, 'backup directory alias refused')
    backup = backup_root / ('goal58-' + uuid.uuid4().hex)
    backup.mkdir(mode=0o700)
    for p in ALIASES:
        stage_path(backup, p).symlink_to(root / SHARED, target_is_directory=True)
    stage = stage_path(backup, LAUNCHER)
    write_once(stage, restored_launcher((root / LAUNCHER).read_bytes()))
    os.chmod(stage, plan['before'][LAUNCHER]['.']['mode'])
    for p in (*ALIASES, LAUNCHER):
        require((root / p).lstat().st_dev == stage_path(backup, p).lstat().st_dev, 'cross-device exchange refused')
    write_once(backup / 'planned.json', raw)
    print(f'Recovery receipt: {backup}', flush=True)
    try:
        verify_before(plan)
        for p in (*ALIASES, LAUNCHER):
            require(snapshot(root / p) == plan['before'][p], f'concurrent drift: {p}')
            exchange(root / p, stage_path(backup, p))
        verify_after(plan)
        for p in (*ALIASES, LAUNCHER):
            require(snapshot(stage_path(backup, p)) == plan['before'][p], f'backup mismatch: {p}')
        write_once(backup / 'deployed.json', json.dumps({'status': 'deployed', 'planSha256': sha(raw)}).encode())
    except BaseException:
        recover(plan, backup)
        raise
    return backup


@contextmanager
def deployment_lock(plan):
    root = verify_common(plan)
    directory = root / '.pi/goal-backups'
    directory.mkdir(mode=0o700, parents=True, exist_ok=True)
    require(directory.resolve() == directory, 'backup directory alias refused')
    fd = os.open(directory / 'goal58.lock', os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600)
    try:
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        yield
    finally:
        os.close(fd)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('mode', nargs='?', default='preflight', choices=['prepare', 'preflight', 'apply', 'verify', 'rollback'])
    parser.add_argument('--root', type=Path)
    parser.add_argument('--plan', type=Path, required=True)
    parser.add_argument('--sha256')
    parser.add_argument('--backup', type=Path)
    args = parser.parse_args()
    if args.mode == 'prepare':
        require(args.root is not None, '--root required')
        plan = prepare(args.root)
        raw = (json.dumps(plan, indent=2, sort_keys=True) + '\n').encode()
        write_once(args.plan, raw)
        print('Prepared plan SHA256:', sha(raw))
        return
    raw = args.plan.read_bytes()
    require(sha(raw) == args.sha256, 'reviewed plan SHA256 required')
    plan = json.loads(raw)
    if args.mode == 'preflight':
        verify_before(plan)
    elif args.mode == 'apply':
        with deployment_lock(plan):
            print('Deployed:', apply(plan, raw))
    elif args.mode == 'verify':
        verify_after(plan)
    else:
        require(args.backup is not None, '--backup required')
        require((args.backup / 'planned.json').read_bytes() == raw, 'backup receipt mismatch')
        with deployment_lock(plan):
            recover(plan, args.backup)
    print('PASS', args.mode)


if __name__ == '__main__':
    main()