#!/usr/bin/env python
import argparse, json, re, time, os, platform, shutil, subprocess, sys, threading
from collections import OrderedDict, deque
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import av
import numpy as np
from datasets import load_dataset

def sec(s):
    p = [int(x) for x in str(s).split(':')]
    return p[-1] + 60 * p[-2] + (3600 * p[-3] if len(p) > 2 else 0)

def norm(x):
    return re.sub(r'\s+', ' ', re.sub(r'<think>.*?</think>', '', str(x or ''), flags=re.S)).strip().lower().strip(' .')

def correct(pred, gold, fmt):
    p, g = norm(pred), norm(gold)
    if fmt == 'binary': return (p.startswith('yes') and g.startswith('yes')) or (p.startswith('no') and g.startswith('no'))
    if fmt in {'number', 'percentage'}:
        a = re.search(r'-?\d+(?:\.\d+)?', p); b = re.search(r'-?\d+(?:\.\d+)?', g)
        return bool(a and b and float(a.group()) == float(b.group()))
    if fmt == 'time':
        a = re.search(r'\d{1,2}:\d{2}:\d{2}', p); b = re.search(r'\d{1,2}:\d{2}:\d{2}', g)
        return bool(a and b and a.group() == b.group())
    return p == g or (g and g in p)


def sample_frames(out, n):
    out = [out[i] for i in np.linspace(0, len(out) - 1, min(n, len(out))).round().astype(int)]
    return out + [out[0]] if len(out) == 1 else out


def frames(src, start, end, n=64):
    def decode():
        c = av.open(str(src)); s = c.streams.video[0]
        c.seek(max(0, int(start / float(s.time_base))), stream=s, any_frame=True)
        out = []
        for index, f in enumerate(c.decode(s)):
            t = float(f.pts * s.time_base) if f.pts is not None else None
            if index == 0 and t is not None and t < max(0, start - 60):
                c.close(); return []
            if t is not None and t > end: break
            image = f.to_ndarray(format='rgb24')
            if t is None or t >= start: out.append(image)
        c.close(); return out
    out = decode()
    if not out:
        # Keep the run complete for corrupt/misindexed segments; mark them via zero-valued frames.
        out = [np.zeros((224, 224, 3), dtype=np.uint8)] * n
    out = sample_frames(out, n) if out else []
    if not out: raise RuntimeError(f'no frames in {src} for [{start},{end}]')
    arr = np.stack(out)
    return arr, {'fps': 5.0, 'duration': len(arr) / 5.0, 'total_num_frames': len(arr),
                 'frames_indices': list(range(len(arr))), 'video_backend': 'opencv',
                 'do_sample_frames': False}


def fixed64_frames(src, output_fps, n=64):
    c = av.open(str(src)); s = c.streams.video[0]
    out = [f.to_ndarray(format='rgb24') for f in c.decode(s)]
    c.close()
    if len(out) != n:
        raise RuntimeError(f'expected {n} fixed frames in {src}, decoded {len(out)}')
    fps = float(output_fps)
    if fps <= 0:
        raise RuntimeError(f'invalid output_fps={output_fps} for {src}')
    arr = np.stack(out)
    return arr, {'fps': fps, 'duration': len(arr) / fps, 'total_num_frames': len(arr),
                 'frames_indices': list(range(len(arr))), 'video_backend': 'opencv',
                 'do_sample_frames': False}


def fixed_absmeta_frames(src, metadata, n=80):
    c = av.open(str(src))
    stream = c.streams.video[0]
    out = [f.to_ndarray(format='rgb24') for f in c.decode(stream)]
    c.close()
    if len(out) != n:
        raise RuntimeError(f'expected {n} fixed frames in {src}, decoded {len(out)}')
    if not isinstance(metadata, dict):
        raise RuntimeError(f'missing video_metadata for {src}')
    frame_indices = metadata.get('frames_indices')
    if metadata.get('video_backend') != 'precomputed_absolute_pts' or not isinstance(frame_indices, list):
        raise RuntimeError(f'invalid absolute video_metadata for {src}')
    if len(frame_indices) != n:
        raise RuntimeError(f'expected {n} frame indices in video_metadata for {src}, got {len(frame_indices)}')
    return np.stack(out), dict(metadata)


class VideoCache:
    def __init__(self, root, max_gib):
        self.root, self.limit = Path(root), int(max_gib * 1024 ** 3)
        self.root.mkdir(parents=True, exist_ok=True)
        self.entries, self.size, self.lock = OrderedDict(), 0, threading.Lock()
        self.hits = self.misses = self.bytes = 0

    def get(self, src):
        src = Path(src)
        key = f'{src.parent.parent.name}_{src.parent.name}_{src.name}'
        dst = self.root / key
        with self.lock:
            if dst.exists():
                self.hits += 1; self.entries[dst] = dst.stat().st_size; self.entries.move_to_end(dst)
                return dst
            self.misses += 1
            size = src.stat().st_size
            while self.entries and self.size + size > self.limit:
                old, old_size = self.entries.popitem(last=False)
                old.unlink(missing_ok=True); self.size -= old_size
            if size > self.limit:
                return src
            tmp = dst.with_suffix(dst.suffix + f'.{os.getpid()}.partial')
            try:
                shutil.copyfile(src, tmp); os.replace(tmp, dst)
                self.entries[dst] = size; self.size += size; self.bytes += size
                return dst
            except OSError:
                tmp.unlink(missing_ok=True)
                return src


def read_jsonl(path):
    if not path.exists(): return []
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def load_rows(data_root):
    rows = []
    for dataset_name, rel in [('heico', 'heico-focus-vqa/data/segment/test.parquet'), ('lapchole', 'lapchole-focus-vqa/data/segment/test.parquet')]:
        ds = load_dataset('parquet', data_files=str(Path(data_root) / rel), split='train')
        rows.extend([{**r, '_dataset': dataset_name, '_input_index': len(rows) + i} for i, r in enumerate(ds)])
    return rows


def load_fixed64_rows(dataset_jsonl):
    path = Path(dataset_jsonl)
    rows = []
    for i, line in enumerate(path.read_text().splitlines()):
        if not line.strip():
            continue
        r = json.loads(line)
        if int(r.get('derived_frame_count', 0)) != 64 or int(r.get('verified_frame_count', 0)) != 64:
            raise ValueError(f'row {i} id={r.get("id")} is not verified fixed64')
        rows.append({**r, '_dataset': str(r.get('dataset', 'unknown')).replace('-focus-vqa', ''),
                     '_input_index': len(rows), '_fixed64_root': str(path.parent)})
    return rows


def row_id(row):
    if row.get('id') is not None:
        return row['id']
    metadata = row.get('segment_metadata') or {}
    if metadata.get('id') is None:
        raise ValueError('row is missing both id and segment_metadata.id')
    return metadata['id']


def row_answer_format(row):
    return row.get('answer_format') or (row.get('segment_metadata') or {}).get('answer_format')


def row_video_name(row):
    return row.get('video') or (row.get('segment_metadata') or {}).get('video', '')


def row_timestamp_start(row):
    return row.get('timestamp_start') or (row.get('segment_metadata') or {}).get('timestamp_start', '00:00:00')


def load_fixed80_absmeta_requests(dataset_jsonl, nframes=80):
    path = Path(dataset_jsonl)
    rows = []
    for i, line in enumerate(path.read_text().splitlines()):
        if not line.strip():
            continue
        r = json.loads(line)
        qid = row_id(r)
        if 'answer' in r:
            raise ValueError(f'row {i} id={qid} unexpectedly contains an answer')
        messages = r.get('messages')
        if not isinstance(messages, list) or [m.get('role') for m in messages] != ['system', 'user']:
            raise ValueError(f'row {i} id={qid} must contain system,user messages only')
        user_text = messages[1].get('content')
        if not isinstance(user_text, str) or user_text.count('<video>') != 1:
            raise ValueError(f'row {i} id={qid} user message must contain exactly one <video> marker')
        if 'answer' in r or any(m.get('role') == 'assistant' for m in messages):
            raise ValueError(f'row {i} id={qid} unexpectedly contains answer content')
        videos = r.get('videos')
        if not isinstance(videos, list) or len(videos) != 1:
            raise ValueError(f'row {i} id={qid} must contain exactly one video')
        video_path = Path(videos[0])
        if not video_path.exists():
            raise FileNotFoundError(f'row {i} id={qid} missing video: {video_path}')
        metadata = (r.get('chat_template_kwargs') or {}).get('video_metadata')
        if not isinstance(metadata, dict):
            raise ValueError(f'row {i} id={qid} missing chat_template_kwargs.video_metadata')
        frame_indices = metadata.get('frames_indices')
        if metadata.get('video_backend') != 'precomputed_absolute_pts' or not isinstance(frame_indices, list):
            raise ValueError(f'row {i} id={qid} has invalid absolute video_metadata')
        if len(frame_indices) != nframes:
            raise ValueError(f'row {i} id={qid} expected {nframes} frame indices, got {len(frame_indices)}')
        rows.append({**r, 'id': qid, '_dataset': str((r.get('segment_metadata') or {}).get('dataset', 'unknown')).replace('-focus-vqa', ''),
                     '_input_index': len(rows), '_fixed80_absmeta': True})
    return rows


def detect_dataset_format(dataset_jsonl):
    path = Path(dataset_jsonl)
    for line in path.read_text().splitlines():
        if line.strip():
            row = json.loads(line)
            if 'chat_template_kwargs' in row and 'segment_metadata' in row:
                return 'fixed80_absmeta_requests'
            return 'fixed64_with_refs'
    raise ValueError(f'empty dataset: {dataset_jsonl}')


def load_dataset_jsonl(dataset_jsonl, dataset_format, nframes):
    fmt = detect_dataset_format(dataset_jsonl) if dataset_format == 'auto' else dataset_format
    if fmt == 'fixed80_absmeta_requests':
        return load_fixed80_absmeta_requests(dataset_jsonl, nframes), fmt
    if fmt == 'fixed64_with_refs':
        return load_fixed64_rows(dataset_jsonl), fmt
    raise ValueError(f'unknown dataset format: {dataset_format}')


def pending_rows(rows, existing, order='input'):
    row_by_id = {str(row_id(r)): r for r in rows}
    ids = [str(row_id(r)) for r in existing]
    if len(ids) != len(set(ids)): raise ValueError('resume file contains duplicate IDs')
    if set(ids) - row_by_id.keys(): raise ValueError('resume file contains IDs outside this split')
    pending = [r for r in rows if str(row_id(r)) not in set(ids)]
    if order == 'video_time':
        pending.sort(key=lambda r: (r['_dataset'], row_video_name(r), sec(row_timestamp_start(r)), r['_input_index']))
    return row_by_id, set(ids), pending


def resolve_checkpoint(checkpoint, train_run_dir=None, require_complete=False, expected_step=None):
    if checkpoint != 'auto-final':
        ckpt = Path(checkpoint)
    else:
        if not train_run_dir:
            raise ValueError('--train-run-dir is required when --checkpoint auto-final')
        roots = [Path(train_run_dir) / 'results' / 'train', Path(train_run_dir)]
        checkpoints = []
        for root in roots:
            checkpoints.extend(root.glob('**/checkpoint-*'))
        if not checkpoints:
            raise FileNotFoundError(f'no checkpoints found under {train_run_dir}')
        ckpt = max(checkpoints, key=lambda p: int(p.name.split('-')[-1]))
    required = ['adapter_config.json', 'adapter_model.safetensors', 'trainer_state.json']
    missing = [name for name in required if not (ckpt / name).exists()]
    if missing:
        raise FileNotFoundError(f'incomplete checkpoint {ckpt}: missing {missing}')
    state = json.loads((ckpt / 'trainer_state.json').read_text())
    step = int(state.get('global_step') or ckpt.name.split('-')[-1])
    max_steps = state.get('max_steps')
    if expected_step is not None and step != expected_step:
        raise ValueError(f'checkpoint {ckpt} step {step} != expected final step {expected_step}')
    if require_complete and max_steps is not None and step != int(max_steps):
        raise ValueError(f'checkpoint {ckpt} is not final: global_step={step}, max_steps={max_steps}')
    return ckpt, {'global_step': step, 'epoch': state.get('epoch'), 'max_steps': max_steps}


def make_messages(row):
    if row.get('_fixed80_absmeta'):
        messages = []
        for message in row['messages']:
            if message['role'] == 'system':
                messages.append({'role': 'system', 'content': message['content']})
            elif message['role'] == 'user':
                text = message['content'].replace('<video>', '', 1).strip()
                messages.append({'role': 'user', 'content': [{'type': 'video'}, {'type': 'text', 'text': text}]})
        return messages
    return [{'role': 'user', 'content': [{'type': 'video'}, {'type': 'text', 'text': row['question']}]}]


def strip_prediction(text):
    return re.sub(r'<think>.*?</think>', '', str(text), flags=re.S).strip()


def validate_rows(rows, nframes, dataset_format, decode_samples=0):
    ids = [str(row_id(r)) for r in rows]
    if len(ids) != len(set(ids)):
        raise ValueError('dataset contains duplicate request IDs')
    missing_videos = 0
    frame_lengths = {}
    backends = {}
    user_video_markers = {}
    message_role_patterns = {}
    system_prompts = set()
    for r in rows:
        if dataset_format == 'fixed80_absmeta_requests':
            videos = r.get('videos') or []
            missing_videos += sum(1 for v in videos if not Path(v).exists())
            roles = tuple(m.get('role') for m in (r.get('messages') or []))
            message_role_patterns[str(roles)] = message_role_patterns.get(str(roles), 0) + 1
            system_prompts.add((r.get('messages') or [{}])[0].get('content'))
            user_content = (r.get('messages') or [{}, {}])[1].get('content', '')
            markers = user_content.count('<video>') if isinstance(user_content, str) else -1
            user_video_markers[markers] = user_video_markers.get(markers, 0) + 1
            metadata = (r.get('chat_template_kwargs') or {}).get('video_metadata') or {}
            frame_lengths[len(metadata.get('frames_indices') or [])] = frame_lengths.get(len(metadata.get('frames_indices') or []), 0) + 1
            backends[metadata.get('video_backend')] = backends.get(metadata.get('video_backend'), 0) + 1
        elif dataset_format == 'fixed64_with_refs':
            src = Path(r['_fixed64_root']) / r['derived_video']
            missing_videos += 0 if src.exists() else 1
    for r in rows[:max(0, decode_samples)]:
        if dataset_format == 'fixed80_absmeta_requests':
            fixed_absmeta_frames(Path(r['videos'][0]), (r.get('chat_template_kwargs') or {})['video_metadata'], nframes)
        elif dataset_format == 'fixed64_with_refs':
            fixed64_frames(Path(r['_fixed64_root']) / r['derived_video'], r['output_fps'], nframes)
    return {'rows': len(rows), 'unique_ids': len(set(ids)), 'duplicate_ids': len(ids) - len(set(ids)),
            'missing_video_files': missing_videos, 'frames_indices_len_counts': frame_lengths,
            'video_metadata_backend_counts': backends, 'message_role_patterns': message_role_patterns,
            'unique_system_prompts': len(system_prompts), 'user_video_marker_counts': user_video_markers,
            'decode_checked_rows': max(0, decode_samples)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', required=True); ap.add_argument('--checkpoint')
    ap.add_argument('--train-run-dir')
    ap.add_argument('--require-complete-checkpoint', action='store_true')
    ap.add_argument('--expected-checkpoint-step', type=int)
    ap.add_argument('--data-root'); ap.add_argument('--dataset-jsonl'); ap.add_argument('--output', required=True)
    ap.add_argument('--dataset-format', choices=('auto', 'fixed64_with_refs', 'fixed80_absmeta_requests'), default='auto')
    ap.add_argument('--limit', type=int, default=0); ap.add_argument('--shard-index', type=int, default=0)
    ap.add_argument('--num-shards', type=int, default=1); ap.add_argument('--nframes', type=int, default=64)
    ap.add_argument('--tensor-parallel-size', type=int, default=4)
    ap.add_argument('--batch-size', type=int, default=2)
    ap.add_argument('--max-model-len', type=int, default=16384)
    ap.add_argument('--max-new-tokens', type=int, default=64)
    ap.add_argument('--gpu-memory-utilization', type=float, default=0.9)
    ap.add_argument('--max-num-seqs', type=int)
    ap.add_argument('--mm-processor-cache-gb', type=float)
    ap.add_argument('--lora-name', default='segmentqa_tr005')
    ap.add_argument('--max-lora-rank', type=int, default=64)
    ap.add_argument('--disable-lora', action='store_true',
                    help='Run a fully merged model without dynamic vLLM LoRA activation.')
    ap.add_argument('--merged-source-checkpoint',
                    help='Source adapter checkpoint used to create a fully merged model.')
    ap.add_argument('--merge-report',
                    help='Merge verification report path for a fully merged model.')
    ap.add_argument('--disable-tower-connector-lora', action='store_true')
    ap.add_argument('--enforce-eager', action='store_true',
                    help='Disable torch.compile/CUDAGraphs for vLLM startup paths that are incompatible with multimodal LoRA.')
    ap.add_argument('--resume', action='store_true')
    ap.add_argument('--resume-from')
    ap.add_argument('--execution-order', choices=('input', 'video_time'), default='input')
    ap.add_argument('--decode-workers', type=int, default=4)
    ap.add_argument('--prefetch-batches', type=int, default=1)
    ap.add_argument('--video-cache-dir')
    ap.add_argument('--video-cache-gib', type=float, default=64)
    ap.add_argument('--results-subdir', action='store_true')
    ap.add_argument('--run-id', default='unknown'); ap.add_argument('--issue', action='append', default=[])
    ap.add_argument('--dry-run', action='store_true')
    ap.add_argument('--dry-run-decode-samples', type=int, default=0)
    args = ap.parse_args()
    if bool(args.data_root) == bool(args.dataset_jsonl):
        ap.error('provide exactly one of --data-root or --dataset-jsonl')
    if not args.disable_lora and not args.checkpoint:
        ap.error('--checkpoint is required unless --disable-lora is set for a fully merged model')
    out = Path(args.output); out.mkdir(parents=True, exist_ok=True)
    results_dir = out / 'results' if args.results_subdir else out
    results_dir.mkdir(parents=True, exist_ok=True)
    if args.checkpoint:
        ckpt, ckpt_info = resolve_checkpoint(
            args.checkpoint, args.train_run_dir, args.require_complete_checkpoint, args.expected_checkpoint_step)
    else:
        ckpt, ckpt_info = None, {}
    if args.dataset_jsonl:
        rows, dataset_format = load_dataset_jsonl(args.dataset_jsonl, args.dataset_format, args.nframes)
    else:
        rows, dataset_format = load_rows(args.data_root), 'source_parquet'
    if args.limit: rows = rows[:args.limit]
    rows = rows[args.shard_index::args.num_shards]
    raw_path = results_dir / 'raw_predictions.jsonl'
    if args.resume_from and not raw_path.exists(): shutil.copyfile(args.resume_from, raw_path)
    existing = read_jsonl(raw_path) if args.resume or args.resume_from else []
    row_by_id, done, pending = pending_rows(rows, existing, args.execution_order)
    validation = validate_rows(rows, args.nframes, dataset_format, args.dry_run_decode_samples)
    if args.dry_run:
        report = {'status': 'dry_run_passed', 'run_id': args.run_id, 'dataset_format': dataset_format,
                  'validation': validation, 'checkpoint': str(ckpt) if ckpt else None, 'checkpoint_info': ckpt_info,
                  'model': args.model, 'max_new_tokens': args.max_new_tokens, 'temperature': 0,
                  'lora': {'enabled': not args.disable_lora, 'max_lora_rank': args.max_lora_rank,
                           'enable_tower_connector_lora': not args.disable_tower_connector_lora},
                  'pending': len(pending), 'completed_at_start': len(done)}
        (results_dir / 'dry_run_report.json').write_text(json.dumps(report, indent=2))
        print(json.dumps(report, indent=2))
        return
    from vllm import LLM, SamplingParams
    commit = subprocess.run(['git', 'rev-parse', 'HEAD'], text=True, capture_output=True).stdout.strip() or 'unknown'
    dirty = bool(subprocess.run(['git', 'status', '--porcelain'], text=True, capture_output=True).stdout)
    runmeta = {'status': 'running', 'run_id': args.run_id, 'run_name': out.name, 'task_type': 'test',
        'backend': 'vllm', 'model': args.model, 'checkpoint': str(ckpt) if ckpt else None, 'checkpoint_info': ckpt_info,
        'count': len(rows), 'completed_at_start': len(done), 'pending': len(pending),
        'nframes': args.nframes, 'batch_size': args.batch_size, 'max_model_len': args.max_model_len,
        'max_new_tokens': args.max_new_tokens, 'dataset_format': dataset_format, 'validation': validation,
        'lora': {'enabled': not args.disable_lora, 'name': args.lora_name, 'max_lora_rank': args.max_lora_rank,
                 'enable_tower_connector_lora': not args.disable_tower_connector_lora},
        'merged_source_checkpoint': args.merged_source_checkpoint,
        'merge_report': args.merge_report,
        'enforce_eager': args.enforce_eager,
        'tensor_parallel_size': args.tensor_parallel_size, 'precision': 'bfloat16', 'temperature': 0,
        'execution_order': args.execution_order, 'decode_workers': args.decode_workers,
        'prefetch_batches': args.prefetch_batches, 'video_cache_dir': args.video_cache_dir,
        'video_cache_gib': args.video_cache_gib, 'data_root': args.data_root,
        'dataset_jsonl': args.dataset_jsonl, 'gpu_memory_utilization': args.gpu_memory_utilization,
        'max_num_seqs': args.max_num_seqs,
        'dataset_split': 'official SegmentQA test', 'dataset_purpose': 'test',
        'python': sys.executable, 'host': platform.node(), 'pid': os.getpid(),
        'devices': os.environ.get('CUDA_VISIBLE_DEVICES', 'all-visible'), 'argv': sys.argv,
        'code_commit': commit, 'dirty': dirty, 'issues': args.issue, 'started_at': time.time(),
        'outputs': {'raw_predictions': str(raw_path), 'answer': str(results_dir / 'answer.json'),
                    'metrics': str(results_dir / 'metrics.json')}}
    (out / 'run.json').write_text(json.dumps(runmeta, indent=2))
    llm_kwargs = {'model': args.model, 'dtype': 'bfloat16', 'tensor_parallel_size': args.tensor_parallel_size,
                  'trust_remote_code': True, 'limit_mm_per_prompt': {'video': 1},
                  'max_model_len': args.max_model_len,
                  'gpu_memory_utilization': args.gpu_memory_utilization,
                  'enforce_eager': args.enforce_eager}
    if not args.disable_lora:
        llm_kwargs.update(enable_lora=True,
                          max_lora_rank=args.max_lora_rank,
                          enable_tower_connector_lora=not args.disable_tower_connector_lora)
    if args.max_num_seqs:
        llm_kwargs['max_num_seqs'] = args.max_num_seqs
    if args.mm_processor_cache_gb is not None:
        llm_kwargs['mm_processor_cache_gb'] = args.mm_processor_cache_gb
    llm = LLM(**llm_kwargs)
    tok = llm.get_tokenizer(); sampling = SamplingParams(max_tokens=args.max_new_tokens, temperature=0)
    lora_request = None
    if not args.disable_lora:
        from vllm.lora.request import LoRARequest
        lora_request = LoRARequest(args.lora_name, 1, str(ckpt))
    rawf = raw_path.open('a' if existing else 'w')
    batch_size = max(1, args.batch_size)
    cache = VideoCache(args.video_cache_dir, args.video_cache_gib) if args.video_cache_dir else None
    decode_pool = ThreadPoolExecutor(max_workers=max(1, args.decode_workers))
    timings = {'queue_wait_seconds': 0.0, 'generate_seconds': 0.0}
    def build_requests(batch):
        def make_request(r):
            if r.get('_fixed80_absmeta'):
                video = fixed_absmeta_frames(Path(r['videos'][0]), (r.get('chat_template_kwargs') or {})['video_metadata'], args.nframes)
            elif 'derived_video' in r:
                src = Path(r['_fixed64_root']) / r['derived_video']
                video = fixed64_frames(src, r['output_fps'], args.nframes)
            else:
                src = Path(args.data_root) / ('heico-focus-vqa/videos' if r['_dataset'] == 'heico' else 'lapchole-focus-vqa/videos') / r['video']
                if cache: src = cache.get(src)
                video = frames(src, sec(r['timestamp_start']), sec(r['timestamp_end']), args.nframes)
            prompt = tok.apply_chat_template(make_messages(r), tokenize=False, add_generation_prompt=True, enable_thinking=False)
            return {'prompt': prompt, 'multi_modal_data': {'video': video}}
        return [decode_pool.submit(make_request, r) for r in batch]
    batches = [pending[i:i + batch_size] for i in range(0, len(pending), batch_size)]
    queued, next_batch = deque(), 0
    def fill_queue():
        nonlocal next_batch
        while next_batch < len(batches) and len(queued) < max(1, args.prefetch_batches):
            queued.append((batches[next_batch], build_requests(batches[next_batch]))); next_batch += 1
    fill_queue()
    processed = len(done)
    for batch_index in range(len(batches)):
        batch, futures = queued.popleft(); wait_start = time.perf_counter()
        requests = [future.result() for future in futures]
        timings['queue_wait_seconds'] += time.perf_counter() - wait_start
        fill_queue()
        t0 = time.perf_counter()
        if lora_request is None:
            results = llm.generate(requests, sampling)
        else:
            results = llm.generate(requests, sampling, lora_request=lora_request)
        elapsed = time.perf_counter() - t0
        timings['generate_seconds'] += elapsed
        for j, (r, result) in enumerate(zip(batch, results)):
            output = result.outputs[0]
            pred = output.text; latency = elapsed / len(batch)
            item = {**r, 'prediction': pred, 'latency': latency,
                    'finish_reason': getattr(output, 'finish_reason', None),
                    'generated_token_count': len(getattr(output, 'token_ids', []) or [])}
            if 'answer' in r:
                item['correct'] = correct(pred, r['answer'], row_answer_format(r))
            rawf.write(json.dumps(item, ensure_ascii=False) + '\n')
            processed += 1
            status = f' correct={item["correct"]}' if 'correct' in item else ''
            print(f'{processed}/{len(rows)} id={row_id(r)}{status} batch_latency={elapsed:.2f}s', flush=True)
        rawf.flush()
    decode_pool.shutdown(wait=True)
    rawf.close()
    all_items = read_jsonl(raw_path)
    all_items.sort(key=lambda item: row_by_id[str(item['id'])]['_input_index'])
    answers = [{'qID': str(row_id(item)), 'content': strip_prediction(item['prediction']),
                'latency': item['latency']} for item in all_items]
    (results_dir / 'answer.json').write_text(json.dumps(answers, ensure_ascii=False, indent=2))
    by_fmt = {}
    for item in all_items:
        if 'correct' in item:
            fmt = row_answer_format(item); by_fmt.setdefault(fmt, [0, 0]); by_fmt[fmt][0] += int(item['correct']); by_fmt[fmt][1] += 1
    metrics = {'count': len(answers), 'unique_ids': len({answer['qID'] for answer in answers})}
    if by_fmt:
        metrics.update({'correct': sum(v[0] for v in by_fmt.values()),
                        'accuracy': sum(v[0] for v in by_fmt.values()) / len(answers) if answers else 0,
                        'by_format': {k: {'correct': v[0], 'total': v[1], 'accuracy': v[0] / v[1]} for k, v in by_fmt.items()}})
    (results_dir / 'metrics.json').write_text(json.dumps(metrics, indent=2))
    runmeta.update(status='completed', finished_at=time.time(), metrics=metrics, timings=timings,
                   cache=None if not cache else {'hits': cache.hits, 'misses': cache.misses,
                                                  'copied_bytes': cache.bytes})
    (out / 'run.json').write_text(json.dumps(runmeta, indent=2))


if __name__ == '__main__':
    try:
        main()
    except BaseException as exc:
        if '--output' in sys.argv:
            status = Path(sys.argv[sys.argv.index('--output') + 1]) / 'run.json'
            if status.exists():
                meta = json.loads(status.read_text()); meta.update(status='failed', finished_at=time.time(), error=repr(exc))
                status.write_text(json.dumps(meta, indent=2))
        raise