#!/usr/bin/env python3
"""
Create a ≤500MB subset for Unsloth Studio.
- Keeps ALL curated + only v1 variations
- Injects explicit shloka_id into assistant answers
- Outputs clean ShareGPT format
"""

import json
import os
import re

def append_shloka_id_to_answer(assistant_content, shloka_id):
    """Append Reference ID if not already present."""
    if not assistant_content:
        return assistant_content
    # Check if already has Reference ID
    if "Reference ID:" in assistant_content or "Reference ID" in assistant_content:
        # Ensure it matches the correct shloka_id
        if shloka_id in assistant_content:
            return assistant_content
        # Replace incorrect or generic one
        lines = assistant_content.split('\n')
        new_lines = []
        for line in lines:
            if "Reference ID:" in line:
                new_lines.append("**Reference ID:** %s" % shloka_id)
            else:
                new_lines.append(line)
        return '\n'.join(new_lines)

    # Append at end
    return assistant_content + "\n\n**Reference ID:** %s" % shloka_id

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

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

    print("Building 500MB subset with shloka_id references...")

    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', '')
                                sid = conv.get('shloka_id', '')
                                is_curated = ctype.startswith('curated_')
                                is_v1 = ctype.endswith('_v1')

                                if is_curated or is_v1:
                                    messages = conv.get("messages", [])
                                    # Inject shloka_id into assistant answer
                                    for msg in messages:
                                        if msg.get('role') == 'assistant' and sid:
                                            msg['content'] = append_shloka_id_to_answer(msg.get('content', ''), sid)

                                    # Write clean format
                                    clean = {"messages": 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 Exception as e:
                                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 WITH REFS 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)

    # Verify sample
    print("\nSample verification:")
    with open(output_file, 'r', encoding='utf-8') as f:
        for _ in range(3):
            line = f.readline()
            if not line:
                break
            obj = json.loads(line)
            msgs = obj.get('messages', [])
            for msg in msgs:
                if msg.get('role') == 'assistant':
                    content = msg.get('content', '')
                    ref_line = [l for l in content.split('\n') if 'Reference ID:' in l]
                    print("  Ref line: %s" % (ref_line[0] if ref_line else "MISSING"))
                    break
            else:
                print("  No assistant msg")

if __name__ == '__main__':
    main()