import re

with open("server/src/state.rs", "r", encoding="utf-8") as f:
    content = f.read()

recover_wal_code = """
    pub fn recover_wal(&self) {
        let wal_path = self.base_dir.join("wal.jsonl");
        if let Ok(content) = std::fs::read_to_string(&wal_path) {
            let mut session = self.session_graph.write().unwrap();
            for line in content.lines() {
                if let Ok(d) = serde_json::from_str::<KnowledgeGraph>(line) {
                    Self::merge_graphs(&mut session, &d);
                }
            }
        }
    }
"""
content = content.replace("    pub fn merge_graphs(dest: &mut KnowledgeGraph, src: &KnowledgeGraph) {", recover_wal_code + "    pub fn merge_graphs(dest: &mut KnowledgeGraph, src: &KnowledgeGraph) {")

get_full_graph_old = """    pub fn get_full_graph(&self) -> KnowledgeGraph {
        let mut master = self.read_master_cached();
        let wal_path = self.base_dir.join("wal.jsonl");
        if let Ok(content) = std::fs::read_to_string(&wal_path) {
            for line in content.lines() {
                if let Ok(d) = serde_json::from_str::<KnowledgeGraph>(line) {
                    Self::merge_graphs(&mut master, &d);
                }
            }
        }
        let session_graph = self.session_graph.read().unwrap();
        Self::merge_graphs(&mut master, &session_graph);
        master
    }"""

get_full_graph_new = """    pub fn get_full_graph(&self) -> KnowledgeGraph {
        let mut master = self.read_master_cached();
        let session_graph = self.session_graph.read().unwrap();
        Self::merge_graphs(&mut master, &session_graph);
        master
    }"""

content = content.replace(get_full_graph_old, get_full_graph_new)

with open("server/src/state.rs", "w", encoding="utf-8") as f:
    f.write(content)

with open("server/src/main.rs", "r", encoding="utf-8") as f:
    main_content = f.read()

main_content = main_content.replace("    state.rebuild_index();\n\n    run_server(state)", "    state.recover_wal();\n    state.rebuild_index();\n\n    run_server(state)")

with open("server/src/main.rs", "w", encoding="utf-8") as f:
    f.write(main_content)