#!/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)

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

section_names = {
    'SU': 'Sutrasthana', 'NI': 'Nidanasthana', 'VI': 'Vimanasthana',
    'SHA': 'Sharirasthana', 'IN': 'Indriyasthana', 'CHI': 'Cikitsasthana',
    'KAL': 'Kalpasthana', 'SID': 'Siddhisthana'
}

def build_system_prompt(lang, section):
    if lang == "hi":
        return f"You are an expert Ayurvedic scholar fluent in Hindi. Answer based on Charaka Samhita {section} with proper citations."
    elif lang == "sa":
        return f"You are an expert Ayurvedic scholar fluent in Sanskrit. Answer based on Charaka Samhita {section} with classical citations."
    return f"You are an expert Ayurvedic scholar trained in Charaka Samhita. Provide accurate responses with precise citations from {section}."

def build_answer(metadata, lang):
    en = metadata.get('translation_en', '')
    hi = metadata.get('translation_hi', '')
    sanskrit = metadata.get('sanskrit', '')
    trans = metadata.get('transliteration', '')
    section = metadata.get('section', '')
    chapter = metadata.get('chapter', '')
    verse = metadata.get('verse', '')

    if lang == 'hi':
        ans = hi or en
        if sanskrit:
            ans += f"\n\n**संस्कृत:** {sanskrit}"
        ans += f"\n\n**सन्दर्भ:** {section}, अध्याय {chapter}, श्लोक {verse}"
    elif lang == 'sa':
        ans = en
        if sanskrit:
            ans += f"\n\n**मूलम्:** {sanskrit}"
        ans += f"\n\n**सन्दर्भः:** {section}, अध्यायः {chapter}, श्लोकः {verse}"
    else:
        ans = en
        if sanskrit:
            ans += f"\n\n**Sanskrit:** {sanskrit}"
        if trans:
            ans += f"\n**Transliteration:** {trans}"
        ans += f"\n\n**Reference:** {section}, Chapter {chapter}, Verse {verse}"
    return ans

output = []
for shloka in selected:
    sid = shloka['shloka_id']
    sec_code = sid.split('_')[1] if '_' in sid else 'SU'
    section_name = section_names.get(sec_code, 'Charaka Samhita')

    meta = {
        'sanskrit': shloka.get('sanskrit', ''),
        'transliteration': shloka.get('transliteration', ''),
        'translation_en': shloka.get('translation_english', ''),
        'translation_hi': shloka.get('translation_hindi', ''),
        'section': section_name,
        'chapter': shloka.get('source', {}).get('chapter', ''),
        'verse': shloka.get('source', {}).get('verse', '')
    }

    # Build sample generated conversations (2 types x 3 langs)
    generated_samples = []
    for ctype_base in ['natural_conceptual', 'reference_translation']:
        for lang in ['en', 'hi', 'sa']:
            conv_type = f"{ctype_base}_{lang}_v1"
            if ctype_base == 'natural_conceptual':
                topic = meta['translation_en'][:40] if meta['translation_en'] else 'this teaching'
                if lang == 'en':
                    q = f"What significance did the Ayurvedic sage attach to {topic} when establishing the framework for {section_name}?"
                elif lang == 'hi':
                    q = f"आयुर्वेदिक ऋषि ने {section_name} में {topic} के विषय में क्या महत्व स्थापित किया?"
                else:
                    q = f"आयुर्वेदर्षिः {section_name} {topic} प्राधान्यं ददाति?"
            else:
                if lang == 'en':
                    q = f"Translate this: {meta['sanskrit'][:80]}"
                elif lang == 'hi':
                    q = f"इसका अनुवाद करें: {meta['sanskrit'][:80]}"
                else:
                    q = f"भाषान्तरितं कुर्वन्तु: {meta['sanskrit'][:80]}"

            a = build_answer(meta, lang)
            generated_samples.append({
                "conversation_type": conv_type,
                "messages": [
                    {"role": "system", "content": build_system_prompt(lang, section_name)},
                    {"role": "user", "content": q},
                    {"role": "model", "content": a}
                ]
            })

    # Curated questions sample
    curated_samples = []
    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 '')[:250]
                }
        if q_sample:
            curated_samples.append(q_sample)

    entry = {
        "shloka_id": sid,
        "sanskrit": meta['sanskrit'],
        "transliteration": meta['transliteration'],
        "translation_english": meta['translation_en'],
        "translation_hindi": meta['translation_hi'],
        "section": section_name,
        "section_code": sec_code,
        "chapter": meta['chapter'],
        "verse": meta['verse'],
        "total_curated_questions": len(shloka.get('questions', [])),
        "curated_questions_sample": curated_samples,
        "generated_conversation_samples": generated_samples
    }
    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")