#!/usr/bin/env python3
"""Conditional startup-guard rollout, with real installed-scanner receive probes."""
import argparse
import datetime
import hashlib
import importlib.util
import json
from pathlib import Path
import time

HERE = Path(__file__).parent
s = importlib.util.spec_from_file_location('template', HERE / 'activate_forge_template.py')
t = importlib.util.module_from_spec(s)
s.loader.exec_module(t)


def native_probe(pod, suffix):
    base = '/tmp/custodian-installed-custody-' + suffix
    created = False
    try:
        t.native.execute(pod, 'sh', '-c', 'umask 077; test ! -e "$1" && mkdir "$1" && printf "%s" custody-probe > "$1/.owner"', 'sh', base)
        created = True
        raw = t.native.execute(pod, 'sh', '-s', '--', base, pod['existing_dispatcher'],
            payload=(HERE / 'probe_forge_installed.sh').read_bytes(), timeout=90)
        return json.loads(raw)
    finally:
        if created:
            t.native.execute(pod, 'sh', '-c', 'test "$(cat "$1/.owner")" = custody-probe && rm -rf -- "$1" && test ! -e "$1"', 'sh', base)


def patch_for(d, candidate, rollback=False):
    current = d['spec']['template']['spec']['initContainers']
    matches = [i for i,c in enumerate(current) if c['name'] == candidate['name']]
    patch = [{'op':'test', 'path':'/metadata/uid', 'value':d['metadata']['uid']},
             {'op':'test', 'path':'/metadata/resourceVersion', 'value':d['metadata']['resourceVersion']}]
    if rollback:
        if matches != [0] or current[0] != candidate:
            raise ValueError('rollback-candidate-differs')
        patch.append({'op':'remove', 'path':'/spec/template/spec/initContainers/0'})
    else:
        if matches:
            raise ValueError('candidate-already-present')
        # apply_locally's JSON array add support appends only; replace exact guarded list.
        patch.append({'op':'replace', 'path':'/spec/template/spec/initContainers', 'value':[candidate,*current]})
    return patch


def verify(pod, suffix):
    t.check_artifacts(pod)
    stats = []
    for item in pod['artifacts'][:2]:
        stat = t.native.execute(pod, 'stat', '-c', '%a %u %g', pod['version_directory']+'/'+item['name']).decode().strip()
        if stat != '755 1000 1000':
            raise ValueError('custody-mode-wrong')
        stats.append({'name':item['name'], 'stat':stat})
    return {'artifacts':stats, 'receive_probe':native_probe(pod, suffix)}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--proposal',type=Path,required=True)
    parser.add_argument('--receipt',type=Path,required=True)
    group=parser.add_mutually_exclusive_group()
    group.add_argument('--apply',action='store_true')
    group.add_argument('--rollback',action='store_true')
    group.add_argument('--restart-proof',action='store_true')
    args=parser.parse_args()
    p=json.loads(args.proposal.read_text())
    for file,expected in p['executor_hashes'].items():
        if hashlib.sha256((HERE/file).read_bytes()).hexdigest()!=expected:
            raise ValueError('executor-changed')
    original=json.loads((HERE/'receive-hook-proposal.json').read_text())
    c=json.loads((HERE/'forge-custody-init.json').read_text())
    d=t.deployment()
    expected=p['patched_template_digest'] if args.rollback or args.restart_proof else p['original_template_digest']
    if d['metadata']['uid']!=p['deployment_uid'] or t.digest(d['spec']['template'])!=expected:
        raise ValueError('deployment-changed')
    pod=t.target_pod(d,original)
    if args.restart_proof:
        patch=[{'op':'test','path':'/metadata/resourceVersion','value':d['metadata']['resourceVersion']},
               {'op':'add','path':'/spec/template/metadata/annotations/custodian.coulomb.social~1custody-proof','value':p['restart_marker']}]
    else:
        patch=patch_for(d,c,args.rollback)
    local=t.apply_locally(d,patch)
    dry=t.patch_remote(patch,dry_run=True)
    if t.digest(local['spec']['template'])!=t.digest(dry['spec']['template']):
        raise ValueError('dry-run-differs')
    receipt={'proposal_id':p['proposal_id'],'revision':p['revision'],
             'proposal_sha256':hashlib.sha256(args.proposal.read_bytes()).hexdigest(),
             'at':datetime.datetime.now(datetime.timezone.utc).isoformat(),
             'mode':'rollback' if args.rollback else 'restart-proof' if args.restart_proof else 'apply' if args.apply else 'dry-run',
             'dry_run_verified':True,'before_pod_uid':pod['pod_uid'],'result':'prepared'}
    args.receipt.write_text(json.dumps(receipt,indent=2)+'\n')
    if not (args.apply or args.rollback or args.restart_proof):
        receipt['before_probe']=verify(pod,'before')
        receipt['result']='dry-run-verified'
    else:
        applied=t.patch_remote(patch)
        receipt.update(result='rollout-started', generation=applied['metadata']['generation'])
        args.receipt.write_text(json.dumps(receipt,indent=2)+'\n')
        after,new=t.wait_ready(original,t.TEMPLATE)
        if new['pod_uid']==pod['pod_uid'] or t.digest(after['spec']['template'])!=t.digest(local['spec']['template']):
            raise ValueError('replacement-not-exact')
        receipt['after_pod_uid']=new['pod_uid']
        if args.rollback:
            # Removing the init guard permits fsGroup to widen modes again; repair explicitly.
            t.native.execute(new,'sh','-s','--','apply',payload=(HERE/'forge-custody-guard.sh').read_bytes())
        else:
            raw=json.loads(t.native.command(['-n','forgejo','get','pod',new['pod'],'-o','json']))
            status=next(x for x in raw['status']['initContainerStatuses'] if x['name']==c['name'])
            if status['state']['terminated']['exitCode']!=0:
                raise ValueError('guard-not-completed')
            receipt['guard_completed']=True
        receipt['after_probe']=verify(new,'after-'+str(after['metadata']['generation']))
        receipt['result']='verified'
    args.receipt.write_text(json.dumps(receipt,indent=2)+'\n')
    print(json.dumps(receipt))


if __name__=='__main__':
    main()