#!/usr/bin/env python3
"""Recover postprocessing after deletion of RUN; never submit integration."""
import argparse
import json
from pathlib import Path
import re
import subprocess
import rcms as core
import operations as ops
CYCLE='2026082100'
END='2027030100'
JOBS={'wrf':('1401038','1401039'),'regcm':('1401040','1401041')}


def check_job(job,model,phase):
    p=subprocess.run(['scontrol','show','job','-o',job],stdout=subprocess.PIPE,stderr=subprocess.PIPE,universal_newlines=True)
    core.need(p.returncode==0,'Cannot verify job '+job+': '+p.stderr.strip())
    row=dict(re.findall(r'(\w+)=(\S+)',p.stdout))
    core.need(row.get('UserId','').startswith('nampq('),'Unexpected job owner: '+job)
    core.need(row.get('JobName')=='rcms-'+model+'-'+CYCLE+'-'+phase+'-initial','Unexpected job identity: '+job)
    return row


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--execute',action='store_true')
    args=parser.parse_args()
    cfg=json.loads((core.HERE/'config.json').read_text())
    core.need(ops.storage(cfg)==core.HERE/'runs','Install workspace config first (runs_root=runs)')
    rows=ops.queue(cfg)
    for model,(integration,oldpost) in JOBS.items():
        check_job(integration,model,'integrate')
        if oldpost in rows:
            old=check_job(oldpost,model,'post')
            core.need(old['JobState']=='PENDING','Old post already started; inspect before recovery')
        raw=Path(cfg['legacy_dirs'][model])/'wdir'/CYCLE/'output'
        core.need(raw.is_dir(),'Missing raw output: '+str(raw))
        print(model+': keep integration '+integration+'; recover post after it exits; require coverage through '+END)
    print('Control files:',ops.storage(cfg))
    print('Replace only pending post jobs 1401039 and 1401041 after replacement submission succeeds.')
    if not args.execute:
        print('PLAN ONLY. Use --execute as nampq to submit replacement post jobs.')
        return
    ops.write_preflight(cfg)
    core.validate_config(cfg,list(JOBS))
    with ops.cycle_lock(cfg,CYCLE):
        receipt=ops.storage(cfg)/CYCLE/'deleted-run-recovery.json'
        if receipt.exists():
            root,man=ops.managed_run(json.loads(receipt.read_text())['run'])
        else:
            decisions={m:{'action':'ADOPT_POST','raw':str(Path(cfg['legacy_dirs'][m])/'wdir'/CYCLE/'output'),
                          'expected_end':END,'original_integration_job':pair[0]} for m,pair in JOBS.items()}
            root,man=ops.new_run(cfg,CYCLE,list(JOBS),decisions)
            core.put(receipt,{'run':str(root),'original_jobs':JOBS})
        print('Recovery RUN:',root,flush=True)
        for model,(integration,oldpost) in JOBS.items():
            prefix=root/(model+'-post-rescue')
            record=Path(str(prefix)+'.job.json')
            if record.exists():
                print(model,'replacement already recorded:',json.loads(record.read_text())['job_id'])
            else:
                core.need(not Path(str(prefix)+'.submit-command.json').exists(),
                          'Previous submission outcome needs inspection; refusing duplicate: '+str(prefix))
                # Deleted runtime can cause a nonzero exit AFTER all outputs are written.
                # afterany waits for exit; expected_end validation then refuses partial output.
                core.submit_job(root,man,model,'post','rescue',integration,dependency_type='afterany')
            if oldpost in ops.queue(cfg):
                old=check_job(oldpost,model,'post')
                core.need(old['JobState']=='PENDING','Old post no longer pending; inspect manually')
                result=subprocess.run(['scancel',oldpost])
                core.need(result.returncode==0,'Cannot cancel obsolete pending post '+oldpost)
                note=root/(model+'-obsolete-post-cancelled.json')
                if not note.exists():core.put(note,{'job_id':oldpost,'time':core.stamp()})
        print('Integration jobs were not cancelled or resubmitted. Raw files remain in place.')


if __name__=='__main__':main()