#!/usr/bin/env python3
"""
Convert combined dataset to Unsloth-compatible ShareGPT format.
Changes role: model -> assistant.
Outputs: gemma4_ayurveda_unsloth_ready.jsonl
"""

import json
import os

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

    total = 0
    errors = 0

    print("Fixing %s -> %s ..." % (input_file, output_file))

    with open(input_file, 'r', encoding='utf-8') as inf:
        with open(output_file, 'w', encoding='utf-8') as outf:
            buffer = ""
            brace = 0
            in_obj = False

            for raw_line in inf:
                for char in raw_line:
                    if char == '{':
                        if brace == 0:
                            in_obj = True
                            buffer = "{"
                        else:
                            buffer += char
                        brace += 1
                    elif char == '}':
                        brace -= 1
                        buffer += char
                        if brace == 0 and in_obj:
                            try:
                                conv = json.loads(buffer)
                                for msg in conv.get('messages', []):
                                    if msg.get('role') == 'model':
                                        msg['role'] = 'assistant'
                                outf.write(json.dumps(conv, ensure_ascii=False) + '\n')
                                total += 1
                                if total % 50000 == 0:
                                    print("  Processed %s" % f"{total:,}")
                            except:
                                errors += 1
                            buffer = ""
                            in_obj = False
                    elif in_obj:
                        buffer += char

    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("Parse errors: %d" % errors)
    print("File: %s" % output_file)
    print("Size: %.1f MB" % size_mb)

    print("\nSample verification:")
    with open(output_file, 'r', encoding='utf-8') as f:
        first = json.loads(f.readline())
        print("  shloka_id: %s" % first.get('shloka_id'))
        print("  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]))

if __name__ == "__main__":
    main()