#!/usr/bin/env python3
"""Prepare one published release record from already pushed immutable images."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import subprocess
import tempfile
from datetime import datetime, timezone

IMAGES = {
    'server-main': 'registry.cn-shanghai.aliyuncs.com/comfy-ai/comfy-server',
    'ws-gateway': 'registry.cn-shanghai.aliyuncs.com/easyaigc/wsgateway',
    'web': 'registry.cn-shanghai.aliyuncs.com/comfy-ai/one-ai',
}
POLICY = ('managedUpdateAllowed', 'scheduledUpdateAllowed', 'rollbackSafe',
          'minimumUpdaterVersion', 'minimumDeploymentSchemaVersion', 'requiredFreeBytes')


def inspect_json(reference, template, run=subprocess.check_output):
    args = ['docker', 'buildx', 'imagetools', 'inspect', reference]
    args += ['--format', template] if template else ['--raw']
    return json.loads(run(args, text=True, timeout=90))


def resolve_component(key, reference, version=None):
    image = IMAGES[key]
    if not reference.startswith(image + ':') and not reference.startswith(image + '@sha256:'):
        raise ValueError('Unexpected image for ' + key)
    manifest = inspect_json(reference, None)
    # Pin a platform manifest; refuse mixed platform indexes rather than guessing a digest.
    if manifest.get('manifests'):
        matching = [m for m in manifest['manifests'] if m.get('platform', {}).get('os') == 'linux'
                    and m.get('platform', {}).get('architecture') == 'amd64']
        if len(matching) != 1:
            raise ValueError('Missing unique linux/amd64 image for ' + key)
        digest = matching[0]['digest']
    else:
        # The formatted descriptor is resolved by buildx, not a local cached image ID.
        descriptor = inspect_json(reference, '{{json .Manifest}}')
        digest = descriptor['digest']
    pinned = image + '@' + digest
    metadata = inspect_json(pinned, '{{json .Image}}')
    if metadata.get('os') != 'linux' or metadata.get('architecture') != 'amd64':
        raise ValueError('Invalid platform for ' + key)
    labels = metadata.get('config', {}).get('Labels', {})
    if version and labels.get('org.opencontainers.image.version') != version:
        raise ValueError('Published image version differs from release: ' + key)
    resolved_manifest = inspect_json(pinned, None)
    return {'key': key, 'image': image, 'digest': digest,
            'buildId': labels.get('org.opencontainers.image.revision', 'retained-formal-image'),
            'imageSizeBytes': sum(layer['size'] for layer in resolved_manifest['layers'])}


def prepare(version, policy, components, now=None):
    if not re.fullmatch(r'\d+\.\d+\.\d+\.\d{8}(?:\.\d+)?', version):
        raise ValueError('Invalid product version')
    if {c['key'] for c in components} != set(IMAGES) or len(components) != 3:
        raise ValueError('A release requires all three components')
    for key in POLICY[:3]:
        if type(policy.get(key)) is not bool:
            raise ValueError('Explicit update policy required: ' + key)
    identity = hashlib.sha256(json.dumps(components, sort_keys=True).encode()).hexdigest()[:12]
    return {'schemaVersion': 1, 'channel': 'stable', 'releaseId': version + '-' + identity,
            'productVersion': version, 'platform': 'linux/amd64',
            'publishedAt': now or datetime.now(timezone.utc).isoformat(),
            **{key: policy[key] for key in POLICY if key in policy}, 'components': components}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--version', required=True)
    parser.add_argument('--policy', type=Path, required=True,
                        help='Reviewed release policy; never infer rollback safety from image labels')
    parser.add_argument('--gateway-ref', help='Retain the last released gateway digest when unchanged')
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    policy = json.loads(args.policy.read_text())
    gateway = args.gateway_ref
    if not gateway:
        retained = next(c for c in policy['components'] if c['key'] == 'ws-gateway')
        gateway = retained['image'] + '@' + retained['digest']
    components = [resolve_component('server-main', IMAGES['server-main'] + ':' + args.version, args.version),
                  resolve_component('ws-gateway', gateway),
                  resolve_component('web', IMAGES['web'] + ':' + args.version, args.version)]
    release = prepare(args.version, policy, components)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    # No partial or one-component release becomes visible if inspection fails.
    handle, name = tempfile.mkstemp(dir=args.output.parent, prefix='.release-')
    try:
        with os.fdopen(handle, 'w') as stream:
            json.dump(release, stream, indent=2)
            stream.write('\n')
        os.replace(name, args.output)
    finally:
        if os.path.exists(name):
            os.unlink(name)
    print(json.dumps({'productVersion': release['productVersion'], 'releaseId': release['releaseId'],
                      'output': str(args.output), 'published': False}))


if __name__ == '__main__':
    main()