import sys

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

# 1. In _render_provider_panel, remove Fetch Models
old_fetch = """  imgui.text("Model")
  imgui.same_line()
  if imgui.button("Fetch Models"):
   self._fetch_models(self.current_provider)
  if imgui.begin_list_box("##models", imgui.ImVec2(-1, 120)):"""
new_fetch = """  imgui.text("Model")
  if imgui.begin_list_box("##models", imgui.ImVec2(-1, 120)):"""
content = content.replace(old_fetch, new_fetch)

# 2. Extract Persona block
# We need to find the start of 'imgui.text("Persona")' and end of 'self._editing_persona_is_new = True'
# Let's be very careful.
old_persona_block = """  imgui.text("Persona")
  if not hasattr(self, 'ui_active_persona'):
   self.ui_active_persona = ""
  personas = getattr(self.controller, 'personas', {})
  if imgui.begin_combo("##persona", self.ui_active_persona or "None"):
   if imgui.selectable("None", not self.ui_active_persona)[0]:
    self.ui_active_persona = ""
   for pname in sorted(personas.keys()):
    if imgui.selectable(pname, pname == self.ui_active_persona)[0]:
     self.ui_active_persona = pname
     if pname in personas:
      persona = personas[pname]
      self._editing_persona_name = persona.name
      self._editing_persona_provider = persona.provider or ""
      self._editing_persona_model = persona.model or ""
      self._editing_persona_system_prompt = persona.system_prompt or ""
      self._editing_persona_temperature = persona.temperature or 0.7
      self._editing_persona_max_tokens = persona.max_output_tokens or 4096
      self._editing_persona_tool_preset_id = persona.tool_preset or ""
      self._editing_persona_bias_profile_id = persona.bias_profile or ""
      import json
      self._editing_persona_preferred_models = json.dumps(persona.preferred_models) if persona.preferred_models else "[]"
      self._editing_persona_is_new = False
      if persona.provider and persona.provider in self.controller.PROVIDERS:
       self.current_provider = persona.provider
      if persona.model:
       self.current_model = persona.model
      if persona.temperature is not None:
       ai_client.temperature = persona.temperature
      if persona.max_output_tokens:
       ai_client.max_output_tokens = persona.max_output_tokens
      if persona.system_prompt:
       ai_client.system_instruction = persona.system_prompt
      if persona.tool_preset:
       self.ui_active_tool_preset = persona.tool_preset
       ai_client.set_tool_preset(persona.tool_preset)
      if persona.bias_profile:
       self.ui_active_bias_profile = persona.bias_profile
       ai_client.set_bias_profile(persona.bias_profile)
   imgui.end_combo()
  imgui.same_line()
  if imgui.button("Manage Personas"):
   self.show_persona_editor_window = True
   if self.ui_active_persona and self.ui_active_persona in personas:
    persona = personas[self.ui_active_persona]
    self._editing_persona_name = persona.name
    self._editing_persona_provider = persona.provider or ""
    self._editing_persona_model = persona.model or ""
    self._editing_persona_system_prompt = persona.system_prompt or ""
    self._editing_persona_temperature = persona.temperature if persona.temperature is not None else 0.7
    self._editing_persona_max_tokens = persona.max_output_tokens if persona.max_output_tokens is not None else 4096
    self._editing_persona_tool_preset_id = persona.tool_preset or ""
    self._editing_persona_bias_profile_id = persona.bias_profile or ""
    self._editing_persona_preferred_models_list = list(persona.preferred_models) if persona.preferred_models else []
    self._editing_persona_scope = self.controller.persona_manager.get_persona_scope(persona.name)
    self._editing_persona_is_new = False
   else:
    self._editing_persona_name = ""
    self._editing_persona_provider = self.current_provider
    self._editing_persona_model = self.current_model
    self._editing_persona_system_prompt = ""
    self._editing_persona_temperature = 0.7
    self._editing_persona_max_tokens = 4096
    self._editing_persona_tool_preset_id = ""
    self._editing_persona_bias_profile_id = ""
    self._editing_persona_preferred_models_list = []
    self._editing_persona_scope = "project"
    self._editing_persona_is_new = True"""

# We need to extract the bias profile block as well
old_bias_block = """  imgui.text("Bias Profile")
  if imgui.begin_combo("##bias", self.ui_active_bias_profile or "None"):
   if imgui.selectable("None", not self.ui_active_bias_profile)[0]:
    self.ui_active_bias_profile = ""
    ai_client.set_bias_profile(None)
   for bname in sorted(self.bias_profiles.keys()):
    if imgui.selectable(bname, bname == self.ui_active_bias_profile)[0]:
     self.ui_active_bias_profile = bname
     ai_client.set_bias_profile(bname)
   imgui.end_combo()"""

# Remove them from their original spots
content = content.replace(old_bias_block, "")
content = content.replace(old_persona_block, "")

# Insert Persona block at the top of _render_provider_panel
old_provider_start = """ def _render_provider_panel(self) -> None:
  if self.perf_profiling_enabled: self.perf_monitor.start_component("_render_provider_panel")
  imgui.text("Provider")"""
new_provider_start = f""" def _render_provider_panel(self) -> None:
  if self.perf_profiling_enabled: self.perf_monitor.start_component("_render_provider_panel")
{old_persona_block}
  imgui.separator()
  imgui.text("Provider")"""
content = content.replace(old_provider_start, new_provider_start)

# Update _render_agent_tools_panel
old_agent_tools_start = """ def _render_agent_tools_panel(self) -> None:
  imgui.text_colored(C_LBL, 'Active Tool Preset')"""
new_agent_tools_start = f""" def _render_agent_tools_panel(self) -> None:
  if imgui.collapsing_header("Active Tool Presets & Biases", imgui.TreeNodeFlags_.default_open):
   imgui.text("Tool Preset")"""
content = content.replace(old_agent_tools_start, new_agent_tools_start)

# Wait, if I do collapsing header, I need to indent the rest of the function.
# Instead of indenting the whole function, I can just use the header. 
# But wait, ImGui collapsing_header doesn't require indenting, it just returns true if open.
# So I should write it properly:
old_agent_tools_func = """ def _render_agent_tools_panel(self) -> None:
  imgui.text_colored(C_LBL, 'Active Tool Preset')
  presets = self.controller.tool_presets
  preset_names = [""] + sorted(list(presets.keys()))
  
  # Gracefully handle None or missing preset
  active = getattr(self, "ui_active_tool_preset", "")
  if active is None: active = ""
  try:
   idx = preset_names.index(active)
  except ValueError:
   idx = 0
  
  ch, new_idx = imgui.combo("##tool_preset_select", idx, preset_names)
  if ch:
   self.ui_active_tool_preset = preset_names[new_idx]
  
  imgui.same_line()
  if imgui.button("Manage Presets##tools"):
   self.show_tool_preset_manager_window = True
  if imgui.is_item_hovered():
   imgui.set_tooltip("Configure tool availability and default modes.")
  
  imgui.dummy(imgui.ImVec2(0, 8))
  active_name = self.ui_active_tool_preset
  if active_name and active_name in presets:
   preset = presets[active_name]
   for cat_name, tools in preset.categories.items():
    if imgui.tree_node(cat_name):
     for tool in tools:
      if tool.weight >= 5:
       imgui.text_colored(vec4(255, 100, 100), "[HIGH]")
       imgui.same_line()
      elif tool.weight == 4:
       imgui.text_colored(vec4(255, 255, 100), "[PREF]")
       imgui.same_line()
      elif tool.weight == 2:
       imgui.text_colored(vec4(255, 150, 50), "[REJECT]")
       imgui.same_line()
      elif tool.weight <= 1:
       imgui.text_colored(vec4(180, 180, 180), "[LOW]")
       imgui.same_line()
      
      imgui.text(tool.name)
      imgui.same_line(180)
      
      mode = tool.approval
      if imgui.radio_button(f"Auto##{cat_name}_{tool.name}", mode == "auto"):
       tool.approval = "auto"
      imgui.same_line()
      if imgui.radio_button(f"Ask##{cat_name}_{tool.name}", mode == "ask"):
       tool.approval = "ask"
     imgui.tree_pop()"""

new_agent_tools_func = """ def _render_agent_tools_panel(self) -> None:
  if imgui.collapsing_header("Active Tool Presets & Biases", imgui.TreeNodeFlags_.default_open):
   imgui.text("Tool Preset")
   presets = self.controller.tool_presets
   preset_names = [""] + sorted(list(presets.keys()))
   
   # Gracefully handle None or missing preset
   active = getattr(self, "ui_active_tool_preset", "")
   if active is None: active = ""
   try:
    idx = preset_names.index(active)
   except ValueError:
    idx = 0
   
   ch, new_idx = imgui.combo("##tool_preset_select", idx, preset_names)
   if ch:
    self.ui_active_tool_preset = preset_names[new_idx]
   
   imgui.same_line()
   if imgui.button("Manage Tools##tools"):
    self.show_tool_preset_manager_window = True
   if imgui.is_item_hovered():
    imgui.set_tooltip("Configure tool availability and default modes.")
   
   imgui.dummy(imgui.ImVec2(0, 4))
""" + "\n   ".join(old_bias_block.split("\n")) + """
   
   imgui.dummy(imgui.ImVec2(0, 8))
   active_name = self.ui_active_tool_preset
   if active_name and active_name in presets:
    preset = presets[active_name]
    for cat_name, tools in preset.categories.items():
     if imgui.tree_node(cat_name):
      for tool in tools:
       if tool.weight >= 5:
        imgui.text_colored(vec4(255, 100, 100), "[HIGH]")
        imgui.same_line()
       elif tool.weight == 4:
        imgui.text_colored(vec4(255, 255, 100), "[PREF]")
        imgui.same_line()
       elif tool.weight == 2:
        imgui.text_colored(vec4(255, 150, 50), "[REJECT]")
        imgui.same_line()
       elif tool.weight <= 1:
        imgui.text_colored(vec4(180, 180, 180), "[LOW]")
        imgui.same_line()
       
       imgui.text(tool.name)
       imgui.same_line(180)
       
       mode = tool.approval
       if imgui.radio_button(f"Auto##{cat_name}_{tool.name}", mode == "auto"):
        tool.approval = "auto"
       imgui.same_line()
       if imgui.radio_button(f"Ask##{cat_name}_{tool.name}", mode == "ask"):
        tool.approval = "ask"
      imgui.tree_pop()"""
content = content.replace(old_agent_tools_func, new_agent_tools_func)

# Fix cache text display in Usage Analytics
content = content.replace('self._gemini_cache_text = f"Gemini Caches: {count} ({size_bytes / 1024:.1f} KB)"', 'self._gemini_cache_text = f"Cache Usage: {count} ({size_bytes / 1024:.1f} KB)"')
content = content.replace('imgui.text_colored(C_LBL, f"Gemini Cache: ACTIVE | Age: {age:.0f}s / {ttl}s | Renews at: {ttl * 0.9:.0f}s")', 'imgui.text_colored(C_LBL, f"Cache Usage: ACTIVE | Age: {age:.0f}s / {ttl}s | Renews at: {ttl * 0.9:.0f}s")')
content = content.replace('imgui.text_disabled("Gemini Cache: INACTIVE")', 'imgui.text_disabled("Cache Usage: INACTIVE")')

# Also, user requested: "The persona should problably just mess with the project system prompt for now."
# Currently in persona selection: `ai_client.system_instruction = persona.system_prompt`
# Let's change that to `self.ui_project_system_prompt = persona.system_prompt` and remove ai_client direct injection
content = content.replace('ai_client.system_instruction = persona.system_prompt', 'self.ui_project_system_prompt = persona.system_prompt')

with open("src/gui_2.py", "w", encoding="utf-8") as f:
    f.write(content)
print("done")