"""Summarize context-benchmark records, keeping first attempts separate from retries."""
import argparse
import json
import statistics
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument('files', nargs='+', type=Path)
args = parser.parse_args()
for path in args.files:
    rows = []
    for line in path.read_text().splitlines():
        if line.startswith('{'):
            rows.append(json.loads(line))
    summary = {'file': path.name,
               'completed': any(r.get('event') == 'CONTEXT_SUITE_COMPLETED' for r in rows),
               'note': 'Request completion is not a host-stability verdict; inspect resource and kernel OOM evidence.',
               'stages': []}
    for size in [65536, 131072, 260000]:
        results = [r for r in rows if r.get('event') == 'RESULT' and r['target_input_tokens'] == size]
        first = [r for r in results if r['attempt'] == 1]
        if not first:
            if any(r.get('event') == 'STAGE_START' and r['target_input_tokens'] == size for r in rows):
                summary['stages'].append({'target_input_tokens': size, 'first_attempts_completed': 0,
                                          'status': 'started_without_completed_response'})
            continue
        cold = [r for r in first if r['case'] == 'front']
        warm = [r for r in first if r['case'] != 'front']
        summary['stages'].append({
            'target_input_tokens': size, 'first_attempts': len(first),
            'correct_first_attempts': sum(r['correct'] for r in first),
            'truncated_first_attempts': sum(r['finish_reason'] == 'length' for r in first),
            'retries': len(results)-len(first),
            'actual_input_range': [min(r['usage']['prompt_tokens'] for r in first), max(r['usage']['prompt_tokens'] for r in first)],
            'first_document_ttft_s': [r['ttft_s'] for r in cold],
            'reused_prefix_ttft_median_s': statistics.median(r['ttft_s'] for r in warm) if warm else None,
            'first_document_elapsed_s': [r['elapsed_s'] for r in cold],
            'reused_prefix_elapsed_median_s': statistics.median(r['elapsed_s'] for r in warm) if warm else None,
            'max_completion_tokens': max(r['usage']['completion_tokens'] for r in first),
            'max_reasoning_tokens': max(r['usage'].get('completion_tokens_details', {}).get('reasoning_tokens', 0) for r in first),
            'prefix_hit_tokens': sum(r['metrics_delta']['vllm:prefix_cache_hits_total'] for r in results),
            'preemptions': sum(r['metrics_delta']['vllm:num_preemptions_total'] for r in results),
        })
    print(json.dumps(summary, ensure_ascii=False, indent=2))