"""Exercise saved Gemma releases through the real approved-key HTTP API."""
import json
import os
from pathlib import Path
os.environ.setdefault('HF_HOME',str(Path(__file__).resolve().parent/'.cache/huggingface'))
import tempfile
import threading
from http.server import ThreadingHTTPServer
import urllib.request
import urllib.error
import torch
import serve_models
from api_access import AccessStore

ROOT=Path(__file__).resolve().parent
EXAMPLES={
    'open-emotions':dict(text='Thank you for helping me! I am relieved and grateful.'),
    'open-morals':dict(text='We should protect vulnerable people and treat everyone fairly.'),
    'open-sensations':dict(text='The aroma of freshly brewed coffee filled the kitchen.',target='aroma'),
}

def main():
    torch.set_num_threads(8)
    results={}
    with tempfile.TemporaryDirectory() as tmp:
        serve_models.ACCESS=AccessStore(Path(tmp)/'access.sqlite3')
        serve_models.ALLOW_ANONYMOUS=False
        server=ThreadingHTTPServer(('127.0.0.1',0),serve_models.Handler)
        thread=threading.Thread(target=server.serve_forever,daemon=True);thread.start()
        def post(route,body,key=None):
            headers={'Content-Type':'application/json'}
            if key:headers['Authorization']='Bearer '+key
            request=urllib.request.Request(f'http://127.0.0.1:{server.server_port}{route}',
                    data=json.dumps(body).encode(),headers=headers)
            try:response=urllib.request.urlopen(request,timeout=1200)
            except urllib.error.HTTPError as error:response=error
            with response:return response.status,json.load(response),dict(response.headers)
        try:
            status,registration,_=post('/v1/access-requests',dict(app_name='Gemma verification',contact='local@example.test'))
            assert status==201
            key=registration['api_key']
            assert post('/v1/emotions',EXAMPLES['open-emotions'],key)[0]==403
            serve_models.ACCESS.change(registration['request_id'],'approved',3,3)
            for route,task in serve_models.ROUTES.items():
                out=ROOT/'outputs/gemma'/task
                run=json.loads((out/'run.json').read_text())
                serve_models.MODEL_PATHS[task]=out/run['decision']['selected_arm']/'model'
                body={**EXAMPLES[task],'text':(EXAMPLES[task]['text']+' ')*12,'window_size':64,'overlap':16}
                status,result,_=post(route,body,key)
                assert status==200,(task,status,result)
                assert result['model']==f'mts-the-architect/{task}-gemma4-e2b-bf16-lora'
                assert len(result['windows'])>1 and result['task']==task
                expected=set(run['labels'])
                for window in result['windows']:
                    assert set(window['scores'])==expected
                    assert set(window['labels'])<=expected
                    assert all(0<=v<=1 for v in window['scores'].values())
                    assert window['input_tokens']<=64
                results[task]=dict(sliding_http=result,short_input=serve_models.predict(task,EXAMPLES[task]))
                print('Saved checkpoint + real HTTP contract passed',task,flush=True)
                serve_models.MODELS.clear();torch.cuda.empty_cache()
            status,_,headers=post('/v1/emotions',EXAMPLES['open-emotions'],key)
            assert status==429 and int(headers['Retry-After'])>0
            serve_models.ACCESS.change(registration['request_id'],'revoked')
            assert post('/v1/emotions',EXAMPLES['open-emotions'],key)[0]==403
        finally:
            server.shutdown();server.server_close();thread.join()
    path=ROOT/'outputs/gemma/inference_verification.json'
    path.write_text(json.dumps(dict(completed=True,approved_key_real_inference=True,
        rate_limit_verified=True,revocation_verified=True,examples=results),indent=2)+'\n')
    print('All inference checks passed; no client credentials retained.',flush=True)

if __name__=='__main__':main()