"""Finite local job runner; one compute worker, sampled process-tree RSS watchdog."""
import argparse, datetime, hashlib, json, os, pathlib, resource, signal, subprocess, time

ROOT = pathlib.Path(__file__).resolve().parent
LEDGER = ROOT / 'attempts.json'

def main():
    p=argparse.ArgumentParser(); p.add_argument('--name',required=True); p.add_argument('--timeout',type=int,default=120); p.add_argument('--cwd'); p.add_argument('command',nargs=argparse.REMAINDER); a=p.parse_args()
    old=json.loads(LEDGER.read_text()) if LEDGER.exists() else []
    cpu=sum(x['cpu_seconds'] for x in old); wall=sum(x['wall_seconds'] for x in old)
    if cpu>=3600 or wall+a.timeout>5400: raise SystemExit('Aggregate budget unavailable')
    command=a.command[1:] if a.command[:1]==['--'] else a.command
    env=dict(os.environ, OMP_NUM_THREADS='1',OPENBLAS_NUM_THREADS='1',MKL_NUM_THREADS='1',VECLIB_MAXIMUM_THREADS='1',NUMEXPR_NUM_THREADS='1',PYTHONDONTWRITEBYTECODE='1')
    start=time.monotonic(); usage=resource.getrusage(resource.RUSAGE_CHILDREN); stamp=datetime.datetime.now(datetime.timezone.utc).isoformat()
    def limits():
        os.setsid(); resource.setrlimit(resource.RLIMIT_CPU,(min(300,a.timeout),min(301,a.timeout+1)))
    log=ROOT/(a.name+'.log'); reason=None; peak=0
    with log.open('w') as output:
        child=subprocess.Popen(command,cwd=a.cwd,env=env,stdout=output,stderr=subprocess.STDOUT,preexec_fn=limits)
        while child.poll() is None:
            rows=subprocess.run(['ps','-axo','pid=,ppid=,rss='],capture_output=True,text=True).stdout.splitlines()
            table=[tuple(map(int,r.split())) for r in rows if len(r.split())==3]
            ids={child.pid}
            for _ in range(8): ids.update(pid for pid,ppid,rss in table if ppid in ids)
            total=sum(rss*1024 for pid,ppid,rss in table if pid in ids); peak=max(peak,total)
            if total>1024**3: reason='RSS watchdog ceiling exceeded'
            if time.monotonic()-start>a.timeout: reason='wall watchdog exceeded'
            if reason: os.killpg(child.pid,signal.SIGKILL); break
            time.sleep(.1)
        rc=child.wait()
    end=resource.getrusage(resource.RUSAGE_CHILDREN)
    record=dict(name=a.name,command=command,started_utc=stamp,returncode=rc,failure=reason,cpu_seconds=end.ru_utime+end.ru_stime-usage.ru_utime-usage.ru_stime,wall_seconds=time.monotonic()-start,peak_sampled_tree_rss_bytes=peak,timeout_seconds=a.timeout,environment={k:env[k] for k in ['OMP_NUM_THREADS','OPENBLAS_NUM_THREADS','MKL_NUM_THREADS','VECLIB_MAXIMUM_THREADS','NUMEXPR_NUM_THREADS']},memory_enforcement='100 ms sampled process-tree RSS kill watchdog, not an address-space or instantaneous hard RSS limit')
    old.append(record); LEDGER.write_text(json.dumps(old,indent=2)+'\n')
    print(log.read_text()[-12000:]); print(json.dumps(record)); raise SystemExit(rc if rc else bool(reason))

if __name__=='__main__': main()
