import json, sys

filepath = r'E:\dataset\gemma4_ayurveda_final_combined.jsonl'
outpath = r'E:\dataset\analyze_jsonl.out'
chunk_size = 8192

buffer = ''
brace_count = 0
in_string = False
escape_next = False

conv_type_counts = {}
unique_shlokas = set()
total_convs = 0
has_curated = False
has_variation = False

sample_curated = None
sample_variation = None

with open(filepath, 'r', encoding='utf-8') as f:
    while True:
        chunk = f.read(chunk_size)
        if not chunk:
            break
        for char in chunk:
            if in_string:
                if escape_next:
                    escape_next = False
                elif char == '\\':
                    escape_next = True
                elif char == '"':
                    in_string = False
            else:
                if char == '"':
                    in_string = True
                elif char == '{':
                    if brace_count == 0:
                        buffer = '{'
                    else:
                        buffer += char
                    brace_count += 1
                    continue
                elif char == '}':
                    brace_count -= 1
                    buffer += char
                    if brace_count == 0:
                        total_convs += 1
                        try:
                            obj = json.loads(buffer)
                            ct = obj.get('conversation_type', 'UNKNOWN')
                            sid = obj.get('shloka_id', '')
                            if sid:
                                unique_shlokas.add(sid)
                            conv_type_counts[ct] = conv_type_counts.get(ct, 0) + 1
                            if 'curated_' in ct:
                                has_curated = True
                                if sample_curated is None:
                                    sample_curated = obj
                            if 'natural_' in ct or 'reference_' in ct or 'v5' in ct:
                                has_variation = True
                                if sample_variation is None:
                                    sample_variation = obj
                        except Exception:
                            pass
                        buffer = ''
                    continue
            if brace_count > 0:
                buffer += char

if buffer.strip():
    try:
        obj = json.loads(buffer)
        ct = obj.get('conversation_type', 'UNKNOWN')
        sid = obj.get('shloka_id', '')
        if sid:
            unique_shlokas.add(sid)
        conv_type_counts[ct] = conv_type_counts.get(ct, 0) + 1
        total_convs += 1
    except Exception:
        pass

with open(outpath, 'w', encoding='utf-8') as out:
    out.write('=== Report ===\n')
    out.write(f'Total conversations: {total_convs}\n')
    out.write(f'Unique shloka_ids: {len(unique_shlokas)}\n')
    out.write(f'Has curated types: {has_curated}\n')
    out.write(f'Has variation types: {has_variation}\n')

    sorted_types = sorted(conv_type_counts.items(), key=lambda x: x[1], reverse=True)[:20]
    out.write('\nTop 20 conversation_type counts:\n')
    for ct, cnt in sorted_types:
        out.write(f'  {ct}: {cnt}\n')

    out.write('\n--- Sample Curated ---\n')
    if sample_curated:
        ct = sample_curated.get('conversation_type', '')
        sid = sample_curated.get('shloka_id', '')
        msgs = sample_curated.get('messages', [])
        first_user = ''
        for m in msgs:
            if m.get('role') == 'user':
                first_user = m.get('content', '')[:80]
                break
        out.write(f'conversation_type: {ct}\n')
        out.write(f'shloka_id: {sid}\n')
        out.write(f'first_user_content: {first_user}\n')
    else:
        out.write('No curated sample found\n')

    out.write('\n--- Sample Variation ---\n')
    if sample_variation:
        ct = sample_variation.get('conversation_type', '')
        sid = sample_variation.get('shloka_id', '')
        msgs = sample_variation.get('messages', [])
        first_user = ''
        for m in msgs:
            if m.get('role') == 'user':
                first_user = m.get('content', '')[:80]
                break
        out.write(f'conversation_type: {ct}\n')
        out.write(f'shloka_id: {sid}\n')
        out.write(f'first_user_content: {first_user}\n')
    else:
        out.write('No variation sample found\n')

print('Done')