#!/usr/bin/env python3
import json

with open('shlokas_export_2026-05-14_SORTED.json', 'r', encoding='utf-8') as f:
    data = json.load(f)

# Pick 10 diverse shlokas with rich content
indices = [0, 50, 100, 500, 1000, 1500, 2000, 2500, 3000, 3575]
selected = [data[i] for i in indices if i < len(data)]

# Also fetch one generated conversation per shloka from the combined dataset
def find_generated(sid):
    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)
                if conv.get('shloka_id') == sid and not conv.get('conversation_type', '').startswith('curated_'):
                    return conv
            except:
                pass
    return None

output = []
for shloka in selected:
    sid = shloka['shloka_id']
    gen_conv = find_generated(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
    }

    # Add up to 3 curated questions across languages
    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("Wrote sample_10_shlokas.json with %d shlokas" % len(output))