#!/usr/bin/env python3
"""
Create a ≤500MB subset for Unsloth Studio upload limit.
Reads from unsloth_ready.jsonl (has conversation_type metadata).
Outputs clean format (messages only).
Strategy: Keep ALL curated + only v1 for generated variations.
"""

import json
import os

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

    total_in = 0
    total_out = 0
    curated_kept = 0
    v1_kept = 0
    dropped = 0

    print("Filtering 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)
                                total_in += 1

                                ctype = conv.get('conversation_type', '')
                                is_curated = ctype.startswith('curated_')
                                is_v1 = ctype.endswith('_v1')

                                if is_curated or is_v1:
                                    # Write clean format: only messages
                                    clean = {"messages": conv.get("messages", [])}
                                    outf.write(json.dumps(clean, ensure_ascii=False) + '\n')
                                    total_out += 1
                                    if is_curated:
                                        curated_kept += 1
                                    else:
                                        v1_kept += 1
                                else:
                                    dropped += 1

                                if total_in % 50000 == 0:
                                    print("  In: %s | Kept: %s | Dropped: %s" % (f"{total_in:,}", f"{total_out:,}", f"{dropped:,}"))
                            except:
                                pass
                            buffer = ""
                            in_obj = False
                    elif in_obj:
                        buffer += char

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

    print("\n" + "="*60)
    print("500MB SUBSET READY")
    print("="*60)
    print("Input conversations:  %s" % f"{total_in:,}")
    print("Output conversations: %s" % f"{total_out:,}")
    print("  Curated kept:       %s" % f"{curated_kept:,}")
    print("  Generated v1 kept:  %s" % f"{v1_kept:,}")
    print("  Dropped (v2-v5):    %s" % f"{dropped:,}")
    print("File: %s" % output_file)
    print("Size: %.1f MB" % size_mb)

if __name__ == '__main__':
    main()