#!/usr/bin/env python3
"""
Merge curated questions from sorted.json into 56variations dataset.
Streams large files to avoid memory issues.
"""

import json
import re
from pathlib import Path
from collections import defaultdict

def load_sorted_questions(filepath):
    """Load sorted.json and build map of shloka_id -> questions"""
    with open(filepath, 'r', encoding='utf-8') as f:
        data = json.load(f)

    question_map = {}
    for item in data:
        sid = item['shloka_id']
        questions = item.get('questions', [])
        if questions:
            question_map[sid] = {
                'questions': questions,
                'metadata': {
                    'sanskrit': item.get('sanskrit', ''),
                    'transliteration': item.get('transliteration', ''),
                    'translation_en': item.get('translation_english', ''),
                    'translation_hi': item.get('translation_hindi', ''),
                    'section': item.get('source', {}).get('section', ''),
                    'section_full': item.get('source', {}).get('text_name', ''),
                    'chapter': item.get('source', {}).get('chapter', ''),
                    'verse': item.get('source', {}).get('verse', '')
                }
            }

    return question_map

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

def build_answer(metadata, lang):
    """Build answer with metadata"""
    en_text = metadata.get('translation_en', '')
    hi_text = metadata.get('translation_hi', '')
    sanskrit = metadata.get('sanskrit', '')
    trans = metadata.get('transliteration', '')
    section = metadata.get('section_full', '')
    chapter = metadata.get('chapter', '')
    verse = metadata.get('verse', '')

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

    return answer

def stream_jsonl(filepath):
    """Stream JSONL file one object at a time"""
    buffer = ""
    brace_count = 0
    in_object = False

    with open(filepath, 'r', encoding='utf-8') as f:
        while True:
            chunk = f.read(8192)  # Read 8KB at a time
            if not chunk:
                break

            for char in chunk:
                if char == '{':
                    if brace_count == 0:
                        in_object = True
                        buffer = "{"
                    else:
                        buffer += char
                    brace_count += 1
                elif char == '}':
                    brace_count -= 1
                    buffer += char
                    if brace_count == 0 and in_object:
                        try:
                            yield json.loads(buffer)
                        except:
                            pass
                        buffer = ""
                        in_object = False
                elif in_object:
                    buffer += char

def merge_curated_questions(variations_file, sorted_file, output_file):
    """Merge curated questions into variations dataset by streaming"""

    print(f"Loading questions from {sorted_file}...")
    question_map = load_sorted_questions(sorted_file)
    print(f"Loaded {len(question_map)} shlokas with curated questions")

    # Count existing conversations
    print(f"Counting existing conversations in {variations_file}...")
    existing_count = 0
    sample_metadata = None
    sample_sid = None

    for conv in stream_jsonl(variations_file):
        existing_count += 1
        if sample_metadata is None:
            sample_metadata = conv.get('metadata', {})
            sample_sid = conv.get('shloka_id', '')
        if existing_count % 50000 == 0:
            print(f"  Counted {existing_count}...")

    print(f"Total existing: {existing_count}")

    # Generate new conversations for each shloka
    print(f"\nGenerating curated conversations...")
    added = 0
    skipped = 0

    # Pre-build new conversations
    new_conversations = []
    for sid, info in question_map.items():
        questions = info['questions']
        metadata = info['metadata']

        for qidx, q in enumerate(questions, 1):
            for lang in ['en', 'hi', 'sa']:
                q_data = q.get(lang, {})
                if not q_data.get('q'):
                    continue

                conv_type = f"curated_q{qidx}_{lang}"

                new_conv = {
                    "shloka_id": sid,
                    "conversation_type": conv_type,
                    "messages": [
                        {"role": "system", "content": build_system_prompt(lang, metadata)},
                        {"role": "user", "content": q_data['q']},
                        {"role": "model", "content": q_data.get('a', '') or build_answer(metadata, lang)}
                    ],
                    "metadata": sample_metadata  # Use sample metadata structure
                }

                new_conversations.append(new_conv)
                added += 1

    print(f"Added {added} curated conversations")

    # Stream original + append new
    print(f"\nWriting to {output_file}...")
    with open(output_file, 'w', encoding='utf-8') as f:
        # Copy existing
        copied = 0
        for conv in stream_jsonl(variations_file):
            f.write(json.dumps(conv, ensure_ascii=False, indent=2) + '\n')
            copied += 1
            if copied % 50000 == 0:
                print(f"  Copied {copied}...")

        # Append new
        for conv in new_conversations:
            f.write(json.dumps(conv, ensure_ascii=False, indent=2) + '\n')

    import os
    size_mb = os.path.getsize(output_file) / (1024 * 1024)
    total = copied + len(new_conversations)
    print(f"\nDone!")
    print(f"Existing: {copied}")
    print(f"Added: {len(new_conversations)}")
    print(f"Total: {total}")
    print(f"File size: {size_mb:.1f} MB")

if __name__ == "__main__":
    merge_curated_questions(
        "gemma4_ayurveda_56variations.jsonl",
        "shlokas_export_2026-05-14_SORTED.json",
        "gemma4_ayurveda_56variations_enriched.jsonl"
    )