"""TeenageWorks batch inference CLI. Built with Qwen; see LICENSE."""
import argparse
import json
from pathlib import Path
from toddler_jev import ToddlerJev

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--input', required=True, help='JSONL: audio, question, options; optional id')
    p.add_argument('--output', required=True)
    p.add_argument('--batch-size', type=int, default=1)
    p.add_argument('--model', default=str(Path(__file__).parent))
    p.add_argument('--revision', default='v1.0.0')
    p.add_argument('--device', default=None)
    args = p.parse_args()
    if args.batch_size < 1:
        p.error('--batch-size must be positive')
    source, target = Path(args.input).resolve(), Path(args.output).resolve()
    if source == target or target.exists():
        p.error('Output must be a new file')
    model = ToddlerJev.from_pretrained(args.model, revision=args.revision, device=args.device)
    def flush(batch, output):
        for row, result in zip(batch, model.predict_batch(batch)):
            if 'id' in row:
                result['id'] = row['id']
            output.write(json.dumps(result) + '\n')
        output.flush()
    with source.open(encoding='utf-8') as source_file, target.open('x', encoding='utf-8') as output:
        batch = []
        for line in source_file:
            if not line.strip():
                continue
            row = json.loads(line)
            audio = Path(row['audio'])
            row['audio'] = str(audio if audio.is_absolute() else source.parent / audio)
            batch.append(row)
            if len(batch) == args.batch_size:
                flush(batch, output)
                batch = []
        if batch:
            flush(batch, output)

if __name__ == '__main__':
    main()