#!/usr/bin/env python3
"""
Convert combined dataset to Unsloth-compatible ShareGPT format.
Changes role: model -> assistant.
Verifies reference accuracy in metadata.
Outputs: gemma4_ayurveda_unsloth_trainable.jsonl
"""

import json
import os
from pathlib import Path

def stream_jsonl(filepath):
    """Stream JSONL 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)
            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 fix_conversation(conv):
    """Fix roles and verify metadata"""
    messages = conv.get('messages', [])
    for msg in messages:
        if msg.get('role') == 'model':
            msg['role'] = 'assistant'

    # Ensure metadata has accurate reference fields
    metadata = conv.get('metadata', {})
    sid = conv.get('shloka_id', '')

    # If metadata is empty but we have shloka_id, try to populate basic fields
    if not metadata and sid:
        parts = sid.split('_')
        if len(parts) >= 3:
            metadata = {
                'shloka_id': sid,
                'section_code': parts[1] if len(parts) > 1 else '',
            }
        conv['metadata'] = metadata

    return conv

def main():
    input_file = "gemma4_ayurveda_final_combined.jsonl"
    output_file = "gemma4_ayurveda_unsloth_trainable.jsonl"

    total = 0
    fixed_roles = 0

    print("Streaming and fixing %s ..." % input_file)

    with open(output_file, 'w', encoding='utf-8') as outf:
        for conv in stream_jsonl(input_file):
            fixed = fix_conversation(conv)

            # Count fixes
            for msg in fixed.get('messages', []):
                if msg.get('role') == 'assistant':
                    fixed_roles += 1
                    break  # Only count once per conversation

            outf.write(json.dumps(fixed, ensure_ascii=False) + '\n')
            total += 1

            if total % 50000 == 0:
                print("  Processed %s..." % f"{total:,}")

    size_mb = os.path.getsize(output_file) / (1024 * 1024)

    print("\n" + "="*60)
    print("UNSLOTH DATASET READY")
    print("="*60)
    print("Total conversations: %s" % f"{total:,}")
    print("Fixed model->assistant: %s conversations" % f"{fixed_roles:,}")
    print("File: %s" % output_file)
    print("Size: %.1f MB" % size_mb)

    # Sample verification
    print("\nSample verification (first conversation):")
    with open(output_file, 'r', encoding='utf-8') as f:
        first = json.loads(f.readline())
        print("  shloka_id: %s" % first.get('shloka_id'))
        print("  conversation_type: %s" % first.get('conversation_type'))
        for i, msg in enumerate(first.get('messages', [])[:3]):
            print("  msg[%d] role=%s -> %s..." % (i, msg.get('role'), msg.get('content','')[:50]))
        print("  metadata keys: %s" % list(first.get('metadata', {}).keys()))

if __name__ == "__main__":
    main()