import sys
import re

def patch_gui(file_path):
    with open(file_path, 'r', encoding='utf-8', newline='') as f:
        content = f.read()

    # 1. Patch _render_provider_panel Session ID
    content = content.replace(
        '   imgui.text(f"Session ID: {sid}")',
        '   imgui.text("Session ID:"); imgui.same_line(); self._render_selectable_label("gemini_cli_sid", sid, width=200)'
    )

    # 2. Patch _render_token_budget_panel Session Telemetry
    content = content.replace(
        '  imgui.text_colored(C_RES, f"Tokens: {total:,} (In: {usage[\'input_tokens\']:,} Out: {usage[\'output_tokens\']:,})")',
        '  self._render_selectable_label("session_telemetry_tokens", f"Tokens: {total:,} (In: {usage[\'input_tokens\']:,} Out: {usage[\'output_tokens\']:,})", width=-1, color=C_RES)'
    )

    # 3. Patch _render_token_budget_panel MMA Tier Costs table
    # This is trickier, let's find the loop
    tier_table_pattern = re.compile(
        r'(for tier, stats in self\.mma_tier_usage\.items\(\):\s+.*?imgui\.table_set_column_index\(0\); )imgui\.text\(tier\)(\s+imgui\.table_set_column_index\(1\); )imgui\.text\(model\.split\(\'-\'\)\[0\]\)(\s+imgui\.table_set_column_index\(2\); )imgui\.text\(f"\{tokens:,\}"\)(\s+imgui\.table_set_column_index\(3\); )imgui\.text_colored\(imgui\.ImVec4\(0\.2, 0\.9, 0\.2, 1\), f"\$\{cost:\.4f\}"\)',
        re.DOTALL
    )
    
    def tier_replacement(match):
        return (match.group(1) + 'self._render_selectable_label(f"tier_{tier}", tier, width=-1)' +
                match.group(2) + 'self._render_selectable_label(f"model_{tier}", model.split("-")[0], width=-1)' +
                match.group(3) + 'self._render_selectable_label(f"tokens_{tier}", f"{tokens:,}", width=-1)' +
                match.group(4) + 'self._render_selectable_label(f"cost_{tier}", f"${cost:.4f}", width=-1, color=imgui.ImVec4(0.2, 0.9, 0.2, 1))')

    content = tier_table_pattern.sub(tier_replacement, content)

    # 4. Patch _render_token_budget_panel Session Total
    content = content.replace(
        '    imgui.text_colored(imgui.ImVec4(0, 1, 0, 1), f"Session Total: ${tier_total:.4f}")',
        '    self._render_selectable_label("session_total_cost", f"Session Total: ${tier_total:.4f}", width=-1, color=imgui.ImVec4(0, 1, 0, 1))'
    )

    with open(file_path, 'w', encoding='utf-8', newline='') as f:
        f.write(content)
    print("Successfully patched src/gui_2.py for selectable metrics")

if __name__ == "__main__":
    patch_gui("src/gui_2.py")