#!/usr/bin/env python3
"""
Strip metadata from Unsloth-ready dataset.
Outputs: gemma4_ayurveda_unsloth_clean.jsonl
Only keeps the 'messages' field that Unsloth trains on.
"""

import json
import os

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

    total = 0

    print("Stripping metadata from %s ..." % input_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)
                                # Keep only messages
                                clean = {"messages": conv.get("messages", [])}
                                outf.write(json.dumps(clean, ensure_ascii=False) + '\n')
                                total += 1
                                if total % 50000 == 0:
                                    print("  Processed %s" % f"{total:,}")
                            except:
                                pass
                            buffer = ""
                            in_obj = False
                    elif in_obj:
                        buffer += char

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

    print("\n" + "="*60)
    print("CLEANED DATASET READY")
    print("="*60)
    print("Total conversations: %s" % f"{total:,}")
    print("Original size: %.1f MB" % orig_mb)
    print("Cleaned size: %.1f MB" % size_mb)
    print("Reduction: %.1f%%" % ((1 - size_mb/orig_mb) * 100))
    print("File: %s" % output_file)

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