"""Verify the exact archive and OCI/config identity before accepting imported images."""
import hashlib
import io
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tarfile

STAGE=Path('/storage/contextra/deploy/ai2-20261006-r01')
SOURCE=Path('/workspace/contextra/source/4f0bd4e')
SITE=Path('/storage/contextra/site/ai2-20261006-r01')
sys.path.insert(0,str(SOURCE))
from deployment_kit.image_identity import resolve

def main():
    os.umask(0o077)
    assert subprocess.check_output(['hostname']).decode().strip()=='icertirnd-ai'
    assert shutil.disk_usage('/').free>=45*1024**3
    manifest=json.loads((STAGE/'assets/cached-manifest.json').read_text())
    archive=STAGE/'assets'/manifest['archive']['file']
    with archive.open('rb') as f:
        assert hashlib.file_digest(f,'sha256').hexdigest()==manifest['archive']['sha256']
    done=STAGE/'image-load-complete.json'
    if not done.exists():
        assert not (STAGE/'image-load-intent.json').exists(), 'Inspect interrupted load before resuming'
        (STAGE/'image-load-intent.json').write_text(json.dumps({'sha256':manifest['archive']['sha256']}))
        before=[json.loads(line) for line in subprocess.check_output(['docker','image','ls','--format','json']).decode().splitlines()]
        old_tags={x['Repository']+':'+x['Tag']:subprocess.check_output(['docker','image','inspect',x['Repository']+':'+x['Tag'],'--format','{{.Id}}']).decode().strip()
                  for x in before if x['Repository']!='<none>' and x['Tag']!='<none>'}
        mapping={value:'contextra-ai2-base-'+role+':20261006-r01' for role,value in manifest['images'].items()}
        assert not set(mapping.values()) & set(old_tags)
        with tarfile.open(archive,'r:') as src:
            index=json.load(src.extractfile('index.json'))
            classic=json.load(src.extractfile('manifest.json'))
            for row in classic:
                row['RepoTags']=[mapping['sha256:'+row['Config'].split('/')[-1]]]
            for row in index['manifests']:
                image=json.load(src.extractfile('blobs/sha256/'+row['digest'].split(':')[1]))
                tag=mapping[image['config']['digest']]
                row['annotations']={'io.containerd.image.name':'docker.io/library/'+tag,
                                    'org.opencontainers.image.ref.name':tag.split(':')[1]}
            replacements={name:json.dumps(value,separators=(',',':')).encode() for name,value in
                          [('manifest.json',classic),('index.json',index)]}
            with (STAGE/'image-load.log').open('x') as log:
                proc=subprocess.Popen(['timeout','--kill-after=15s','1200s','docker','image','load'],stdin=subprocess.PIPE,stdout=log,stderr=log)
                try:
                    with tarfile.open(fileobj=proc.stdin,mode='w|') as dest:
                        for item in src:
                            if item.name in replacements:
                                raw=replacements[item.name];item.size=len(raw);dest.addfile(item,io.BytesIO(raw))
                            else:dest.addfile(item,src.extractfile(item) if item.isfile() else None)
                    proc.stdin.close()
                    assert proc.wait(timeout=1200)==0, 'ARCHIVE_IMPORT_FAILED'
                finally:
                    if proc.poll() is None:
                        proc.terminate()
                        try:proc.wait(timeout=15)
                        except subprocess.TimeoutExpired:proc.kill();proc.wait(timeout=10)
        for tag,expected in old_tags.items():
            assert subprocess.check_output(['docker','image','inspect',tag,'--format','{{.Id}}']).decode().strip()==expected,'UNRELATED_TAG_CHANGED'
        done.write_text(json.dumps({'sha256':manifest['archive']['sha256'],'loaded':True}))
    resolved={}
    for role,expected in manifest['images'].items():
        actual=resolve('contextra-ai2-base-'+role+':20261006-r01',expected,search=False)
        assert actual, role+' IMAGE_NOT_FOUND'
        resolved[role]=actual
        tag='contextra-ai2-base-'+role+':20261006-r01'
        prior=subprocess.run(['docker','image','inspect',tag],capture_output=True,text=True,timeout=10)
        if prior.returncode==0:
            assert json.loads(prior.stdout)[0]['Id']==actual
        else:
            subprocess.run(['docker','image','tag',actual,tag],check=True,timeout=15)
    target=STAGE/'images-resolved.json'
    if target.exists():assert json.loads(target.read_text())==resolved
    else:target.write_text(json.dumps(resolved,indent=2)+'\n')
    env=SITE/'runtime.env';raw=env.read_text()
    for role,key in [('analytics','CONTEXTRA_IMAGE'),('analytics','CONTEXTRA_SECURE_ANALYTICS_IMAGE'),
                     ('ai-control-sidecar','CONTEXTRA_SIDECAR_IMAGE'),('ai-control-sidecar','CONTEXTRA_SIDECAR_BASE_IMAGE'),
                     ('ai-gateway','CONTEXTRA_GATEWAY_IMAGE')]:
        before=key+'='+manifest['target_test_images'][role]
        after=key+'='+resolved[role]
        assert before in raw or after in raw
        raw=raw.replace(before,after)
    env.write_text(raw)
    overlay=SITE/'site.compose.json';doc=json.loads(overlay.read_text())
    for role in ('postgresql','redis'):
        assert doc['services'][role]['image'] in (manifest['images'][role],resolved[role])
        doc['services'][role]['image']=resolved[role]
    overlay.write_text(json.dumps(doc,indent=2)+'\n')
    print(json.dumps({'status':'BASE_IMAGES_VERIFIED_CURRENT_SOURCE_REBUILD_PENDING','resolved':resolved}))

if __name__=='__main__':main()