"""Two-stage local authenticated restart check; synthetic payloads only."""
import argparse
import json
import os
from pathlib import Path
import httpx
from benchmarks.fixtures import fixture


def main():
    p=argparse.ArgumentParser();p.add_argument('stage',choices=['initialize','after-restart']);p.add_argument('--state',default='/tmp/needle-qualified-restart.json');a=p.parse_args()
    key=os.environ['NEEDLE_API_KEY']
    with httpx.Client(base_url='http://127.0.0.1:7860',headers={'Authorization':'Bearer '+key},timeout=60) as client:
        assert client.get('/ready').status_code==200
        assert client.get('/health').status_code==200
        assert client.post('/v1/route',headers={'Authorization':'Bearer invalid-test-key'},json={'input':'git_status','tools':[{'name':'git_status','parameters':{'type':'object'}}]}).status_code==401
        if a.stage=='initialize':
            c=fixture(1000);body={'model':'needle3','tools':c['tools'],'messages':c['messages'],'decision_context':{**c['decision_context'],'session_id':'restart-qualification','client_id':'local-qualification','workspace':'/src'}}
            r=client.post('/v1/chat/completions',json=body);r.raise_for_status();data=r.json()
            assert data['needle']['outcome']=='selected'
            body['messages']=[c['messages'][1]];body['decision_context'].update(incremental=True,previous_revision=data['needle']['revision'])
            Path(a.state).write_text(json.dumps(body))
            print(json.dumps({'initialized':True,'authentication_enforced':True}))
        else:
            body=json.loads(Path(a.state).read_text());r=client.post('/v1/chat/completions',json=body);r.raise_for_status();data=r.json()
            assert data['needle']['outcome']=='abstain' and data['needle']['reason']=='stale_state'
            assert not data['choices'][0]['message']['tool_calls']
            body['decision_context'].pop('incremental');body['decision_context'].pop('previous_revision')
            body['messages']=fixture(1000)['messages']
            recovered=client.post('/v1/chat/completions',json=body);recovered.raise_for_status()
            recovered=recovered.json();assert recovered['needle']['outcome']=='selected'
            calls=recovered['choices'][0]['message']['tool_calls']
            assert len(calls)==1 and calls[0]['function']['name']=='read_file'
            assert json.loads(calls[0]['function']['arguments'])=={'path':'/src/auth.py'}
            print(json.dumps({'restart_missing_session_failed_closed':True,'full_history_recovered_exact_route':True,'authentication_enforced':True,'ready':True}))

if __name__=='__main__':main()