#!/usr/bin/env python3
from pathlib import Path
import argparse,json,hashlib,urllib.request,time,datetime,os
ROOT=Path('/data/video-pipeline/HaTramAudio/project/031-Buoi-Dien-Cuoi-Truoc-Ngay-Tot-Nghiep');RUN='run-20260723T010035Z-09393c41'
CFG={'footage':(['http://192.168.1.104:8022','http://192.168.1.104:8025','http://192.168.1.104:8028','http://192.168.1.104:8031'],'/v1/footage-render-huyenan/jobs'),'final':(['http://192.168.1.104:8023','http://192.168.1.104:8026','http://192.168.1.104:8029','http://192.168.1.104:8032'],'/v1/huyenan-render-final/jobs'),'transcode':(['http://192.168.1.104:8024','http://192.168.1.104:8027','http://192.168.1.104:8030','http://192.168.1.104:8033'],'/v1/huyenan-transcode-upload/jobs')}
def now():return datetime.datetime.now(datetime.timezone.utc).isoformat()
def sha(p):
 h=hashlib.sha256()
 with p.open('rb') as f:
  for b in iter(lambda:f.read(8<<20),b''):h.update(b)
 return h.hexdigest()
def atomic(p,o):
 p.parent.mkdir(parents=True,exist_ok=True);q=p.with_suffix('.tmp');q.write_text(json.dumps(o,ensure_ascii=False,indent=2)+'\n');os.replace(q,p)
def get(u):
 with urllib.request.urlopen(urllib.request.Request(u,headers={'Accept':'application/json'}),timeout=30) as r:return json.loads(r.read())
def post(u,p):
 d=json.dumps(p,ensure_ascii=False).encode();q=urllib.request.Request(u,data=d,headers={'Content-Type':'application/json'},method='POST')
 with urllib.request.urlopen(q,timeout=60) as r:return r.status,json.loads(r.read())
def free(base,route):
 h=get(base+'/health')
 if h.get('status')!='ok' or h.get('h264_nvenc') is not True:return False
 j=get(base+route);items=j if isinstance(j,list) else j.get('jobs')
 if not isinstance(items,list):raise RuntimeError('unknown jobs envelope')
 return not h.get('running') and not any(x.get('status') in ('queued','running') for x in items)
def main():
 a=argparse.ArgumentParser();a.add_argument('stage',choices=CFG);a.add_argument('--poll',type=int,default=5);x=a.parse_args();bases,route=CFG[x.stage];w=ROOT/'work/render'/x.stage;req=w/'request.json';intent=w/'intent.json';terminal=w/'terminal.json';bundle=json.loads(req.read_text());payload=bundle.get('request',bundle);assert isinstance(payload,dict);rh=sha(req)
 if intent.exists():
  i=json.loads(intent.read_text());assert i['request_sha256']==rh and i['run_id']==RUN
  if i.get('job_id'):
   base=i['base_url'];job=i['job_id']
  elif i.get('status') in ('ready_to_create','reconciled_no_job_id') and i.get('verified_no_job_id'):
   intent.unlink();base=None
  else:raise RuntimeError('ambiguous intent without job_id')
 if not intent.exists() or base is None:
  base=None
  while base is None:
   for b in bases:
    try:
     if free(b,route):base=b;break
    except Exception:pass
   if base is None:time.sleep(10)
  if not free(base,route):raise RuntimeError('selected replica changed capacity')
  i={'project_id':'031','run_id':RUN,'stage':x.stage,'status':'creating','base_url':base,'route':route,'request_sha256':rh,'request_path':str(req),'created_at':now()};atomic(intent,i)
  code,res=post(base+route,payload);job=res.get('job_id') or res.get('id')
  if not job:raise RuntimeError('ambiguous create response missing job id')
  i.update({'job_id':job,'http_status':code,'status':res.get('status','accepted'),'accepted_at':now()});atomic(intent,i)
 print(json.dumps({'accepted_or_resumed':True,'stage':x.stage,'base_url':base,'job_id':job}),flush=True)
 while True:
  try:r=get(base+route+'/'+str(job))
  except Exception as e:print('transient '+type(e).__name__,flush=True);time.sleep(x.poll);continue
  s=r.get('status');i.update({'status':s,'last_polled_at':now()});atomic(intent,i)
  if s in ('completed','failed','cancelled'):
   atomic(terminal,r);ok=s=='completed' and r.get('verification',{}).get('verified') is True;i.update({'verified':ok,'terminal_at':now()});atomic(intent,i);print(json.dumps({'stage':x.stage,'base_url':base,'job_id':job,'status':s,'server_verified':ok}),flush=True);return 0 if ok else 1
  time.sleep(x.poll)
if __name__=='__main__':raise SystemExit(main())