#!/usr/bin/env python3
"""Verify the resolved installation images, without pulling layers or using ACR credentials."""
import argparse
from concurrent.futures import ThreadPoolExecutor
import json
import subprocess
import sys
import tempfile


STABLE_KEYS = {'easyai-server': 'SERVER_STABLE_DIGEST', 'easyai-web': 'WEB_STABLE_DIGEST', 'ws-gateway': 'WS_GATEWAY_STABLE_DIGEST'}


def verify(image, docker, config, expected=None):
    # Official installation images must be available to a new anonymous user.
    # Never reuse a maintainer's credential helper as evidence of public access.
    official = image.startswith(('registry.cn-shanghai.aliyuncs.com/comfy-ai/',
                                 'registry.cn-shanghai.aliyuncs.com/easyaigc/'))
    command = docker + (['--config', config] if official else [])
    try:
        result = subprocess.run(command + ['manifest', 'inspect', '--verbose', image],
                                capture_output=True, text=True, timeout=90)
        if result.returncode:
            message = result.stderr.lower()
            reason = ('标签或摘要不存在' if any(s in message for s in ('not found', 'manifest unknown', 'no such manifest'))
                      else '未开放匿名拉取' if any(s in message for s in ('unauthorized', 'denied', 'authentication required'))
                      else '仓库不可访问，请检查网络或仓库状态')
            raise ValueError(reason)
        manifests = json.loads(result.stdout)
        manifests = manifests if isinstance(manifests, list) else [manifests]
        platforms = [m.get('Descriptor', {}).get('platform', {}) for m in manifests]
        if not any(p.get('os') == 'linux' and p.get('architecture') == 'amd64' for p in platforms):
            raise ValueError('没有 linux/amd64 产物')
        if expected and not any(m.get('Descriptor', {}).get('digest') == expected
                                and m.get('Descriptor', {}).get('platform', {}).get('architecture') == 'amd64'
                                for m in manifests):
            raise ValueError('官方稳定别名与发布清单尚未同步；保留当前服务')
        return image, None
    except (subprocess.TimeoutExpired, OSError):
        return image, '镜像检查超时或 Docker 不可用'
    except (ValueError, KeyError, TypeError) as error:
        return image, str(error) if isinstance(error, ValueError) and not isinstance(error, json.JSONDecodeError) else '仓库返回了无效的镜像信息'


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--sudo', action='store_true')
    parser.add_argument('--services', nargs='*', default=[])
    args = parser.parse_args()
    configuration = json.load(sys.stdin)
    images = sorted({value['image'] for service, value in configuration['services'].items()
                     if not args.services or service in args.services})
    expected = {value['image']: configuration.get('x-easyai-release', {}).get(service)
                for service, value in configuration['services'].items() if service in STABLE_KEYS}
    if not images:
        parser.error('没有可检查的部署镜像')
    docker = ['sudo', 'docker'] if args.sudo else ['docker']
    with tempfile.TemporaryDirectory(prefix='easyai-anonymous-registry-') as config:
        with ThreadPoolExecutor(max_workers=4) as executor:
            results = list(executor.map(lambda image: verify(image, docker, config, expected.get(image)), images))
    failures = [(image, error) for image, error in results if error]
    for image, error in failures:
        print(f'❌ {image}: {error}', file=sys.stderr)
    if failures:
        print('镜像检查失败，尚未拉取或重建服务；请核对发布记录和组件实际版本。', file=sys.stderr)
        return 1
    print(f'✅ 已核验 {len(images)} 个部署镜像的可访问性和 linux/amd64 平台')
    return 0


if __name__ == '__main__':
    sys.exit(main())