"""Read native JEV state/question JSON lines and write native decision answers."""
import argparse,json,sys
from pathlib import Path
from jev_fp4 import FP4DecisionModel
from kev.model import answer
from kev.decide import decide

def main():
    parser=argparse.ArgumentParser()
    parser.add_argument('--checkpoint',default=str(Path(__file__).resolve().parent))
    parser.add_argument('--device',default='cuda:0')
    args=parser.parse_args()
    model=FP4DecisionModel(args.checkpoint,device=args.device)
    for line in sys.stdin:
        if not line.strip():continue
        row=json.loads(line)
        if 'questions' in row:
            probabilities,tokens=decide(model,row['state'],row['questions'],temperature=model.temperature,
                                        max_tokens=131072,token_budget=131072,batch_size=64,images=row.get('images',()))
            result={'answers':{k:answer(row['questions'][k],v) for k,v in probabilities.items()},'usage':{'input_tokens':tokens}}
        else:
            result=answer(row['question'],model.predict([row],batch_size=1)[0])
        print(json.dumps(result,ensure_ascii=False),flush=True)

if __name__=='__main__':main()