import os
import re

def fix_nvim_msgid():
    filepath = 'nvim-core/src/lib.rs'
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()

    # Add atomic import and static var if not exists
    if 'static NEXT_MSGID' not in content:
        atomic_def = "use std::sync::atomic::{AtomicU64, Ordering};\nstatic NEXT_MSGID: AtomicU64 = AtomicU64::new(1);\n"
        # Find NVIM_CONN
        conn_idx = content.find('static NVIM_CONN')
        if conn_idx != -1:
            content = content[:conn_idx] + atomic_def + content[conn_idx:]

    # Replace all hardcoded msgid
    # e.g., RmpValue::Integer(1.into()), // msgid
    # with: let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst); ... RmpValue::Integer(msgid.into()),

    # We need to insert `let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);` before `let req = RmpValue::Array(vec![`
    
    # We will use regex to find the blocks
    funcs = [
        ('send_nvim_command', '1'),
        ('get_nvim_active_buffer', '2'),
        ('get_nvim_cursor', '3'),
        ('get_nvim_visual_selection', '4'),
        ('set_nvim_diagnostics', '5'),
        ('execute_nvim_lua', '6'),
    ]

    for func, old_id in funcs:
        pattern = rf"let req = RmpValue::Array\(vec!\[\s*RmpValue::Integer\(0\.into\(\)\),\s*RmpValue::Integer\({old_id}\.into\(\)\), // msgid"
        replacement = f"let msgid = NEXT_MSGID.fetch_add(1, Ordering::SeqCst);\n    let req = RmpValue::Array(vec![\n        RmpValue::Integer(0.into()),\n        RmpValue::Integer(msgid.into()), // msgid"
        
        content = re.sub(pattern, replacement, content)

    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(content)

    print('Fixed msgid allocations')

fix_nvim_msgid()