import os

def fix_main_stats():
    filepath = 'server/src/main.rs'
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()

    old_stats = """                        let (entities, relations) = {
                            let graph = state_clone.get_full_graph();
                            (graph.entities.len(), graph.relations.len())
                        };
                        let tasks = state_clone.tasks.read().len();
                        let snippets = state_clone.snippets.read().len();
                        let tech_debts = state_clone.tech_debts.read().len();
                        let adrs = state_clone.adrs.read().len();

                        let ledger = state_clone.ledger.read().len();
                        let sticky = state_clone.sticky.read().len();
                        let error_fixes = state_clone.error_fixes.read().len();
                        let pinned_files = state_clone.pinned_files.read().len();
                        let session_summaries = state_clone.session_summaries.read().len();
                        let handoff_memos = state_clone.handoff_memos.read().len();
                        let env_fingerprints = state_clone.env_fingerprints.read().len();
                        let env_requirements = state_clone.env_requirements.read().len();
                        let milestones = state_clone.milestones.read().len();
                        let environments = state_clone.environments.read().len();
                        let pr_checklists = state_clone.pr_checklists.read().len();
                        let gates = state_clone.gates.read().len();
                        let context_workspaces = state_clone.context_workspaces.read().len();"""

    new_stats = """                        let (entities, relations) = state_clone.read_graph(|g| (g.entities.len(), g.relations.len()));
                        let tasks = state_clone.tasks.read_with(|items| items.len());
                        let snippets = state_clone.snippets.read_with(|items| items.len());
                        let tech_debts = state_clone.tech_debts.read_with(|items| items.len());
                        let adrs = state_clone.adrs.read_with(|items| items.len());

                        let ledger = state_clone.ledger.read_with(|items| items.len());
                        let sticky = state_clone.sticky.read_with(|items| items.len());
                        let error_fixes = state_clone.error_fixes.read_with(|items| items.len());
                        let pinned_files = state_clone.pinned_files.read_with(|items| items.len());
                        let session_summaries = state_clone.session_summaries.read_with(|items| items.len());
                        let handoff_memos = state_clone.handoff_memos.read_with(|items| items.len());
                        let env_fingerprints = state_clone.env_fingerprints.read_with(|items| items.len());
                        let env_requirements = state_clone.env_requirements.read_with(|items| items.len());
                        let milestones = state_clone.milestones.read_with(|items| items.len());
                        let environments = state_clone.environments.read_with(|items| items.len());
                        let pr_checklists = state_clone.pr_checklists.read_with(|items| items.len());
                        let gates = state_clone.gates.read_with(|items| items.len());
                        let context_workspaces = state_clone.context_workspaces.read_with(|items| items.len());"""

    if old_stats in content:
        content = content.replace(old_stats, new_stats)
        print('Replaced stats reads')
    else:
        print('Could not find old_stats')

    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(content)

fix_main_stats()