#!/usr/bin/env python3
"""Resolve Compose images without changing project names or local bind paths."""

import json
import os
import re
import subprocess
import sys
from pathlib import Path


def output(*args, **kwargs):
    return subprocess.check_output(args, text=True, **kwargs).strip()  # noqa: S603


def resolve(reference):
    if '@sha256:' in reference:
        return reference
    descriptor = json.loads(
        output('docker', 'buildx', 'imagetools', 'inspect', reference, '--format', '{{json .Manifest}}')
    )
    digest = descriptor['digest']
    if not re.fullmatch(r'sha256:[0-9a-f]{64}', digest):
        raise ValueError(f'Invalid registry digest for {reference}')
    # Strip tag only from the final path segment (registry ports are preserved).
    repository = reference.rsplit('/', 1)
    repository[-1] = repository[-1].split(':')[0]
    return '/'.join(repository) + '@' + digest


def prepare(source_file):
    config_repo = Path(os.environ['CONFIG_REPO'])
    source_repo = Path(os.environ['REPO'])
    directory = Path(os.environ['RUN_DIR'])
    relative = source_file.relative_to(source_repo)
    project_dir = config_repo / relative.parent
    base = ['docker', 'compose', '--project-directory', str(project_dir), '-f', str(source_file)]
    config = json.loads(output(*base, 'config', '--format', 'json', cwd=config_repo))
    project = config['name']
    previous_file = directory / 'previous.json'
    previous = json.loads(previous_file.read_text()) if previous_file.exists() else {}
    images_file = directory / 'compose-images.json'
    locks = json.loads(images_file.read_text()) if images_file.exists() else previous.get('compose-images', {})
    release = json.loads((directory / 'release.json').read_text())
    before = json.loads(json.dumps(config))
    for service, settings in config['services'].items():
        reference = settings.get('image')
        nextcloud_aio_master = project == 'nextcloud' and service == 'nextcloud-aio-mastercontainer'
        if not reference or settings.get('build'):
            raise ValueError(f'{project}/{service}: Compose deploy requires a published image')
        image_repo = reference.split('@')[0].rsplit('/', 1)
        image_repo[-1] = image_repo[-1].split(':')[0]
        image_repo = '/'.join(image_repo)
        # Nextcloud AIO validates the mastercontainer image reference and rejects
        # a digest. Keep its configured tag so AIO can start and manage its stack.
        if nextcloud_aio_master:
            pinned = reference
        elif image_repo in release['images']:
            pinned = image_repo + '@' + release['images'][image_repo]
        elif os.environ.get('REFRESH_IMAGES') != 'true' and reference in locks:
            pinned = locks[reference]
        else:
            pinned = resolve(reference)
        settings['image'] = pinned
        locks[reference] = pinned
        # Capture what is running, not the current value of its mutable tag.
        ids = output(
            'docker',
            'ps',
            '-aq',
            '--filter',
            f'label=com.docker.compose.project={project}',
            '--filter',
            f'label=com.docker.compose.service={service}',
        ).splitlines()
        actual = set()
        for container in ids:
            image_id = output('docker', 'inspect', container, '--format', '{{.Image}}')
            digests = json.loads(output('docker', 'image', 'inspect', image_id, '--format', '{{json .RepoDigests}}'))
            actual.add(next((d for d in digests or [] if d.split('@')[0] == image_repo), image_id))
        if len(actual) > 1:
            raise ValueError(f'{project}/{service}: mixed running images, cannot capture one recovery config')
        # AIO also rejects a digest in its recovery config. Preserve its tag in
        # both deploy and recovery files.
        if nextcloud_aio_master:
            before['services'][service]['image'] = reference
        else:
            before['services'][service]['image'] = next(iter(actual)) if actual else reference
    for name, data in (('compose', config), ('compose-before', before)):
        folder = directory / name
        folder.mkdir(mode=0o700, exist_ok=True)
        destination = folder / f'{relative.parent.name}.json'
        destination.write_text(json.dumps(data, indent=2) + '\n')
        destination.chmod(0o600)
    images_file.write_text(json.dumps(locks, indent=2) + '\n')
    print(f'Compose {project}: images pinned; local paths preserved')
    print(
        f'Recovery: docker compose --project-directory {project_dir} -p {project} -f {directory}/compose-before/{relative.parent.name}.json up -d --pull never'
    )


if __name__ == '__main__':
    prepare(Path(sys.argv[1]))