#!/usr/bin/env python3
import argparse,json,subprocess,tempfile
from pathlib import Path
HERE=Path(__file__).resolve().parent
BIN=HERE/'bin/slime-bench'
def run(*args,ok=True):
    p=subprocess.run([str(BIN),*map(str,args)],capture_output=True,text=True,timeout=180)
    if ok:
        assert p.returncode==0,p.stderr
        return json.loads(p.stdout)
    assert p.returncode!=0,p.stdout

def term(x,z):
    def i32(v):return ((v+2**31)%2**32)-2**31
    return i32(x*x*4987142)+i32(x*5947611)+i32(z*z)*4392871+i32(z*389711)
def slime(seed,x,z):
    s=((seed+term(x,z))^987234911^0x5DEECE66D)&((1<<48)-1)
    while True:
        s=(s*0x5DEECE66D+11)&((1<<48)-1);q=s>>17
        if q<2147483640:return q%10==0

def main():
    ap=argparse.ArgumentParser();ap.add_argument('--gpu',action='store_true');args=ap.parse_args()
    for radius in [0,1,2,31,1000000]:
        r=run('--x',-123,'--z',456,'--radius',radius,'--size',6,'--dry-run')
        assert r['anchors']==1+2*radius*(radius+1)
        assert r['first_anchor']==[-123,456-radius] and r['last_anchor']==[-123,456+radius]
    for bad in [[],['--radius','-1'],['--radius','1','--size','0'],['--radius','1','--size','7'],['--radius','abc'],['--radius','0','--x','2147483647'],['--radius','1','--repeat','0'],['--radius','1','--position-batch','0']]:run(*bad,'--dry-run',ok=False)
    summary={'cli_scope_and_errors':True}
    if args.gpu:
        common=['--x',51069,'--z',419124]
        reference=run(*common,'--radius',0,'--size',5,'--repeat',2,'--max-results',100)
        expected={61692151016558,15959335054446}
        assert all(r['matches']==24 and r['root_jobs']==84 for r in reference['rounds'])
        samples=[set(map(tuple,r['sample'])) for r in reference['rounds']]
        assert samples[0]==samples[1] and expected <= {s for s,x,z in samples[0]}
        sliced=run(*common,'--radius',0,'--size',5,'--job-batch',17,'--max-results',100,'--cache-classes',0)
        assert set(map(tuple,sliced['rounds'][0]['sample']))==samples[0]
        partial=run(*common,'--radius',0,'--size',5,'--max-results',3)
        assert partial['rounds'][0]['matches']==24 and len(partial['rounds'][0]['sample'])==3 and partial['rounds'][0]['sample_truncated']
        for n in [4,6]:
            r=run(*common,'--radius',0,'--size',n,'--max-results',1000)
            for seed,x,z in r['rounds'][0]['sample']:assert all(slime(seed,x+dx,z+dz) for dx in range(n) for dz in range(n))
            if n>=5:
                wanted={t for t in samples[0] if all(slime(t[0],t[1]+dx,t[2]+dz) for dx in range(n) for dz in range(n))}
                assert set(map(tuple,r['rounds'][0]['sample']))==wanted
            summary[f'size_{n}']=r['rounds'][0]['matches']
        diamond=run(*common,'--radius',2,'--size',5,'--max-results',1000,'--position-batch',3)
        split=[]
        for dz in range(-2,3):
            for dx in range(-2,3):
                if abs(dx)+abs(dz)>2:continue
                r=run('--x',51069+dx,'--z',419124+dz,'--radius',0,'--size',5,'--max-results',1000)
                split.extend(map(tuple,r['rounds'][0]['sample']))
        assert sorted(split)==sorted(map(tuple,diamond['rounds'][0]['sample']))
        assert diamond['anchors']==13
        run('--radius',0,'--devices','0,0',ok=False)
        with tempfile.TemporaryDirectory() as tmp:
            out=Path(tmp)/'report.json'
            r=run(*common,'--radius',0,'--output',out)
            assert json.loads(out.read_text())==r
            run(*common,'--radius',0,'--output',out,ok=False)
        summary['gpu_known_matches']=24;summary['diamond_vs_individual_anchors']=True
    print(json.dumps(summary,indent=2))
if __name__=='__main__':main()