import re

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

unique_items_block = re.search(r'    pub fn unique_items<.*?>\(.*?\) -> Vec<T> \{\n(?:.*?\n){1,15}?    \}\n', content, re.DOTALL)
if unique_items_block:
    content = content.replace(unique_items_block.group(0), "")

merge_graphs_old = """    pub fn merge_graphs(dest: &mut KnowledgeGraph, src: &KnowledgeGraph) {
        for (name, src_ent) in &src.entities {
            let dest_ent = dest
                .entities
                .entry(name.clone())
                .or_insert_with(|| src_ent.clone());
            if dest_ent.name == src_ent.name {
                dest_ent.observations.extend(src_ent.observations.clone());
                dest_ent.observations = Self::unique_items(dest_ent.observations.clone());
            }
        }
        dest.relations.extend(src.relations.clone());
        dest.relations = Self::unique_items(dest.relations.clone());
    }"""

merge_graphs_new = """    pub fn merge_graphs(dest: &mut KnowledgeGraph, src: &KnowledgeGraph) {
        for (name, src_ent) in &src.entities {
            let dest_ent = dest
                .entities
                .entry(name.clone())
                .or_insert_with(|| crate::models::Entity {
                    name: src_ent.name.clone(),
                    entity_type: src_ent.entity_type.clone(),
                    observations: Vec::new(),
                });
            
            for obs in &src_ent.observations {
                if !dest_ent.observations.contains(obs) {
                    dest_ent.observations.push(obs.clone());
                }
            }
        }
        for rel in &src.relations {
            if !dest.relations.contains(rel) {
                dest.relations.push(rel.clone());
            }
        }
    }"""

content = content.replace(merge_graphs_old, merge_graphs_new)

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