import re
import os

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

# Replace MemoryHandler struct
struct_pattern = r'pub struct MemoryHandler \{\s*pub state: Arc<MemoryState>,\s*\}'

new_struct = """use crate::router::McpTool;

pub struct MemoryHandler {
    pub state: Arc<MemoryState>,
    pub tools: std::collections::HashMap<String, Box<dyn McpTool>>,
}

impl MemoryHandler {
    pub fn new(state: Arc<MemoryState>) -> Self {
        let mut tools: std::collections::HashMap<String, Box<dyn McpTool>> = std::collections::HashMap::new();
        
        macro_rules! register {
            ($module:ident::$handler:ident) => {
                let h = crate::handlers_v2::$module::$handler;
                tools.insert(h.name().to_string(), Box::new(h));
            };
        }

        register!(graph::QueryGraphPathHandler);
        register!(graph::CreateEntitiesHandler);
        register!(graph::CreateRelationsHandler);
        register!(graph::AddObservationsHandler);
        register!(graph::DeleteEntitiesHandler);
        register!(graph::DeleteObservationsHandler);
        register!(graph::DeleteRelationsHandler);
        register!(graph::ReadGraphHandler);
        register!(graph::SearchNodesHandler);
        register!(graph::OpenNodesHandler);
        register!(graph::VisualizeGraphHandler);
        register!(graph::CondenseEntityHandler);
        register!(graph::MergeEntitiesHandler);
        register!(graph::FindOrphansHandler);

        register!(tasks::AddTaskHandler);
        register!(tasks::DeleteTaskHandler);
        register!(tasks::UpdateTaskStatusHandler);
        register!(tasks::ListActiveTasksHandler);
        register!(tasks::SetAcceptanceCriteriaHandler);
        register!(tasks::VerifyAcceptanceCriteriaHandler);
        register!(tasks::AddMilestoneHandler);
        register!(tasks::UpdateMilestoneHandler);
        register!(tasks::ListMilestonesHandler);

        register!(notes::AddStickyNoteHandler);
        register!(notes::ReadStickyNotesHandler);
        register!(notes::DeleteStickyNoteHandler);
        register!(notes::ClearStickyNotesHandler);
        register!(notes::LeaveHandoffMemoHandler);
        register!(notes::ReadHandoffMemosHandler);
        register!(notes::ClearHandoffMemosHandler);
        register!(notes::AddSessionSummaryHandler);
        register!(notes::GenerateStandupReportHandler);

        register!(meta::LogDecisionHandler);
        register!(meta::QueryDecisionsHandler);
        register!(meta::LogErrorFixHandler);
        register!(meta::SearchErrorFixesHandler);
        register!(meta::LogCodeChangeHandler);
        register!(meta::QueryRecentChangesHandler);
        register!(meta::LearnPreferenceHandler);
        register!(meta::ReadPreferencesHandler);
        register!(meta::LogTechDebtHandler);
        register!(meta::ResolveTechDebtHandler);
        register!(meta::ListTechDebtHandler);
        register!(meta::OmniSearchHandler);
        register!(meta::GetProjectHealthHandler);

        register!(env::UpdateEnvFingerprintHandler);
        register!(env::ReadEnvFingerprintHandler);
        register!(env::LogEnvRequirementHandler);
        register!(env::RegisterEnvironmentHandler);
        register!(env::GetEnvironmentDetailsHandler);

        register!(workspaces::PinFileHandler);
        register!(workspaces::UnpinFileHandler);
        register!(workspaces::ListPinnedFilesHandler);
        register!(workspaces::StoreSnippetHandler);
        register!(workspaces::SearchSnippetsHandler);
        register!(workspaces::DeleteSnippetHandler);
        register!(workspaces::SaveContextWorkspaceHandler);
        register!(workspaces::LoadContextWorkspaceHandler);
        register!(workspaces::ListContextWorkspacesHandler);
        register!(workspaces::AddPrChecklistItemHandler);
        register!(workspaces::GetPrChecklistHandler);
        register!(workspaces::ClearPrChecklistHandler);

        Self { state, tools }
    }
"""

content = re.sub(struct_pattern, new_struct, content)
content = content.replace("impl MemoryHandler {\n    pub async fn handle_request", "    pub async fn handle_request")

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

print("MemoryHandler struct updated.")