"""Native JEV JSON-lines inference; accepts state/question or state/questions."""
import argparse
import json
from pathlib import Path
import sys

ROOT=Path(__file__).resolve().parent
sys.path.insert(0,str(ROOT/'runtime'))
from maincode_jev_serve.model import DecisionModel, answer
from maincode_jev_serve.decide import decide

def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--checkpoint',default=str(ROOT))
    p.add_argument('--device',default='cuda:0')
    p.add_argument('--max-tokens',type=int,default=131072)
    p.add_argument('--token-budget',type=int,default=131072)
    p.add_argument('--batch-size',type=int,default=64)
    args=p.parse_args()
    model=DecisionModel(checkpoint=args.checkpoint,device=args.device)
    for line in sys.stdin:
        if not line.strip():continue
        row=json.loads(line)
        questions=row.get('questions')
        if questions is None:questions={'answer':row['question']}
        probs,tokens=decide(model,row['state'],questions,temperature=model.temperature,
            max_tokens=args.max_tokens,token_budget=args.token_budget,batch_size=args.batch_size,
            images=row.get('images',()))
        print(json.dumps({'model':'MATILDA JEV by Maincode','answers':{k:answer(questions[k],v) for k,v in probs.items()},
                          'usage':{'input_tokens':tokens}},ensure_ascii=False),flush=True)

if __name__=='__main__':main()