import sys

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

# Remove WS code
ws_start = content.find('async fn ws_handler')
if ws_start != -1:
    content = content[:ws_start]

# Remove /ws route
route_idx = content.find('.route("/ws", get(ws_handler))')
if route_idx != -1:
    content = content.replace('.route("/ws", get(ws_handler))\n            ', '')

# Remove WS import
import_idx = content.find('use axum::extract::ws::{WebSocketUpgrade, WebSocket, Message};\n')
if import_idx != -1:
    content = content.replace('use axum::extract::ws::{WebSocketUpgrade, WebSocket, Message};\n', '')

# Fix GC task state
gc_start = content.find('        // Background Garbage Collection for old tasks')
if gc_start != -1:
    content = content.replace('let state_gc = state_clone.clone();', 'let state_gc = Arc::clone(&state);')

# Fix Git task state and CodeChange fields
git_start = content.find('        // Git Native Sync Background Task')
if git_start != -1:
    content = content.replace('let state_git = state_clone.clone();', 'let state_git = Arc::clone(&state);')
    content = content.replace('commit_hash: current_id.clone(),', 'git_commit: Some(current_id.clone()),')
    content = content.replace('branch: branch,', 'git_branch: Some(branch),')
    content = content.replace('files_changed: vec![],', 'file_path: "".to_string(),')

with open(main_path, "w", encoding="utf-8") as f:
    f.write(content)
print("Cleaned up WS and fixed main.rs")