"""Load a saved ARTI reading-QA program and answer your own JSONL questions.

Input: one {id, context, question} object per line. Public model weights remain
in --work. Answers retain original IDs; this is separate from the benchmark.
"""
import argparse
import json
from pathlib import Path

from huggingface_hub import snapshot_download
from arti.language_recipe import CausalLMPredictor
from reading_qa import BASE, REVISION, TEMPLATE, messages


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--work', type=Path, default=Path('reading-qa-work'))
    parser.add_argument('--asset', type=Path, required=True)
    parser.add_argument('--language', choices=['en','zh'], required=True)
    parser.add_argument('--input', type=Path, required=True)
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    base = snapshot_download(BASE, revision=REVISION, cache_dir=args.work / 'hub-cache', local_files_only=True)
    predictor = CausalLMPredictor.from_pretrained(args.asset, device='cuda', base_model=base, local_files_only=True)
    records = []
    seen = set()
    for line in args.input.read_text(encoding='utf8').splitlines():
        if not line.strip():
            continue
        item = json.loads(line)
        if not isinstance(item.get('id'), str) or not item['id'] or item['id'] in seen:
            raise ValueError('each question needs a distinct nonempty string id')
        seen.add(item['id'])
        for field in ('context', 'question'):
            if not isinstance(item.get(field), str) or not item[field].strip():
                raise ValueError(f'{field} must be a nonempty string')
        turns = messages('squad' if args.language=='en' else 'cmrc', item['context'], item['question'])
        prompt = predictor.tokenizer.apply_chat_template(turns, tokenize=True, add_generation_prompt=True, **TEMPLATE)
        if len(prompt) + 32 > 384:
            raise ValueError('this tutorial supports complete prompt + generation budget <= 384 tokens; shorten the source deliberately')
        records.append({'id': item['id'], 'messages': turns})
    if not records:
        raise ValueError('input contains no questions')
    results = predictor.predict_records(records, batch_size=8, max_new_tokens=32)
    assert [x['id'] for x in results] == [x['id'] for x in records]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open('w', encoding='utf8') as stream:
        for row in results:
            stream.write(json.dumps(row, ensure_ascii=False) + '\n')
    print(json.dumps({'written':len(results),'output':str(args.output),'base':BASE,'revision':REVISION}))


if __name__ == '__main__':
    main()