#!/usr/bin/env python3
"""
Inject correct shloka_id references into assistant answers.
Reads from gemma4_ayurveda_unsloth_ready.jsonl (has metadata + shloka_id).
Outputs: gemma4_ayurveda_unsloth_clean_WITH_REFS.jsonl
Only keeps messages field, but injects correct ID into every assistant answer.
"""

import json
import os
import re

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 inject_reference(conv):
    """Inject correct shloka_id into assistant answer"""
    sid = conv.get('shloka_id', '')
    if not sid:
        return None

    messages = conv.get("messages", [])

    for msg in messages:
        if msg.get('role') == 'assistant':
            content = msg.get('content', '')
            if not content:
                continue

            # Check if already has Reference ID
            if 'Reference ID:' in content and sid in content:
                # Already correct
                pass
            elif 'Reference ID:' in content:
                # Has a wrong/old ID, replace it
                content = re.sub(
                    r'\*\*Reference ID:\*\*\s*CH_[A-Z]+_[\d.]+/(\S+)',
                    f'**Reference ID:** {sid}',
                    content
                )
                msg['content'] = content
            else:
                # No Reference ID yet - inject at the end
                content = content.rstrip()
                content += f"\n\n**Reference ID:** {sid}"
                msg['content'] = content

    return {"messages": messages}

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

    total = 0
    injected = 0
    already_ok = 0

    print(f"Injecting references from {input_file} ...")

    with open(output_file, 'w', encoding='utf-8') as outf:
        for conv in stream_jsonl(input_file):
            sid = conv.get('shloka_id', '')
            if not sid:
                total += 1
                continue

            clean = inject_reference(conv)
            if clean:
                outf.write(json.dumps(clean, ensure_ascii=False) + '\n')
                total += 1

                # Check what we did
                content = clean['messages'][-1].get('content', '')
                if f'**Reference ID:** {sid}' in content:
                    if 'Reference ID:' in content and content.count('Reference ID:') == 1:
                        injected += 1
                    else:
                        already_ok += 1

                if total % 50000 == 0:
                    print(f"  Processed {total:,} | Injected: {injected:,} | Already OK: {already_ok:,}")

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

    print(f"\n{'='*60}")
    print("CLEAN DATASET WITH REFS READY")
    print(f"{'='*60}")
    print(f"Total conversations: {total:,}")
    print(f"Injected refs:      {injected:,}")
    print(f"Already had refs:   {already_ok:,}")
    print(f"File: {output_file}")
    print(f"Size: {size_mb:.1f} MB")

    # Verify 5 random samples
    print(f"\n{'='*60}")
    print("VERIFICATION (5 random samples)")
    print(f"{'='*60}")
    with open(output_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()
        import random
        random.seed(42)
        for idx in random.sample(range(len(lines)), 5):
            obj = json.loads(lines[idx])
            msgs = obj.get('messages', [])
            for m in msgs:
                if m.get('role') == 'assistant':
                    content = m.get('content', '')
                    ref_lines = [l for l in content.split('\n') if 'Reference ID:' in l]
                    if ref_lines:
                        print(f"  Line {idx}: {ref_lines[0]}")
                    else:
                        print(f"  Line {idx}: NO REF FOUND!")
                    break

if __name__ == '__main__':
    main()