#!/usr/bin/env python3
import json

# 1. Load 10 shlokas from sorted.json
with open('shlokas_export_2026-05-14_SORTED.json', 'r', encoding='utf-8') as f:
    data = json.load(f)

indices = [0, 50, 100, 500, 1000, 1500, 2000, 2500, 3000, 3575]
selected = [data[i] for i in indices if i < len(data)]
target_ids = {s['shloka_id'] for s in selected}

# 2. Stream combined JSONL once, find first generated conv per target sid
gen_samples = {}
with open('gemma4_ayurveda_final_combined.jsonl', 'r', encoding='utf-8') as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        try:
            conv = json.loads(line)
            sid = conv.get('shloka_id')
            ctype = conv.get('conversation_type', '')
            if sid in target_ids and not ctype.startswith('curated_') and sid not in gen_samples:
                gen_samples[sid] = conv
                if len(gen_samples) == len(target_ids):
                    break
        except:
            pass

# 3. Build output
output = []
for shloka in selected:
    sid = shloka['shloka_id']
    gen_conv = gen_samples.get(sid)

    entry = {
        "shloka_id": sid,
        "sanskrit": shloka.get('sanskrit', ''),
        "transliteration": shloka.get('transliteration', ''),
        "translation_english": shloka.get('translation_english', ''),
        "translation_hindi": shloka.get('translation_hindi', ''),
        "section": shloka.get('source', {}).get('section', ''),
        "chapter": shloka.get('source', {}).get('chapter', ''),
        "verse": shloka.get('source', {}).get('verse', ''),
        "total_curated_questions": len(shloka.get('questions', [])),
        "curated_questions_sample": [],
        "generated_conversation_sample": gen_conv
    }

    for q in shloka.get('questions', [])[:3]:
        q_sample = {}
        for lang in ['en', 'hi', 'sa']:
            qd = q.get(lang, {})
            if qd.get('q'):
                q_sample[lang] = {
                    "question": qd['q'],
                    "answer_preview": (qd.get('a', '') or '')[:200]
                }
        if q_sample:
            entry["curated_questions_sample"].append(q_sample)

    output.append(entry)

with open('sample_10_shlokas.json', 'w', encoding='utf-8') as f:
    json.dump(output, f, ensure_ascii=False, indent=2)

print("Done: sample_10_shlokas.json with %d shlokas" % len(output))