#!/usr/bin/env python3
"""Durable component transaction; also runnable by old workers and at boot.

The journal is mirrored into the original and candidate data roots before any
system mutation. The worker may rename either root; one journal always survives.
The recovery executable is independent of releases that the old worker deletes.
Only stdlib imports: recovery must not depend on a candidate venv.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import stat
import subprocess
import time

DATA = Path('/var/lib/matrix-screen-controller')
TARGET = Path('/opt/matrix-screen-controller')
RELEASES = Path('/opt/matrix-screen-controller.releases')
RUNTIME = Path('/run/matrix-screen-controller')
HELPER = Path('/opt/matrix-screen-controller-component-recovery.py')
RECOVERY_UNIT = Path('/etc/systemd/system/matrix-screen-component-recovery.service')
ACCOUNT_POLICY = Path('/etc/sudoers.d/90-matrix-screen-controller-account')
WORK_ROOT = Path('/opt/matrix-screen-controller-ota')
RUNTIME_GUARD = Path('/run/systemd/system/matrix-screen-controller.service.d/90-matrix-ota-runtime.conf')
RUNTIME_GUARD_BODY = b'# Managed by the OTA component transaction\n[Service]\nRuntimeDirectoryPreserve=yes\n'
SERVICE = 'matrix-screen-controller.service'
WORKER = 'matrix-screen-controller-ota.service'
FRP = 'matrix-screen-frpc.service'
JOURNAL = Path('ota/component-transaction')
FILES = {
    'frpc': Path('/usr/local/bin/frpc'),
    'frpc-unit': Path('/etc/systemd/system/matrix-screen-frpc.service'),
    'frpc-dropin': Path('/etc/systemd/system/matrix-screen-frpc.service.d/user.conf'),
    'app-unit': Path('/etc/systemd/system/matrix-screen-controller.service'),
    'ota-unit': Path('/etc/systemd/system/matrix-screen-controller-ota.service'),
}


def run(args: list[str], *, check=True, **kwargs):
    return subprocess.run(args, check=check, capture_output=True, timeout=60, **kwargs)


def sync_directory(path: Path):
    fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)


def atomic(path: Path, body: bytes, mode=0o600):
    path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_name(path.name + '.tmp')
    with temp.open('wb') as handle:
        os.chmod(temp, mode)
        handle.write(body)
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temp, path)
    sync_directory(path.parent)


def read(path: Path):
    return json.loads(path.read_text(encoding='utf-8'))


def protect_runtime():
    if RUNTIME_GUARD.exists():
        raise RuntimeError('OTA runtime protection path already exists; refusing to overwrite')
    atomic(RUNTIME_GUARD, RUNTIME_GUARD_BODY, 0o644)
    run(['systemctl', 'daemon-reload'])
    value = run(['systemctl', 'show', SERVICE, '--property=RuntimeDirectoryPreserve', '--value']).stdout.strip()
    if value != b'yes':
        raise RuntimeError('OTA runtime directory protection did not take effect; refusing to stop service')


def release_runtime_guard():
    if RUNTIME_GUARD.exists():
        if RUNTIME_GUARD.read_bytes() != RUNTIME_GUARD_BODY:
            raise RuntimeError('OTA runtime protection was modified; preserving it for inspection')
        RUNTIME_GUARD.unlink()
        try:
            RUNTIME_GUARD.parent.rmdir()
        except OSError:
            pass
        run(['systemctl', 'daemon-reload'])


def metadata(path: Path):
    info = path.stat(follow_symlinks=False)
    if stat.S_ISLNK(info.st_mode):
        raise RuntimeError('managed component paths must not be symbolic links')
    return {'mode': stat.S_IMODE(info.st_mode), 'uid': info.st_uid, 'gid': info.st_gid}


def apply_metadata(path: Path, item: dict):
    os.chown(path, item['uid'], item['gid'], follow_symlinks=False)
    path.chmod(item['mode'])


def permission_snapshot(root: Path):
    paths = [root]
    if (root / 'frp').exists():
        paths += [root / 'frp', *(root / 'frp').rglob('*')]
    return {str(p.relative_to(root)): metadata(p) for p in paths}


def restore_permissions(root: Path, saved: dict):
    for name, item in saved.items():
        path = root / name
        if path.exists():
            apply_metadata(path, item)
    # Remove only empty directories created by the installer, never user files.
    for name in ('frp/profiles', 'frp'):
        if name not in saved:
            try:
                (root / name).rmdir()
            except OSError:
                pass


def mirror_permissions(candidate: Path):
    for name, item in permission_snapshot(DATA).items():
        path = candidate / name
        if not path.exists() and (DATA / name).is_dir():
            path.mkdir(parents=True)
        if path.exists():
            apply_metadata(path, item)


def account() -> str:
    import grp
    import pwd
    dropin = FILES['frpc-dropin']
    if dropin.is_file():
        users = re.findall(r'^User=([a-zA-Z0-9_-]+)$', dropin.read_text(encoding='utf-8'), re.M)
        if len(users) == 1:
            try:
                if pwd.getpwnam(users[0]).pw_uid != 0:
                    return users[0]
            except KeyError:
                pass
    if ACCOUNT_POLICY.is_file():
        match = re.fullmatch(r'([a-zA-Z0-9_-]+)\s+ALL=\(ALL:ALL\)\s+ALL\s*',
                             ACCOUNT_POLICY.read_text(encoding='utf-8').strip())
        if match:
            try:
                if pwd.getpwnam(match[1]).pw_uid != 0:
                    return match[1]
            except KeyError:
                pass
        raise RuntimeError('项目维护账户配置无效；尚未修改系统组件')
    group = grp.getgrnam('sudo')
    users = [p.pw_name for p in pwd.getpwall() if 1000 <= p.pw_uid < 65534
             and (p.pw_name in group.gr_mem or p.pw_gid == group.gr_gid)
             and p.pw_shell not in ('/usr/sbin/nologin', '/bin/false')]
    if len(users) != 1:
        raise RuntimeError('无法唯一确定非 root 维护账户，尚未修改系统组件')
    return users[0]


def check_installed(source: Path, version: str):
    import pwd
    import grp
    entries = read(source / 'UPGRADE_POLICY.json')['checkpoints']
    required = {}
    version_tuple = tuple(map(int, version.split('.')))
    for entry in entries:
        if tuple(map(int, entry['version'].split('.'))) <= version_tuple:
            required.update(entry['components'])
    if not required:
        return
    expected = required['frpc']
    binary = FILES['frpc']
    error = 'frp 系统组件缺失或不完整；请先修复软件安装包组件后重试'
    if not binary.is_file() or hashlib.sha256(binary.read_bytes()).hexdigest() != expected['sha256']:
        raise RuntimeError(error)
    if metadata(binary)['mode'] != 0o755 or run([str(binary), '--version']).stdout.decode().strip() != expected['version']:
        raise RuntimeError(error)
    if not FILES['frpc-unit'].is_file() or FILES['frpc-unit'].read_bytes() != (source / 'systemd/matrix-screen-frpc.service').read_bytes():
        raise RuntimeError(error)
    user = account()
    group = grp.getgrgid(pwd.getpwnam(user).pw_gid).gr_name
    wanted = f'[Service]\nUser={user}\nGroup={group}\n'.encode()
    if not FILES['frpc-dropin'].is_file() or FILES['frpc-dropin'].read_bytes() != wanted:
        raise RuntimeError(error)
    for name in ('frpc-unit', 'frpc-dropin'):
        if metadata(FILES[name]) != {'mode': 0o644, 'uid': 0, 'gid': 0}:
            raise RuntimeError(error)
    if binary.stat().st_uid != 0 or binary.stat().st_gid != 0:
        raise RuntimeError(error)
    gid = pwd.getpwnam(user).pw_gid
    if stat.S_IMODE(DATA.stat().st_mode) != 0o711:
        raise RuntimeError(error)
    for name in ('frp', 'frp/profiles'):
        if not (DATA / name).is_dir():
            raise RuntimeError(error)
    for path in [DATA / 'frp', *(DATA / 'frp').rglob('*')]:
        mode = 0o2750 if path.is_dir() else (0o600 if path == DATA / 'frp/active.env' else 0o640)
        if metadata(path) != {'mode': mode, 'uid': 0, 'gid': gid}:
            raise RuntimeError(error)


def begin(source: Path, candidate: Path):
    request_path = RUNTIME / 'ota-request.json'
    # A regular migration / local test is not authorization to touch the host.
    if not request_path.is_file():
        return
    request = read(request_path)
    job = request.get('job_id', '')
    if not re.fullmatch(r'[a-zA-Z0-9_-]+', job):
        raise RuntimeError('invalid component transaction id')
    version = (source / 'VERSION').read_text(encoding='utf-8').strip()
    expected_release = RELEASES / f'{version}-{job}'
    expected_candidate = DATA.with_name(f'matrix-screen-controller.ota.{job}')
    if source != expected_release or candidate != expected_candidate or version != request['target_version']:
        return
    if os.geteuid() != 0 or os.uname().machine != 'aarch64':
        raise RuntimeError('component transaction requires AArch64 root')
    from app.ota.policy import check_upgrade
    from app.ota.versioning import SoftwareVersion
    check_upgrade(SoftwareVersion.parse(request['current_version']), SoftwareVersion.parse(version))
    if RUNTIME_GUARD.exists():
        raise RuntimeError('存在未清理的 OTA 运行目录保护，请先恢复')
    bundle = source / 'system-dependencies/frpc'
    if not bundle.is_dir():
        check_installed(source, version)
        return
    # Verify the pinned binary, not just a self-reported checksum file.
    from app.ota.policy import required_components
    expected = required_components(source, SoftwareVersion.parse(version))['frpc']
    if hashlib.sha256((bundle / 'frpc').read_bytes()).hexdigest() != expected['sha256']:
        raise RuntimeError('frpc payload differs from the registered dependency')
    user = account()
    for root in (DATA, candidate):
        if (root / JOURNAL).exists():
            raise RuntimeError('存在未完成组件事务，请先恢复')
    journal = DATA / JOURNAL
    journal.mkdir(parents=True, mode=0o700)
    record = {'job_id': job, 'version': version, 'old_version': request['current_version'],
              'accepted_at': request.get('accepted_at', ''),
              'previous_target': os.readlink(TARGET) if TARGET.is_symlink() else None,
              'permissions': permission_snapshot(DATA), 'files': {},
              'active': run(['systemctl', 'is-active', '--quiet', FRP], check=False).returncode == 0,
              'enabled': run(['systemctl', 'is-enabled', '--quiet', FRP], check=False).returncode == 0}
    try:
        for name, path in FILES.items():
            record['files'][name] = metadata(path) if path.exists() else None
            if path.exists():
                atomic(journal / name, path.read_bytes())
        atomic(journal / 'state.json', json.dumps(record).encode())
        shutil.copytree(journal, candidate / JOURNAL)
        # copytree itself is not durable across power loss.
        for path in (candidate / JOURNAL).iterdir():
            with path.open('rb') as handle:
                os.fsync(handle.fileno())
        sync_directory(candidate / JOURNAL)
        atomic(HELPER, Path(__file__).read_bytes(), 0o700)
        unit = ('[Unit]\nDescription=Recover interrupted OTA component transaction\n'
                'After=local-fs.target\nBefore=matrix-screen-controller.service\n'
                '[Service]\nType=oneshot\nExecStart=/usr/bin/python3 ' + str(HELPER) + ' recover\n'
                '[Install]\nWantedBy=multi-user.target\n')
        atomic(RECOVERY_UNIT, unit.encode(), 0o644)
        run(['systemctl', 'daemon-reload'])
        run(['systemctl', 'enable', RECOVERY_UNIT.name])
        run(['systemd-run', '--quiet', '--collect', '--unit=matrix-screen-component-watch',
             '/usr/bin/python3', str(HELPER), 'watch'])
        protect_runtime()
        env = os.environ.copy()
        env.update(FRPC_BUNDLE=str(bundle), FRPC_RUN_USER=user)
        run(['/bin/sh', str(source / 'scripts/install_frpc_system.sh')], env=env)
        mirror_permissions(candidate)
        check_installed(source, version)
        print('FRP component transaction prepared; waiting for OTA health result')
    except BaseException:
        # The watcher handles subsequent worker failure. Restore immediately too,
        # so a failed migration never leaves new system files while rolling back.
        if (journal / 'state.json').is_file():
            restore_components(journal, record)
        else:
            shutil.rmtree(journal)
        raise


def restore_components(journal: Path, record: dict):
    run(['systemctl', 'stop', FRP], check=False)
    for name in ('frpc', 'frpc-unit', 'frpc-dropin'):
        path = FILES[name]
        item = record['files'][name]
        if item is None:
            path.unlink(missing_ok=True)
        else:
            atomic(path, (journal / name).read_bytes(), item['mode'])
            apply_metadata(path, item)
    restore_permissions(DATA, record['permissions'])
    run(['systemctl', 'daemon-reload'])
    # A previously absent unit cannot be disabled; avoid treating absence as error.
    if record['files']['frpc-unit'] is not None:
        run(['systemctl', 'enable' if record['enabled'] else 'disable', FRP])
    else:
        run(['systemctl', 'disable', FRP], check=False)
        link = Path('/etc/systemd/system/multi-user.target.wants') / FRP
        link.unlink(missing_ok=True)
    if record['active']:
        run(['systemctl', 'start', FRP])


def locate_journal():
    if (DATA / JOURNAL / 'state.json').is_file():
        return DATA / JOURNAL
    candidates = list(DATA.parent.glob('matrix-screen-controller.rollback.*/ota/component-transaction/state.json'))
    if len(candidates) == 1:
        return candidates[0].parent
    if candidates:
        raise RuntimeError('multiple recovery journals; refusing ambiguous recovery')
    return None


def committed(record: dict) -> bool:
    try:
        result = read(DATA / 'ota/state.json')['last_result']
        return (result['status'] == 'success' and result['target_version'] == record['version']
                and (TARGET / 'VERSION').read_text(encoding='utf-8').strip() == record['version']
                and result['installed_at'] >= record['accepted_at'])
    except (OSError, ValueError, KeyError, TypeError):
        return False


def recover_application(journal: Path, record: dict):
    job = record['job_id']
    release = RELEASES / f"{record['version']}-{job}"
    backup = DATA.with_name(f'matrix-screen-controller.rollback.{job}')
    displaced = RELEASES / f"{record['old_version']}-pre-ota-{job}"
    switched = TARGET.is_symlink() and TARGET.resolve() == release.resolve()
    moved = not TARGET.exists() and (displaced.exists() or record['previous_target'])
    if switched or moved or backup.exists():
        run(['systemctl', 'stop', SERVICE], check=False)
        if backup.exists():
            if DATA.exists():
                failed = DATA.with_name(f'matrix-screen-controller.failed.{job}')
                if failed.exists():
                    raise RuntimeError('failed data recovery path already exists')
                DATA.rename(failed)
                backup.rename(DATA)
                shutil.rmtree(failed)
            else:
                backup.rename(DATA)
            journal = DATA / JOURNAL
        if switched:
            TARGET.unlink()
        if switched or moved:
            if record['previous_target']:
                os.symlink(record['previous_target'], TARGET, target_is_directory=True)
            elif displaced.exists():
                displaced.rename(TARGET)
            else:
                raise RuntimeError('old application is missing; retaining recovery journal')
    # Legacy rollback restores only its main unit, not its OTA unit.
    for name in ('app-unit', 'ota-unit'):
        if record['files'][name]:
            atomic(FILES[name], (journal / name).read_bytes(), record['files'][name]['mode'])
            apply_metadata(FILES[name], record['files'][name])
    return journal


def finish(*, boot=False):
    journal = locate_journal()
    if journal is None:
        return
    record = read(journal / 'state.json')
    success = committed(record)
    if not success:
        journal = recover_application(journal, record)
        restore_components(journal, record)
        runtime_log = RUNTIME / 'ota-worker.log'
        try:
            if runtime_log.is_file():
                body = runtime_log.read_bytes()[-1024 * 1024:]
                body = body.decode('utf-8', errors='replace').encode('utf-8')[-1024 * 1024:]
                if body:
                    atomic(DATA / 'ota/last-failure.log', body)
        except OSError:
            # Diagnostic storage must never prevent recovery or its cleanup.
            pass
        result = {'schema_version': 1, 'last_result': {
            'status': 'failed', 'target_version': record['version'],
            'installed_at': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
            'packaged_at': '', 'release_notes': '', 'error': 'OTA 未完成，组件与应用已恢复'}}
        # Preserve a more detailed failure result produced by the worker.
        try:
            existing = read(DATA / 'ota/state.json')['last_result']
        except (OSError, ValueError, KeyError):
            existing = {}
        if existing.get('status') != 'failed' or existing.get('target_version') != record['version']:
            atomic(DATA / 'ota/state.json', json.dumps(result).encode())
        status_path = RUNTIME / 'ota-status.json'
        try:
            status = read(status_path)
        except (OSError, ValueError):
            status = {}
        if status.get('job', {}).get('id') == record['job_id'] and status.get('job', {}).get('phase') != 'failed':
            status['active'] = False
            status['job'].update(phase='failed', percent=0, message='更新中断，已恢复原版本',
                                 error='组件事务已恢复', finished_at=result['last_result']['installed_at'])
            atomic(status_path, json.dumps(status).encode())
        run(['systemctl', 'daemon-reload'])
        if not boot:
            run(['systemctl', 'start', '--no-block', SERVICE])
    job = record['job_id']
    # Finish cleanup even if the old worker died after persisting its success.
    release = RELEASES / f"{record['version']}-{job}"
    candidate = DATA.with_name(f'matrix-screen-controller.ota.{job}')
    backup = DATA.with_name(f'matrix-screen-controller.rollback.{job}')
    if success:
        old_release = RELEASES / f"{record['old_version']}-pre-ota-{job}"
        if record['previous_target']:
            previous = Path(record['previous_target'])
            previous = previous if previous.is_absolute() else TARGET.parent / previous
            if previous.parent == RELEASES and previous != release:
                old_release = previous
        if old_release.exists() and old_release != TARGET.resolve():
            shutil.rmtree(old_release)
        if backup.exists():
            shutil.rmtree(backup)
    elif release.exists() and release != TARGET.resolve():
        shutil.rmtree(release)
    if candidate.exists():
        shutil.rmtree(candidate)
    work = WORK_ROOT / f'work.{job}'
    if work.exists():
        shutil.rmtree(work)
    request_path = RUNTIME / 'ota-request.json'
    try:
        request = read(request_path)
    except (OSError, ValueError):
        request = {}
    if request.get('job_id') == job:
        package = Path(request.get('package_path', ''))
        if package.parent == WORK_ROOT / 'uploads' and package.name == job + '.ota':
            package.unlink(missing_ok=True)
        request_path.unlink(missing_ok=True)
    release_runtime_guard()
    for root in (DATA, DATA.with_name(f'matrix-screen-controller.ota.{job}'),
                 DATA.with_name(f'matrix-screen-controller.rollback.{job}')):
        if (root / JOURNAL).exists():
            shutil.rmtree(root / JOURNAL)
    # Installer payloads are transaction inputs, not persistent application data.
    if success and (release / 'system-dependencies').is_dir():
        shutil.rmtree(release / 'system-dependencies')
    run(['systemctl', 'disable', RECOVERY_UNIT.name])
    RECOVERY_UNIT.unlink(missing_ok=True)
    run(['systemctl', 'daemon-reload'])
    HELPER.unlink(missing_ok=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('mode', choices=['watch', 'recover'])
    args = parser.parse_args()
    if os.geteuid() != 0:
        raise RuntimeError('root required')
    if args.mode == 'watch':
        while True:
            state = run(['systemctl', 'show', WORKER, '--property=ActiveState', '--value']).stdout.strip()
            if state not in (b'active', b'activating', b'deactivating'):
                break
            time.sleep(0.5)
    finish(boot=args.mode == 'recover')


if __name__ == '__main__':
    main()