import re
import sys

with open('C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\stub\\src\\main.rs', 'r', encoding='utf-8') as f:
    content = f.read()

start1 = r'let mut log_prefix = String::from\("Message"\);\n\s*if let Ok\(v\) = serde_json::from_str::<serde_json::Value>\(&msg\) \{[\s\S]*?log_prefix = format!\("Response id=\{\}", res_id\);\n\s*\}\n\s*\}'
content = re.sub(start1, r'let log_prefix = logger::extract_log_prefix(&msg, false);', content)

start2 = r'let mut log_prefix = String::from\("Message"\);\n\s*if let Ok\(v\) = serde_json::from_str::<serde_json::Value>\(&text\) \{[\s\S]*?log_prefix = format!\("Response id=\{\} \[Error: \{\}\]", id, is_error\);\n\s*\}'
content = re.sub(start2, r'let log_prefix = logger::extract_log_prefix(&text, true);', content)

with open('C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\stub\\src\\main.rs', 'w', encoding='utf-8') as f:
    f.write(content)
print("Patched stub")