"""
Generate comprehensive Gemma-4 training dataset for Ayurveda
Includes all question types: natural, reference, cross-reference, entity-based
Preserves all metadata, cross-references, chapter names
Output format: ShareGPT JSONL (compatible with Unsloth)
"""

import json
import re
from typing import Dict, List, Any
from dataclasses import dataclass, asdict
from pathlib import Path

# Section mappings with full Sanskrit names
SECTIONS = {
    "SU": "Sutrasthana (सूत्रस्थान)",
    "NI": "Nidanasthana (निदानस्थान)",
    "VI": "Vimanasthana (विमानस्थान)",
    "SHA": "Sharirasthana (शारीरस्थान)",
    "IND": "Indriyasthana (इन्द्रियस्थान)",
    "CI": "Cikitsasthana (चिकित्सास्थान)",
    "CHI": "Cikitsasthana (चिकित्सास्थान)",
    "KA": "Kalpasthana (कल्पस्थान)",
    "SI": "Siddhisthana (सिद्धिस्थान)",
    "SID": "Siddhisthana (सिद्धिस्थान)"
}

SECTION_DESCRIPTIONS = {
    "SU": "General Principles and Fundamentals",
    "NI": "Pathology and Disease Causation",
    "VI": "Quantitative and Measurement Science",
    "SHA": "Human Body and Anatomy",
    "IND": "Sensory Organs and Clinical Signs",
    "CI": "Treatment and Therapeutics",
    "CHI": "Treatment and Therapeutics",
    "KA": "Pharmaceutical Preparations",
    "SI": "Attainment of Perfection and Results",
    "SID": "Attainment of Perfection and Results"
}

# Hindi section names
SECTIONS_HI = {
    "SU": "सूत्रस्थान",
    "NI": "निदानस्थान",
    "VI": "विमानस्थान",
    "SHA": "शारीरस्थान",
    "IND": "इन्द्रियस्थान",
    "CI": "चिकित्सास्थान",
    "CHI": "चिकित्सास्थान",
    "KA": "कल्पस्थान",
    "SI": "सिद्धिस्थान",
    "SID": "सिद्धिस्थान"
}

@dataclass
class ShlokaData:
    id: int
    shloka_id: str
    sanskrit: str
    transliteration: str
    translation_en: str
    translation_hi: str
    section_code: str
    section_name: str
    section_name_hi: str
    section_description: str
    chapter: int
    verse: int
    chapter_name: str = ""  # Will be extracted or generated

class Gemma4DatasetGenerator:
    def __init__(self, source_file: str):
        self.source_file = source_file
        self.shlokas: List[ShlokaData] = []
        self.concept_graph: Dict[str, List[str]] = {}

    def load_source(self):
        """Parse source JSON and create sorted list"""
        with open(self.source_file, 'r', encoding='utf-8') as f:
            data = json.load(f)

        for item in data:
            # Parse shloka_id format: CH_SU_1/1
            match = re.match(r'CH_([A-Z]+)_([\d]+)/([\d]+)', item['shloka_id'])
            if not match:
                continue

            section_code, chapter, verse = match.groups()

            # Create structured data
            shloka = ShlokaData(
                id=item['id'],
                shloka_id=item['shloka_id'],
                sanskrit=item['sanskrit_shloka'],
                transliteration=item['transliteration'],
                translation_en=item.get('translations', {}).get('english', ''),
                translation_hi=item.get('translations', {}).get('hindi', ''),
                section_code=section_code,
                section_name=SECTIONS.get(section_code, section_code),
                section_name_hi=SECTIONS_HI.get(section_code, section_code),
                section_description=SECTION_DESCRIPTIONS.get(section_code, ''),
                chapter=int(chapter),
                verse=int(verse),
                chapter_name=self._generate_chapter_name(section_code, int(chapter))
            )
            self.shlokas.append(shloka)

        # Sort by shloka_id
        self.shlokas.sort(key=lambda x: (x.section_code, x.chapter, x.verse))
        print(f"Loaded {len(self.shlokas)} shlokas")

    def _generate_chapter_name(self, section: str, chapter: int) -> str:
        """Generate chapter name based on section and chapter number"""
        chapter_names = {
            "SU": {
                1: "Dīrghañjīvitīyamadhyāyaḥ (दीर्घञ्जीवितीयाध्यायः) - Longevity",
                2: "Aprāptīyamadhyāyaḥ (अप्राप्तीयाध्यायः)",
                3: "Ārūḍhivibhaktīyamadhyāyaḥ (आरूढिविभक्तीयाध्यायः)",
            },
            "NI": {
                1: "Triśārīrakam (त्रिशारीरकम्) - Three Bodies",
                2: "Jvaranidānam (ज्वरनिदानम्) - Fever Diagnosis",
            },
            "CI": {
                1: "Jvaracikitsitam (ज्वरचिकित्सितम्) - Fever Treatment",
                2: "Raktapittacikitsitam (रक्तपित्तचिकित्सितम्)",
            },
            "SHA": {
                1: "Katikāśarīram (कटिकाशरीरम्) - Pelvis and Body",
                2: "Śukraśonitāṣṭhīyamadhyāyaḥ (शुक्रशोणिताष्ट्हीयाध्यायः)",
            }
        }
        return chapter_names.get(section, {}).get(chapter, f"Chapter {chapter}")

    def _extract_key_concepts(self, shloka: ShlokaData) -> Dict[str, str]:
        """Extract key entities from shloka text"""
        concepts = {
            'speaker': 'Ātreya',  # Default for SU
            'topic': 'longevity',
            'action': 'expound',
            'section_theme': shloka.section_description
        }

        # Extract from Sanskrit
        if 'ātreya' in shloka.sanskrit.lower() or 'ātreya' in shloka.transliteration.lower():
            concepts['speaker'] = 'Ātreya'
        if 'agni' in shloka.sanskrit.lower():
            concepts['topic'] = 'agni'
            concepts['speaker'] = 'Agniveśa'
        if 'bhela' in shloka.sanskrit.lower():
            concepts['speaker'] = 'Bhela'

        return concepts

    def generate_natural_questions(self, shloka: ShlokaData) -> List[Dict]:
        """Generate natural conversation questions (no text references)"""
        conversations = []
        concepts = self._extract_key_concepts(shloka)

        # English Conceptual
        conversations.append({
            "conversation_type": "natural_conceptual_en",
            "messages": [
                {
                    "role": "user",
                    "content": f"What significance did {concepts['speaker']} attach to {concepts['topic']} when establishing his teachings on {concepts['section_theme']}?"
                },
                {
                    "role": "model",
                    "content": self._build_conceptual_answer(shloka, "en")
                }
            ]
        })

        # Hindi Conceptual
        conversations.append({
            "conversation_type": "natural_conceptual_hi",
            "messages": [
                {
                    "role": "user",
                    "content": f"{concepts['speaker']} ने {concepts['section_theme']} के उपदेश स्थापित करते समय {concepts['topic']} को क्या महत्त्व दिया?"
                },
                {
                    "role": "model",
                    "content": self._build_conceptual_answer(shloka, "hi")
                }
            ]
        })

        # Sanskrit Conceptual
        conversations.append({
            "conversation_type": "natural_conceptual_sa",
            "messages": [
                {
                    "role": "user",
                    "content": f"{concepts['speaker']} {concepts['section_theme']} उपदेशस्य स्थापने {concepts['topic']} का प्राधान्यं ददाति?"
                },
                {
                    "role": "model",
                    "content": self._build_conceptual_answer(shloka, "sa")
                }
            ]
        })

        # English Analytical
        conversations.append({
            "conversation_type": "natural_analytical_en",
            "messages": [
                {
                    "role": "user",
                    "content": f"How does {concepts['speaker']}'s emphasis on {concepts['topic']} in the {shloka.section_name} relate to the broader Ayurvedic understanding of {concepts['section_theme']}?"
                },
                {
                    "role": "model",
                    "content": self._build_analytical_answer(shloka, "en")
                }
            ]
        })

        # English Application
        conversations.append({
            "conversation_type": "natural_application_en",
            "messages": [
                {
                    "role": "user",
                    "content": f"A practitioner seeks to apply {concepts['speaker']}'s teachings on {concepts['topic']} in clinical practice. What fundamental principle from the {shloka.section_name} should guide this application?"
                },
                {
                    "role": "model",
                    "content": self._build_application_answer(shloka, "en")
                }
            ]
        })

        return conversations

    def generate_reference_questions(self, shloka: ShlokaData) -> List[Dict]:
        """Generate reference lookup questions (with text references)"""
        conversations = []

        # Direct Translation
        conversations.append({
            "conversation_type": "reference_translation",
            "messages": [
                {
                    "role": "user",
                    "content": f"Translate this Sanskrit text: {shloka.sanskrit}"
                },
                {
                    "role": "model",
                    "content": self._build_translation_answer(shloka)
                }
            ]
        })

        # ID Lookup
        conversations.append({
            "conversation_type": "reference_id_lookup",
            "messages": [
                {
                    "role": "user",
                    "content": f"What does {shloka.shloka_id} say?"
                },
                {
                    "role": "model",
                    "content": self._build_lookup_answer(shloka)
                }
            ]
        })

        # Complete Details
        conversations.append({
            "conversation_type": "reference_complete_details",
            "messages": [
                {
                    "role": "user",
                    "content": f"Give me complete details for {shloka.shloka_id}"
                },
                {
                    "role": "model",
                    "content": self._build_complete_answer(shloka)
                }
            ]
        })

        # Chapter Context
        conversations.append({
            "conversation_type": "reference_chapter_context",
            "messages": [
                {
                    "role": "user",
                    "content": f"Which chapter is {shloka.shloka_id} from?"
                },
                {
                    "role": "model",
                    "content": self._build_chapter_context_answer(shloka)
                }
            ]
        })

        # Multi-turn
        conversations.append({
            "conversation_type": "reference_multi_turn",
            "messages": [
                {
                    "role": "user",
                    "content": f"Explain this teaching from {shloka.shloka_id}"
                },
                {
                    "role": "model",
                    "content": self._build_lookup_answer(shloka)
                },
                {
                    "role": "user",
                    "content": "What is the Sanskrit formulation?"
                },
                {
                    "role": "model",
                    "content": f"The Sanskrit is: {shloka.sanskrit}"
                },
                {
                    "role": "user",
                    "content": "Where does this appear in the text?"
                },
                {
                    "role": "model",
                    "content": self._build_chapter_context_answer(shloka)
                }
            ]
        })

        return conversations

    def generate_cross_reference_questions(self, shloka: ShlokaData) -> List[Dict]:
        """Generate cross-reference questions"""
        conversations = []

        # Find related shlokas
        related = self._find_related_shlokas(shloka)

        if len(related) >= 1:
            # Intra-chapter connection
            conversations.append({
                "conversation_type": "cross_reference_intra_chapter",
                "messages": [
                    {
                        "role": "user",
                        "content": f"How does {shloka.shloka_id} connect to the broader teaching in {shloka.chapter_name}?"
                    },
                    {
                        "role": "model",
                        "content": self._build_cross_reference_answer(shloka, related)
                    }
                ],
                "references": [s.shloka_id for s in related[:3]]
            })

        return conversations

    def _find_related_shlokas(self, shloka: ShlokaData) -> List[ShlokaData]:
        """Find related shlokas for cross-referencing"""
        related = []

        # Same chapter
        for s in self.shlokas:
            if s.section_code == shloka.section_code and s.chapter == shloka.chapter:
                if s.verse != shloka.verse:
                    related.append(s)

        return related[:5]  # Limit to 5

    def _build_conceptual_answer(self, shloka: ShlokaData, lang: str) -> str:
        """Build conceptual understanding answer"""
        if lang == "en":
            return f"""{shloka.translation_en}

This teaching from **{shloka.section_name}** establishes the foundational importance of the subject matter. Located in **{shloka.chapter_name}**, Chapter {shloka.chapter}, this verse ({shloka.shloka_id}) contributes to the understanding of {shloka.section_description}.

**Sanskrit:** {shloka.sanskrit}

**Transliteration:** {shloka.transliteration}

**Full Reference:** {shloka.section_name} Chapter {shloka.chapter}, Verse {shloka.verse}"""
        elif lang == "hi":
            return f"""{shloka.translation_hi}

**{shloka.section_name_hi}** का यह उपदेश विषय के मूलभूत महत्त्व को स्थापित करता है। **{shloka.chapter_name}**, अध्याय {shloka.chapter} में स्थित, यह श्लोक ({shloka.shloka_id}) {shloka.section_description} की समझ में योगदान देता है।

**संस्कृत:** {shloka.sanskrit}

**पूर्ण सन्दर्भ:** {shloka.section_name_hi} अध्याय {shloka.chapter}, श्लोक {shloka.verse}"""
        else:  # Sanskrit
            return f"""{shloka.translation_en}

**{shloka.section_name}** एतत् उपदेशं {shloka.section_description} इति मूलभूतं महत्त्वं स्थापयति।

**मूलम्:** {shloka.sanskrit}

**सम्पूर्ण सन्दर्भः:** {shloka.section_name} अध्यायः {shloka.chapter}, श्लोकः {shloka.verse}"""

    def _build_analytical_answer(self, shloka: ShlokaData, lang: str) -> str:
        """Build analytical reasoning answer"""
        return f"""**Analysis of {shloka.shloka_id}:**

{shloka.translation_en}

This verse from **{shloka.section_name}** (dealing with {shloka.section_description}) presents a key doctrinal concept. The systematic placement in **{shloka.chapter_name}** suggests a pedagogical progression from foundational principles to specific applications.

**Key Insight:** The teaching connects to {shloka.section_description} by establishing the theoretical framework necessary for subsequent practical applications.

**Sanskrit Text:** {shloka.sanskrit}

**Full Citation:** {shloka.section_name} Chapter {shloka.chapter}, Verse {shloka.verse}"""

    def _build_application_answer(self, shloka: ShlokaData, lang: str) -> str:
        """Build practical application answer"""
        return f"""**Clinical Application of {shloka.shloka_id}:**

{shloka.translation_en}

**Practical Guidance:** This principle from **{shloka.section_name}** guides Ayurvedic practice by establishing that {shloka.section_description} must consider the broader context of health preservation.

**Implementation:** Practitioners apply this by ensuring that interventions align with the fundamental principles established in **{shloka.chapter_name}**.

**Sanskrit Source:** {shloka.sanskrit}

**Reference:** {shloka.shloka_id} in {shloka.section_name}, Chapter {shloka.chapter}"""

    def _build_translation_answer(self, shloka: ShlokaData) -> str:
        """Build translation answer with full metadata"""
        return f"""**Translation of {shloka.shloka_id}:**

**Sanskrit:**
{shloka.sanskrit}

**Transliteration (IAST):**
{shloka.transliteration}

**English Translation:**
{shloka.translation_en}

**Hindi Translation:**
{shloka.translation_hi}

**Source:** **{shloka.section_name}** - {shloka.section_description}
**Location:** Chapter {shloka.chapter} ({shloka.chapter_name}), Verse {shloka.verse}
**Full Citation:** {shloka.section_name} Chapter {shloka.chapter}, Verse {shloka.verse}"""

    def _build_lookup_answer(self, shloka: ShlokaData) -> str:
        """Build ID lookup answer"""
        return f"""**Complete text of {shloka.shloka_id}:**

**Sanskrit:**
{shloka.sanskrit}

**Transliteration:**
{shloka.transliteration}

**Translation:**
{shloka.translation_en}

**Source:** **{shloka.section_name}** ({shloka.section_name_hi})
**Section:** {shloka.section_description}
**Chapter:** {shloka.chapter} - {shloka.chapter_name}
**Verse:** {shloka.verse}

**Full Citation:** {shloka.section_name} Chapter {shloka.chapter}, Verse {shloka.verse}"""

    def _build_complete_answer(self, shloka: ShlokaData) -> str:
        """Build complete details answer"""
        return f"""**Complete Information for {shloka.shloka_id}:**

**Sanskrit (Devanagari):**
{shloka.sanskrit}

**Transliteration (IAST):**
{shloka.transliteration}

**English Translation:**
{shloka.translation_en}

**Hindi Translation:**
{shloka.translation_hi}

**Source Details:**
- Text: Charaka Samhita (चरक संहिता)
- Section: {shloka.section_name} ({shloka.section_name_hi})
- Section Code: {shloka.section_code}
- Section Description: {shloka.section_description}
- Chapter: {shloka.chapter}
- Chapter Name: {shloka.chapter_name}
- Verse: {shloka.verse}
- Full Citation: {shloka.section_name} Chapter {shloka.chapter}, Verse {shloka.verse}
- Short Reference: {shloka.shloka_id}"""

    def _build_chapter_context_answer(self, shloka: ShlokaData) -> str:
        """Build chapter context answer"""
        return f"""**{shloka.shloka_id}** is from **{shloka.section_name}** ({shloka.section_name_hi}), Chapter {shloka.chapter}, Verse {shloka.verse}.

**Section:** {shloka.section_name} ({shloka.section_code}) - {shloka.section_description}

**Chapter:** {shloka.chapter_name}

**Verse Text:**
{shloka.sanskrit}

**Translation:**
{shloka.translation_en}

**Full Reference:** {shloka.section_name} Chapter {shloka.chapter}, Verse {shloka.verse}"""

    def _build_cross_reference_answer(self, shloka: ShlokaData, related: List[ShlokaData]) -> str:
        """Build cross-reference answer"""
        refs = ", ".join([s.shloka_id for s in related[:3]])
        return f"""**{shloka.shloka_id}** in context:

{shloka.translation_en}

**Connection to Chapter {shloka.chapter}:** This verse in {shloka.chapter_name} relates to nearby teachings including {refs}. Together they form a comprehensive exposition on {shloka.section_description}.

**Sequential Understanding:**
1. Previous verses establish preliminary concepts
2. {shloka.shloka_id} introduces the specific teaching
3. Subsequent verses expand on applications

**Sanskrit:** {shloka.sanskrit}

**Reference:** {shloka.section_name} Chapter {shloka.chapter}, Verse {shloka.verse}"""

    def generate_advanced_questions(self, shloka: ShlokaData) -> List[Dict]:
        """Generate advanced question types"""
        conversations = []
        concepts = self._extract_key_concepts(shloka)

        # Clinical Case Study
        conversations.append({
            "conversation_type": "advanced_clinical_case",
            "messages": [
                {
                    "role": "user",
                    "content": f"A patient presents with symptoms related to {concepts['topic']}. Using {concepts['speaker']}'s framework from the {shloka.section_name}, how would you approach diagnosis and treatment?"
                },
                {
                    "role": "model",
                    "content": self._build_clinical_answer(shloka)
                }
            ]
        })

        # Word Analysis
        conversations.append({
            "conversation_type": "advanced_word_analysis",
            "messages": [
                {
                    "role": "user",
                    "content": f"Analyze the key terms in {shloka.shloka_id} from {shloka.section_name} and explain their technical significance in Ayurvedic medicine."
                },
                {
                    "role": "model",
                    "content": self._build_word_analysis_answer(shloka)
                }
            ]
        })

        # Historical Context
        conversations.append({
            "conversation_type": "advanced_historical_context",
            "messages": [
                {
                    "role": "user",
                    "content": f"How does {shloka.shloka_id} reflect the classical Ayurvedic approach to {concepts['section_theme']}, and what philosophical principles underlie this teaching?"
                },
                {
                    "role": "model",
                    "content": self._build_historical_answer(shloka)
                }
            ]
        })

        # Comparative Analysis
        conversations.append({
            "conversation_type": "advanced_comparative",
            "messages": [
                {
                    "role": "user",
                    "content": f"Compare the approach in {shloka.shloka_id} ({shloka.section_name}) with similar concepts in other Ayurvedic texts. What distinguishes Charaka's perspective?"
                },
                {
                    "role": "model",
                    "content": self._build_comparative_answer(shloka)
                }
            ]
        })

        return conversations

    def _build_clinical_answer(self, shloka: ShlokaData) -> str:
        return f"""**Clinical Application of {shloka.shloka_id}:**

{shloka.translation_en}

**Diagnostic Framework:**
This teaching from **{shloka.section_name}** ({shloka.section_description}) provides the theoretical basis for clinical assessment. The practitioner applies this by:

1. **Assessment Phase:** Understanding the patient's condition within the framework of {shloka.section_description}
2. **Analysis Phase:** Applying the principles from this verse to interpret signs and symptoms
3. **Treatment Planning:** Developing interventions that align with the verse's teaching

**Practical Implementation:**
- Evaluate patient constitution (prakṛti)
- Assess current pathological state (vikṛti)
- Apply therapeutic principles from {shloka.chapter_name}

**Sanskrit Reference:** {shloka.sanskrit}

**Source:** {shloka.section_name}, Chapter {shloka.chapter}, Verse {shloka.verse}"""

    def _build_word_analysis_answer(self, shloka: ShlokaData) -> str:
        return f"""**Technical Analysis of {shloka.shloka_id}:**

**Sanskrit Text:**
{shloka.sanskrit}

**Transliteration:**
{shloka.transliteration}

**Key Terms:**
The terminology in this verse from **{shloka.section_name}** reflects precise Ayurvedic technical language:

**Translation:**
{shloka.translation_en}

**Significance in {shloka.section_description}:**
These terms establish the conceptual framework for understanding {shloka.section_description}, demonstrating Charaka's systematic approach to medical knowledge.

**Reference:** {shloka.section_name} Chapter {shloka.chapter}, Verse {shloka.verse}"""

    def _build_historical_answer(self, shloka: ShlokaData) -> str:
        return f"""**Historical and Philosophical Context of {shloka.shloka_id}:**

**The Teaching:**
{shloka.translation_en}

**Philosophical Basis:**
This verse from **{shloka.section_name}** ({shloka.section_description}) reflects the classical Ayurvedic worldview that integrates:

1. **Empirical Observation:** Practical medical knowledge
2. **Theoretical Framework:** Systematic understanding of health and disease
3. **Applied Methodology:** Clinical implementation

**Classical Context:**
Located in **{shloka.chapter_name}** (Chapter {shloka.chapter}), this teaching represents the authoritative exposition on {shloka.section_description} as presented by Charaka.

**Sanskrit Original:** {shloka.sanskrit}

**Source:** {shloka.section_name}, Chapter {shloka.chapter}, Verse {shloka.verse}"""

    def _build_comparative_answer(self, shloka: ShlokaData) -> str:
        return f"""**Comparative Analysis of {shloka.shloka_id}:**

**Charaka's Approach:**
{shloka.translation_en}

**Distinctive Features:**
This verse from **{shloka.section_name}** demonstrates Charaka's systematic methodology for {shloka.section_description}. The approach is characterized by:

1. **Comprehensive Scope:** Addressing theory and practice
2. **Systematic Organization:** Logical progression of concepts
3. **Clinical Orientation:** Practical applicability

**Within Charaka Samhita:**
As part of **{shloka.chapter_name}**, this teaching connects to broader themes in Ayurvedic medicine while maintaining the text's distinctive emphasis on {shloka.section_description}.

**Technical Source:** {shloka.sanskrit}

**Location:** {shloka.section_name} Chapter {shloka.chapter}, Verse {shloka.verse}"""

    def generate_all_conversations(self) -> List[Dict]:
        """Generate complete dataset with all question types"""
        all_conversations = []

        for idx, shloka in enumerate(self.shlokas):
            if idx % 500 == 0:
                print(f"Processing {idx}/{len(self.shlokas)}...")

            # Generate all question types
            natural = self.generate_natural_questions(shloka)
            reference = self.generate_reference_questions(shloka)
            cross_ref = self.generate_cross_reference_questions(shloka)
            advanced = self.generate_advanced_questions(shloka)

            # Combine with metadata
            for conv in natural + reference + cross_ref + advanced:
                conversation = {
                    "shloka_id": shloka.shloka_id,
                    "conversation_type": conv["conversation_type"],
                    "messages": conv["messages"],
                    "metadata": {
                        "sanskrit": shloka.sanskrit,
                        "transliteration": shloka.transliteration,
                        "section": shloka.section_name,
                        "section_hi": shloka.section_name_hi,
                        "section_code": shloka.section_code,
                        "section_description": shloka.section_description,
                        "chapter": shloka.chapter,
                        "chapter_name": shloka.chapter_name,
                        "verse": shloka.verse,
                        "translation_en": shloka.translation_en,
                        "translation_hi": shloka.translation_hi
                    }
                }

                # Add references if present
                if "references" in conv:
                    conversation["references"] = conv["references"]

                all_conversations.append(conversation)

        return all_conversations

    def save_dataset(self, output_file: str, conversations: List[Dict]):
        """Save in ShareGPT JSONL format"""
        with open(output_file, 'w', encoding='utf-8') as f:
            for conv in conversations:
                f.write(json.dumps(conv, ensure_ascii=False, indent=2) + '\n')

        print(f"Saved {len(conversations)} conversations to {output_file}")

        # Save statistics
        stats = {}
        for conv in conversations:
            ctype = conv["conversation_type"]
            stats[ctype] = stats.get(ctype, 0) + 1

        print("\nConversation type distribution:")
        for ctype, count in sorted(stats.items(), key=lambda x: -x[1])[:10]:
            print(f"  {ctype}: {count}")

    def save_alpaca_format(self, output_file: str, conversations: List[Dict]):
        """Save in Alpaca CSV format for training"""
        import csv

        rows = []
        for conv in conversations:
            messages = conv["messages"]
            if len(messages) >= 2:
                # Extract user and assistant messages
                user_msg = ""
                assistant_msg = ""

                for msg in messages:
                    if msg["role"] == "user":
                        user_msg = msg["content"]
                    elif msg["role"] == "model":
                        assistant_msg = msg["content"]

                if user_msg and assistant_msg:
                    rows.append({
                        "instruction": user_msg,
                        "input": "",
                        "output": assistant_msg,
                        "shloka_id": conv["shloka_id"],
                        "conversation_type": conv["conversation_type"],
                        "sanskrit": conv["metadata"]["sanskrit"],
                        "transliteration": conv["metadata"]["transliteration"]
                    })

        with open(output_file, 'w', encoding='utf-8', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=["instruction", "input", "output", "shloka_id", "conversation_type", "sanskrit", "transliteration"])
            writer.writeheader()
            writer.writerows(rows)

        print(f"Saved {len(rows)} training rows to {output_file}")


def main():
    print("Initializing Gemma-4 Dataset Generator...")
    generator = Gemma4DatasetGenerator("shlokas_export_2026-04-30.json")

    print("Loading source data...")
    generator.load_source()

    print("Generating conversations...")
    conversations = generator.generate_all_conversations()

    print("\nSaving ShareGPT format...")
    generator.save_dataset("gemma4_ayurveda_complete.jsonl", conversations)

    print("\nSaving Alpaca CSV format...")
    generator.save_alpaca_format("gemma4_ayurveda_training.csv", conversations)

    print("\nDone!")
    print(f"Total conversations: {len(conversations)}")
    print(f"Estimated size: ~{len(conversations) * 500 / 1024 / 1024:.1f} MB")


if __name__ == "__main__":
    main()