#!/usr/bin/env python3
"""Update installed CUPS project files without reinstalling or changing config."""

import argparse
import ast
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path, PurePosixPath
import shutil
import signal
import stat
import subprocess
import sys
import tempfile
import time
import urllib.parse
import urllib.request
import uuid
import zipfile


ARCHIVE_URL = 'https://git.workyai.cn/237899745/S905L3A/archive/master.zip'
MAX_DOWNLOAD = 20 * 1024 * 1024
MAX_EXTRACTED = 100 * 1024 * 1024
TIMERS = ('cups-network-watchdog.timer', 'cups-print-watchdog.timer')
WATCHDOGS = ('cups-network-watchdog.service', 'cups-print-watchdog.service')
MANAGER = 'cups-driver-manager.service'
ACTIVE = {'active', 'activating', 'reloading', 'deactivating'}
PROGRAMS = {
    'cups-driver-manager/driver_manager.py': 'opt/cups-driver-manager/driver_manager.py',
    'watchdog/network-watchdog.sh': 'opt/cups-watchdog/network-watchdog.sh',
    'watchdog/print-watchdog.sh': 'opt/cups-watchdog/print-watchdog.sh',
}


class UpdateError(Exception):
    pass


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


def safe_path(root, relative):
    relative = PurePosixPath(relative)
    if relative.is_absolute() or '..' in relative.parts or not relative.parts:
        raise UpdateError(f'Unsafe relative path: {relative}')
    root = root.resolve()
    candidate = root
    for part in relative.parts:
        candidate = candidate / part
        if candidate.is_symlink():
            raise UpdateError(f'Refusing symbolic link: {candidate}')
    if not candidate.resolve().is_relative_to(root):
        raise UpdateError(f'Path escaped update root: {candidate}')
    return candidate


def fetch_archive(url, output):
    if urllib.parse.urlsplit(url).scheme != 'https':
        raise UpdateError('The update archive URL must use HTTPS')
    started = time.monotonic()
    size = 0
    with urllib.request.urlopen(url, timeout=30) as response:
        if urllib.parse.urlsplit(response.geturl()).scheme != 'https':
            raise UpdateError('Refusing an HTTPS-to-HTTP redirect')
        with output.open('wb') as destination:
            while True:
                chunk = response.read(65536)
                if not chunk:
                    break
                size += len(chunk)
                if size > MAX_DOWNLOAD or time.monotonic() - started > 180:
                    raise UpdateError('Update download exceeded the size/time limit')
                destination.write(chunk)
    if not size:
        raise UpdateError('Downloaded archive is empty')


def extract_archive(archive_path, destination):
    with zipfile.ZipFile(archive_path) as archive:
        entries = archive.infolist()
        if len(entries) > 5000 or sum(entry.file_size for entry in entries) > MAX_EXTRACTED:
            raise UpdateError('Update archive is too large')
        seen = set()
        for entry in entries:
            if '\\' in entry.filename or ':' in entry.filename:
                raise UpdateError(f'Invalid archive path: {entry.filename}')
            target = safe_path(destination, entry.filename)
            if target in seen:
                raise UpdateError(f'Duplicate archive path: {entry.filename}')
            seen.add(target)
            mode = entry.external_attr >> 16
            if stat.S_IFMT(mode) not in (0, stat.S_IFREG, stat.S_IFDIR):
                raise UpdateError(f'Unsupported archive entry: {entry.filename}')
        archive.extractall(destination)
    candidates = [destination] + [path for path in destination.iterdir() if path.is_dir()]
    candidates = [path for path in candidates if (path / 'setup_cups.sh').is_file()]
    if len(candidates) != 1:
        raise UpdateError('Archive must contain exactly one project root')
    return candidates[0]


@dataclass
class Change:
    relative: str
    content: bytes
    previous: bytes
    metadata: object


def build_plan(source, root, bash='bash'):
    # Validate the release before inspecting or changing any installed file.
    for relative in ('setup_cups.sh', *PROGRAMS):
        path = safe_path(source, relative)
        if not path.is_file():
            raise UpdateError(f'Incomplete release: {relative}')
        if relative.endswith('.sh'):
            result = subprocess.run([bash, '-n', str(path)], capture_output=True, timeout=30)
            if result.returncode:
                raise UpdateError(f'Invalid shell syntax: {relative}')
        else:
            ast.parse(path.read_bytes(), filename=relative)

    if not any(safe_path(root, relative).is_file() for relative in PROGRAMS.values()) and not safe_path(root, 'usr/share/cups/templates-zh_CN').is_dir():
        raise UpdateError('No existing installation of this CUPS project was found; use the installer first')

    mapping = dict(PROGRAMS)
    groups = (
        ('cups-templates-zh_CN', 'usr/share/cups/templates', ('.tmpl',)),
        ('cups-templates-zh_CN', 'usr/share/cups/templates-zh_CN', ('.tmpl',)),
        ('cups-docroot-zh_CN', 'usr/share/cups/doc-root', ('.html', '.js')),
    )
    # Each source may be installed at more than one destination.
    pairs = list(mapping.items())
    for source_dir, installed_dir, suffixes in groups:
        directory = safe_path(source, source_dir)
        if not directory.is_dir():
            raise UpdateError(f'Incomplete release: {source_dir}')
        installed = safe_path(root, installed_dir)
        if installed.is_dir():
            for path in sorted(directory.iterdir()):
                if path.suffix in suffixes:
                    safe_path(source, f'{source_dir}/{path.name}')
                    pairs.append((f'{source_dir}/{path.name}', f'{installed_dir}/{path.name}'))

    css = safe_path(root, 'usr/share/cups/doc-root/cups.css')
    modern_css = css.is_file() and b'.cups-header' in css.read_bytes()
    changes = []
    for origin, relative in pairs:
        target = safe_path(root, relative)
        if not target.exists():
            continue
        if not target.is_file():
            raise UpdateError(f'Installed target is not a regular file: {target}')
        content = safe_path(source, origin).read_bytes()
        if modern_css and origin in ('cups-templates-zh_CN/header.tmpl', 'cups-templates-zh_CN/trailer.tmpl'):
            for old, new in ((b'header', b'cups-header'), (b'body', b'cups-body'), (b'footer', b'cups-footer')):
                content = content.replace(b'class="' + old + b'"', b'class="' + new + b'"')
        previous = target.read_bytes()
        if previous != content:
            changes.append(Change(relative, content, previous, target.stat()))
    return changes


class Services:
    def __init__(self, runner=subprocess.run, sleeper=time.sleep, proc=Path('/proc')):
        self.runner = runner
        self.sleeper = sleeper
        self.proc = proc
        self.states = {}
        self.stopped = []

    def command(self, *args):
        result = self.runner(['systemctl', *args], capture_output=True, text=True, timeout=120)
        if result.returncode:
            raise UpdateError(f'systemctl {" ".join(args)} failed: {result.stderr.strip()}')
        return result.stdout

    def state(self, unit):
        result = self.runner(['systemctl', 'show', unit, '--property=LoadState,ActiveState,MainPID,UnitFileState'],
                             capture_output=True, text=True, timeout=120)
        state = dict(line.split('=', 1) for line in result.stdout.splitlines() if '=' in line)
        if state.get('LoadState') == 'not-found':
            return state
        if result.returncode or 'ActiveState' not in state:
            raise UpdateError(f'Cannot inspect {unit}: {result.stderr.strip()}')
        return state

    def snapshot(self, manager_changed):
        units = list(TIMERS) + list(WATCHDOGS)
        if manager_changed:
            units.append(MANAGER)
        for unit in units:
            state = self.state(unit)
            if state.get('LoadState') == 'not-found':
                continue
            if state.get('LoadState') not in ('loaded', 'masked'):
                raise UpdateError(f'Cannot inspect unit: {unit}')
            if state.get('ActiveState') in ACTIVE and state.get('UnitFileState', '').startswith('masked'):
                raise UpdateError(f'Cannot safely restore a masked active unit: {unit}')
            self.states[unit] = state

    def busy_processes(self, manager_pid=0):
        # Read only process names/parents, never environment variables or credentials.
        parents = {}
        watchdog_pids = []
        for directory in self.proc.glob('[0-9]*'):
            try:
                pid = int(directory.name)
                fields = dict(line.split(':', 1) for line in (directory / 'status').read_text().splitlines() if ':' in line)
                parents[pid] = int(fields.get('PPid', '0').strip())
                args = (directory / 'cmdline').read_bytes().split(b'\0')
                if any(Path(os.fsdecode(arg)).name in ('network-watchdog.sh', 'print-watchdog.sh') for arg in args if arg):
                    watchdog_pids.append(pid)
            except (OSError, ValueError):
                continue
        descendants = {manager_pid} if manager_pid else set()
        while True:
            children = {pid for pid, parent in parents.items() if parent in descendants}
            expanded = descendants | children
            if expanded == descendants:
                break
            descendants = expanded
        return watchdog_pids, descendants - {manager_pid}

    def pause(self):
        for timer in TIMERS:
            if self.states.get(timer, {}).get('ActiveState') in ACTIVE:
                self.stopped.append(timer)
                self.command('stop', timer)
        # Never kill a watchdog halfway through changing network configuration.
        for unit in WATCHDOGS:
            if unit not in self.states:
                continue
            for _ in range(60):
                if self.state(unit).get('ActiveState') not in ACTIVE:
                    break
                self.sleeper(2)
            else:
                raise UpdateError(f'Watchdog is still running; update cancelled: {unit}')
        manager = self.states.get(MANAGER, {})
        watchdogs, children = self.busy_processes(int(manager.get('MainPID', '0')))
        if watchdogs or children:
            raise UpdateError('A manual watchdog or driver installation is running; retry when idle')
        if manager.get('ActiveState') in ACTIVE:
            self.stopped.append(MANAGER)
            self.command('stop', MANAGER)

    def restore(self):
        # Keep watchdog timers paused until the manager has passed startup checks.
        order = ([MANAGER] if MANAGER in self.stopped else []) + [unit for unit in TIMERS if unit in self.stopped]
        errors = []
        for unit in order:
            try:
                self.command('start', unit)
                for _ in range(3):
                    self.sleeper(1)
                    if self.state(unit).get('ActiveState') != 'active':
                        raise UpdateError(f'Startup health check failed: {unit}')
            except Exception as error:
                errors.append(str(error))
                if unit == MANAGER:
                    break
        if errors:
            raise UpdateError('; '.join(errors))

    def stop_for_rollback(self):
        errors = []
        for unit in self.stopped:
            try:
                self.command('stop', unit)
            except Exception as error:
                errors.append(str(error))
        if errors:
            raise UpdateError('; '.join(errors))


def atomic_write(target, data, metadata):
    fd, temporary = tempfile.mkstemp(prefix=f'.{target.name}.update-', dir=target.parent)
    temporary = Path(temporary)
    try:
        with os.fdopen(fd, 'wb') as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())
        shutil.copystat(target, temporary)
        os.chmod(temporary, stat.S_IMODE(metadata.st_mode))
        if os.name == 'posix':
            os.chown(temporary, metadata.st_uid, metadata.st_gid)
        os.replace(temporary, target)
    finally:
        if temporary.exists():
            temporary.unlink()


def clear_old_locks(root):
    for kind in ('network', 'print'):
        lock = safe_path(root, f'run/cups-watchdog/{kind}.lock')
        if lock.exists():
            if not lock.is_dir() or any(lock.iterdir()):
                raise UpdateError(f'Refusing to remove an unexpected watchdog lock: {lock}')
            lock.rmdir()


def apply_update(changes, root, services, release):
    backup_root = safe_path(root, 'var/backups/cups-project')
    backup_root.mkdir(parents=True, exist_ok=True, mode=0o700)
    backup = backup_root / (datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ-') + uuid.uuid4().hex[:8])
    backup.mkdir(mode=0o700)
    manifest = {'release': release, 'status': 'preparing', 'files': []}

    def save_status(status):
        manifest['status'] = status
        (backup / 'manifest.json').write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding='utf-8')

    for change in changes:
        original = safe_path(root, change.relative)
        copy = safe_path(backup / 'files', change.relative)
        copy.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(original, copy)
        if copy.read_bytes() != change.previous:
            raise UpdateError(f'File changed during backup; retry: {original}')
        manifest['files'].append({'path': change.relative, 'before': digest(change.previous), 'after': digest(change.content)})
    save_status('prepared')
    print(f'Backup: {backup}', flush=True)

    attempted = []
    try:
        services.snapshot(any(change.relative == PROGRAMS['cups-driver-manager/driver_manager.py'] for change in changes))
        manifest['services_before'] = getattr(services, 'states', {})
        save_status('prepared')
        services.pause()
        clear_old_locks(root)
        for change in changes:
            if safe_path(root, change.relative).read_bytes() != change.previous:
                raise UpdateError(f'File changed after preflight: {change.relative}')
        save_status('applying')
        for change in changes:
            attempted.append(change)
            target = safe_path(root, change.relative)
            atomic_write(target, change.content, change.metadata)
            if target.read_bytes() != change.content:
                raise UpdateError(f'Written file verification failed: {change.relative}')
        services.restore()
        save_status('complete')
    except (Exception, KeyboardInterrupt) as error:
        rollback_errors = []
        if attempted:
            try:
                services.stop_for_rollback()
            except Exception as rollback_error:
                rollback_errors.append(str(rollback_error))
            for change in reversed(attempted):
                try:
                    previous = safe_path(backup / 'files', change.relative).read_bytes()
                    if digest(previous) != digest(change.previous):
                        raise UpdateError(f'Backup checksum mismatch: {change.relative}')
                    atomic_write(safe_path(root, change.relative), previous, change.metadata)
                except Exception as rollback_error:
                    rollback_errors.append(str(rollback_error))
        try:
            services.restore()
        except Exception as rollback_error:
            rollback_errors.append(str(rollback_error))
        try:
            save_status('rollback_failed' if rollback_errors else 'rolled_back')
        except Exception as rollback_error:
            rollback_errors.append(str(rollback_error))
        detail = '; '.join(rollback_errors) if rollback_errors else 'Original files/service states restored'
        raise UpdateError(f'Update failed: {error}. {detail}. Backup: {backup}') from error
    return backup


@contextmanager
def update_lock(path):
    import fcntl
    with path.open('a') as handle:
        try:
            fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as error:
            raise UpdateError('Another update is already running') from error
        yield


def confirm():
    try:
        with open('/dev/tty', 'r+') as terminal:
            terminal.write('Apply this update? [y/N]: ')
            terminal.flush()
            return terminal.readline().strip().lower() == 'y'
    except OSError as error:
        raise UpdateError('No interactive terminal. Review with --dry-run, then use --yes to confirm') from error


def check_prerequisites():
    if sys.version_info < (3, 9):
        raise UpdateError('Python 3.9 or later is required')
    if os.name != 'posix' or os.geteuid() != 0:
        raise UpdateError('Run this updater as root on the Linux print server')
    if not shutil.which('systemctl') or not shutil.which('bash'):
        raise UpdateError('systemctl and Bash are required; no packages will be installed automatically')


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--archive-url', default=ARCHIVE_URL, help='HTTPS archive URL (default: project master)')
    parser.add_argument('--source-dir', type=Path, help='Use a trusted local release instead of downloading')
    parser.add_argument('--dry-run', action='store_true', help='Validate and list changes without applying them')
    parser.add_argument('--yes', action='store_true', help='Confirm the displayed file update without prompting')
    args = parser.parse_args(argv)
    check_prerequisites()
    root = Path('/')
    with update_lock(safe_path(root, 'run/cups-project-update.lock')), tempfile.TemporaryDirectory(prefix='cups-project-update-') as temporary:
        if args.source_dir:
            source = args.source_dir.resolve()
            release = {'source': str(source)}
        else:
            archive = Path(temporary) / 'release.zip'
            print('Downloading the current project release...', flush=True)
            fetch_archive(args.archive_url, archive)
            source = extract_archive(archive, Path(temporary) / 'release')
            release = {'source': args.archive_url, 'archive_sha256': digest(archive.read_bytes())}
        changes = build_plan(source, root)
        if not changes:
            print('No installed project files need updating. No services were changed.')
            return 0
        print('Only the following installed files will be replaced:')
        for change in changes:
            print(f'  /{change.relative}')
        print('CUPS queues/config, network settings, credentials, systemd units and packages are preserved.')
        print('The driver-manager web service may be briefly unavailable; CUPS will not be restarted.')
        if args.dry_run or not (args.yes or confirm()):
            print('No installed files or services were changed.')
            return 0
        backup = apply_update(changes, root, Services(), release)
        print(f'Update completed. Backup retained at: {backup}')
    return 0


def interrupted(signum, frame):
    raise UpdateError(f'Update interrupted by signal {signum}')


if __name__ == '__main__':
    signal.signal(signal.SIGTERM, interrupted)
    try:
        sys.exit(main())
    except (Exception, KeyboardInterrupt) as error:
        print(f'ERROR: {error}', file=sys.stderr)
        sys.exit(1)