import json,os,pathlib,selectors,subprocess,time,urllib.request
from preflight import g,R,IMAGE
os.umask(0o077)
plan=json.loads((R/'preflight.json').read_text())
env=g.load_env_file(str(g.ENV));base='http://127.0.0.1:'+env['FASTAPI_HOST_PORT']
opener=urllib.request.build_opener(urllib.request.ProxyHandler({}))
def cmd(args,timeout=30):return subprocess.run(args,capture_output=True,text=True,check=True,timeout=timeout)
def api(path):
 with opener.open(base+path,timeout=3) as f:return json.load(f)
def capture(args,path,limit,seconds,monitor=None):
 fd=os.open(path,os.O_WRONLY|os.O_CREAT|os.O_EXCL,0o600)
 p=subprocess.Popen(args,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,start_new_session=True)
 sel=selectors.DefaultSelector();sel.register(p.stdout,selectors.EVENT_READ)
 started=time.monotonic();size=0;reason=None;checked=0.
 try:
  with os.fdopen(fd,'wb') as output:
   while True:
    for key,_ in sel.select(.1):
     data=os.read(key.fileobj.fileno(),65536)
     if not data:sel.unregister(key.fileobj);continue
     allowed=min(len(data),limit-size);output.write(data[:allowed]);output.flush();size+=allowed
     if size>=limit:reason='LOG_CAP'
    if monitor and time.monotonic()-checked>=1:
     checked=time.monotonic()
     try:
      if monitor():reason='READY_OBSERVED'
     except Exception:reason='RESOURCE_OR_STATE_GUARD'
    if time.monotonic()-started>seconds:reason='TIME_CAP'
    if reason is not None and p.poll() is None:
     if monitor:cmd(['docker','stop','--timeout','20',g.NAME],35)
     else:p.terminate()
     p.wait(15)
    if p.poll() is not None and not sel.get_map():break
   os.fsync(output.fileno())
  return dict(returncode=p.wait(5),bytes=size,reason=reason,elapsed_s=time.monotonic()-started)
 finally:
  sel.close()
  if p.poll() is None:
   p.terminate()
   try:p.wait(5)
   except subprocess.TimeoutExpired:p.kill();p.wait(5)
def main():
 assert not (R/'run.json').exists()
 probe=capture(['python3','-c','import sys;print("OUT");print("ERR",file=sys.stderr);sys.exit(1)'],R/'collector-test.log',1024,3)
 assert probe['returncode']==1 and b'OUT' in (R/'collector-test.log').read_bytes() and b'ERR' in (R/'collector-test.log').read_bytes()
 state=dict(stage='prepared',start=0,stream_start=0,redis_changed=False,analytics_changed=False,collector_selftest='PASS')
 def save():(R/'run.json').write_text(json.dumps(state,indent=2))
 save()
 try:
  assert g.sha(g.ENV)==plan['env_hash'] and g.inspect()['Id']==plan['analytics']['Id']
  redis=json.loads(g.sh(['docker','inspect','contextra-redis-1']))[0];assert redis['Id']==plan['redis']['Id']
  assert g.sh(['docker','exec',redis['Id'],'redis-cli','DBSIZE']).strip()=='0'
  assert api('/api/v1/streams')['streams'][0]['state']=='stopped'
  assert g.gpu()<16384
  mem={l.split(':')[0]:int(l.split()[1]) for l in pathlib.Path('/proc/meminfo').read_text().splitlines()}
  assert mem['MemAvailable']>6*1024**2
  state['redis_changed']=True;save()
  cmd(plan['compose']+['up','-d','--no-deps','--force-recreate','--pull','never','redis'],120)
  end=time.monotonic()+45
  while time.monotonic()<end:
   redis=json.loads(g.sh(['docker','inspect','contextra-redis-1']))[0]
   assert redis['State']['Status']=='running' and not redis['RestartCount']
   if redis['State'].get('Health',{}).get('Status')=='healthy':break
   time.sleep(1)
  assert redis['State']['Health']['Status']=='healthy' and not redis['HostConfig']['PortBindings']
  state['analytics_changed']=True;save()
  cmd(plan['compose']+['up','--no-start','--no-deps','--force-recreate','--pull','never','analytics'],120)
  c=g.inspect();assert c['State']['Status']=='created' and c['Image']==IMAGE
  assert c['HostConfig']['RestartPolicy']['Name']=='no'
  state['candidate_id']=c['Id'];state['start']=1;state['stage']='observing';save()
  def monitor():
   c=g.inspect();assert c['Id']==state['candidate_id'] and not c['State']['OOMKilled']
   assert g.gpu()-plan['gpu_mib']<8192
   mem={l.split(':')[0]:int(l.split()[1]) for l in pathlib.Path('/proc/meminfo').read_text().splitlines()}
   assert mem['MemAvailable']>4*1024**2
   try:return c['State'].get('Health',{}).get('Status')=='healthy' and api('/ready')['ready']
   except Exception:return False
  state['capture']=capture(['docker','start','-a',g.NAME],R/'startup.private.log',2*1024*1024,150,monitor)
  c=g.inspect()
  state['candidate_final']={k:c['State'].get(k) for k in ['Status','ExitCode','OOMKilled','StartedAt','FinishedAt']}
  state['candidate_restart_count']=c['RestartCount']
  state['candidate_image']=c['Image'];state['stage']='captured';save()
 except Exception as error:state['stage']='failed';state['error']=type(error).__name__;save()
 finally:
  try:
   if state['analytics_changed']:
    c=g.inspect()
    if c['State']['Status']=='running':cmd(['docker','stop','--timeout','20',g.NAME],35)
    if not state.get('candidate_final'):
     state['candidate_final']={k:c['State'].get(k) for k in ['Status','ExitCode','OOMKilled']}
     log=subprocess.run(['docker','logs','--tail','150',g.NAME],capture_output=True,timeout=10)
     (R/'fallback.private.log').write_bytes((log.stdout+log.stderr)[-1024*1024:])
    save()
   for name,changed in [('redis','redis_changed'),('analytics','analytics_changed')]:
    if not state[changed]:continue
    previous=plan[name];files=plan['previous'][previous['Id']]
    assert all(g.sha(p)==h for p,h in files.items())
    args=['docker','compose','--env-file',str(g.ENV)]
    for p in files:args+=['-f',p]
    cmd(args+['up','-d','--no-deps','--force-recreate','--pull','never',name],120)
   policy=plan['analytics']['HostConfig']['RestartPolicy']
   value=policy['Name']+(':'+str(policy['MaximumRetryCount']) if policy.get('MaximumRetryCount') else '')
   cmd(['docker','update','--restart='+value,g.NAME])
   end=time.monotonic()+120
   while time.monotonic()<end:
    c=g.inspect()
    try:
     if c['State'].get('Health',{}).get('Status')=='healthy' and api('/ready')['ready']:break
    except Exception:pass
    time.sleep(2)
   assert c['State']['Health']['Status']=='healthy'
   assert api('/api/v1/streams')['streams'][0]['state']=='stopped'
   redis=json.loads(g.sh(['docker','inspect','contextra-redis-1']))[0]
   assert not redis['HostConfig']['PortBindings']
   assert g.sha(g.ENV)==plan['env_hash'] and g.sha(g.CFG)==plan['config_hash']
   assert g.sha(g.CFG.parent/'2/model.plan')==plan['plan_hash']
   assert g.sha(pathlib.Path(env['CONTEXTRA_RUNTIME_STATE_ROOT'])/'streams_runtime.yaml')==plan['overlay_hash']
   assert {k:v for k,v in g.others().items() if k!=redis['Id']}==plan['other_containers']
   state['restored']=dict(analytics='healthy',stream='stopped',redis_port=0,env_model_overlay_unchanged=True,
                         non_targets_unchanged=True,restart_policy=c['HostConfig']['RestartPolicy'])
  except Exception as error:
   state['restore_error']=type(error).__name__
   try:cmd(['docker','stop','--timeout','20',g.NAME],35)
   except Exception:state['safety_stop_failed']=True
  save()
 print(json.dumps({k:state.get(k) for k in ['stage','start','stream_start','capture','candidate_final','restored','error','restore_error']}))
if __name__=='__main__':main()