import re

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

rebuild_index_old = """    pub fn rebuild_index(&self) {
        if let Ok(new_idx) = MemoryIndex::new(&self.base_dir) {
            let session = self.session_graph.read().unwrap();
            let mut full = {
                let cache = self.master_cache.read().unwrap();
                cache.0.clone()
            };
            Self::merge_graphs(&mut full, &session);
            for e in full.entities.values() {
                let _ = new_idx.index_entity(e);
            }
            for t in self.tasks.read() {
                let _ = new_idx.index_task(&t);
            }
            for s in self.snippets.read() {
                let _ = new_idx.index_snippet(&s);
            }
            for a in self.adrs.read() {
                let _ = new_idx.index_adr(&a);
            }
            if let Ok(mut w) = self.search_index.write() {
                *w = new_idx;
            }
        }
    }"""

rebuild_index_new = """    pub fn rebuild_index(&self) {
        if let Ok(new_idx) = MemoryIndex::new(&self.base_dir) {
            let session = self.session_graph.read().unwrap();
            let cache = self.master_cache.read().unwrap();
            
            // Index entities that are only in master, or merge if they are in both
            for (name, e) in &cache.0.entities {
                if let Some(session_e) = session.entities.get(name) {
                    let mut merged_e = e.clone();
                    for obs in &session_e.observations {
                        if !merged_e.observations.contains(obs) {
                            merged_e.observations.push(obs.clone());
                        }
                    }
                    let _ = new_idx.index_entity(&merged_e);
                } else {
                    let _ = new_idx.index_entity(e);
                }
            }
            // Index entities that are only in session
            for (name, session_e) in &session.entities {
                if !cache.0.entities.contains_key(name) {
                    let _ = new_idx.index_entity(session_e);
                }
            }
            
            for t in self.tasks.read() {
                let _ = new_idx.index_task(&t);
            }
            for s in self.snippets.read() {
                let _ = new_idx.index_snippet(&s);
            }
            for a in self.adrs.read() {
                let _ = new_idx.index_adr(&a);
            }
            if let Ok(mut w) = self.search_index.write() {
                *w = new_idx;
            }
        }
    }"""

content = content.replace(rebuild_index_old, rebuild_index_new)

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