"""Release identity checks and promotion of the exact digests that passed E2E.

Uses only the standard library and Docker Buildx. No registry mutation happens
until every image in the set has been resolved and its build labels validated.
"""
import argparse
import json
import os
from pathlib import Path
import re
import subprocess
import time
from urllib.request import Request, urlopen

COMPONENTS = ('BACKEND', 'FRONTEND', 'PIPELINE')


def docker(*args):
    return subprocess.check_output(['docker', 'buildx', 'imagetools', *args], text=True).strip()


def check_sha(sha):
    if not re.fullmatch(r'[0-9a-f]{40}', sha):
        raise ValueError('Expected a full commit SHA')
    return sha


def check_digest(digest):
    if not re.fullmatch(r'sha256:[0-9a-f]{64}', digest):
        raise ValueError('Expected an immutable image digest')
    return digest


def image_identity(ref):
    image = json.loads(docker('inspect', ref, '--format', '{{json .Image}}'))
    configs = [image] if 'config' in image else list(image.values())
    identities = set()
    for config in configs:
        labels = config['config']['Labels']
        identities.add((labels['io.schoolcompare.commit'], labels['io.schoolcompare.build-id']))
    if len(identities) != 1:
        raise ValueError('Image platforms disagree about their release identity')
    return next(iter(identities))


def resolve_images(sha, verified=False, expected_build_id=None):
    check_sha(sha)
    refs = []
    build_ids = set()
    for component in COMPONENTS:
        image = f"{os.environ['REGISTRY']}/{os.environ[component + '_IMAGE_NAME']}"
        if verified:
            manifest = json.loads(docker('inspect', f'{image}:verified-{sha}', '--format', '{{json .Manifest}}'))
            digest = check_digest(manifest['digest'])
        else:
            digest = check_digest(os.environ[component + '_DIGEST'])
        ref = f'{image}@{digest}'
        actual_sha, build_id = image_identity(ref)
        if actual_sha != sha or not re.fullmatch(r'[0-9a-f]{32}', build_id):
            raise ValueError(f'Unrecognised release identity for {component}')
        if expected_build_id is not None and build_id != expected_build_id:
            raise ValueError(f'Build identity mismatch for {component}')
        build_ids.add(build_id)
        refs.append((image, ref))
    if len(build_ids) != 1:
        raise ValueError('Refusing a mixed image set')
    return {'sha': sha, 'build_id': build_ids.pop(), 'images': refs}


def verify(sha, build_id):
    release = resolve_images(sha, expected_build_id=build_id)
    for image, ref in release['images']:
        docker('create', '-t', f'{image}:verified-{sha}', ref)
    return release


def promote(sha):
    release = resolve_images(sha, verified=True)
    # Resolve all targets first; never discover a missing candidate halfway through.
    for image, _ in release['images']:
        try:
            docker('create', '-t', f'{image}:prod-previous', f'{image}:prod')
        except subprocess.CalledProcessError:
            print(f'No rollback pointer saved for {image}', flush=True)
    for image, ref in release['images']:
        docker('create', '-t', f'{image}:prod', ref)
    return release


def matches(payload, sha, build_id):
    return all(payload.get(component) == {'sha': sha, 'build_id': build_id}
               for component in ('frontend', 'backend'))


def wait(base_url, sha, build_id, timeout):
    check_sha(sha)
    if not re.fullmatch(r'[0-9a-f]{32}', build_id):
        raise ValueError('Missing expected build identity')
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            req = Request(f'{base_url.rstrip("/")}/release.json?check={time.time_ns()}',
                          headers={'Cache-Control': 'no-cache'})
            with urlopen(req, timeout=min(10, max(.1, deadline - time.monotonic()))) as response:
                payload = json.load(response)
            if matches(payload, sha, build_id):
                print(f'Verified deployed release {sha} / {build_id}')
                return
        except (OSError, ValueError):
            pass
        time.sleep(min(5, max(0, deadline - time.monotonic())))
    raise RuntimeError('Deployment did not report the expected frontend/backend release')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=['wait', 'verify', 'promote'])
    parser.add_argument('--timeout', type=float, default=300)
    parser.add_argument('--release', type=Path)
    parser.add_argument('--output', type=Path)
    args = parser.parse_args()
    identity = json.loads(args.release.read_text()) if args.release else {
        'sha': os.environ.get('EXPECTED_SHA', ''),
        'build_id': os.environ.get('EXPECTED_BUILD_ID', ''),
    }
    if args.action == 'wait':
        wait(os.environ['BASE_URL'], identity['sha'], identity['build_id'], args.timeout)
        return
    result = (verify(identity['sha'], identity['build_id']) if args.action == 'verify'
              else promote(identity['sha']))
    if args.output:
        args.output.write_text(json.dumps(result))


if __name__ == '__main__':
    main()