import ast
import sys
import re

def restore():
 file_path = "src/gui_2.py"
 with open(file_path, "r", encoding="utf-8") as f:
  content = f.read()
 
 # Comprehensive region mapping for ALL render functions
 region_map = {
  "Main Interface": ["render_main_interface", "render_custom_title_bar"],
  "Diagnostics & Analytics": ["render_usage_analytics_panel", "render_diagnostics_panel", "render_cache_panel", "render_tool_analytics_panel", "render_token_budget_panel"],
  "Logging": ["render_log_management"],
  "Project Management": ["render_project_settings_hub", "render_projects_panel", "render_paths_panel", "render_path_field"],
  "AI Settings": ["render_ai_settings_hub", "render_rag_panel", "render_system_prompts_panel", "render_agent_tools_panel", "render_provider_panel", "render_persona_selector_panel", "render_base_prompt_diff_modal", "render_save_preset_modal", "render_preset_manager_content", "render_preset_manager_window", "render_tool_preset_manager_content", "render_tool_preset_manager_window", "render_persona_editor_window"],
  "Context Management": ["render_files_and_media", "render_files_panel", "render_screenshots_panel", "render_context_batch_actions", "render_add_context_files_modal", "render_context_composition_panel", "render_ast_inspector_modal", "render_save_workspace_profile_modal", "render_context_presets_panel", "render_context_screenshots", "render_context_files_table", "render_context_presets", "render_snapshot_tab"],
  "Discussions": ["render_discussion_hub", "render_discussion_entry", "render_discussion_entry_read_mode", "render_history_window", "render_session_insights_panel", "render_prior_session_view", "render_thinking_indicator", "render_synthesis_panel", "render_comms_history_panel", "render_takes_panel", "render_discussion_entries", "render_discussion_entry_controls", "render_discussion_metadata", "render_discussion_panel", "render_discussion_roles", "render_discussion_selector", "render_discussion_tab"],
  "Operations Monitor": ["render_operations_hub", "render_message_panel", "render_response_panel", "render_tool_calls_panel", "render_external_tools_panel", "render_text_viewer_window", "render_patch_modal", "render_external_editor_panel", "render_approve_script_modal"],
  "Misc Tools": ["render_theme_panel", "render_shader_live_editor", "render_markdown_test", "render_error_tint", "render_text_viewer", "render_heavy_text", "render_thinking_trace", "render_selectable_label", "render_window_if_open"],
  "MMA": ["render_mma_dashboard", "render_mma_modals", "render_mma_track_summary", "render_mma_epic_planner", "render_mma_conductor_setup", "render_mma_track_browser", "render_mma_global_controls", "render_mma_usage_section", "render_mma_ticket_editor", "render_mma_agent_streams", "render_tier_stream_panel", "render_track_proposal_modal", "render_ticket_queue", "render_task_dag_panel", "render_beads_tab", "render_mma_focus_selector"]
 }

 # Parse to get exact line ranges
 lines = content.splitlines()
 tree = ast.parse(content)
 
 func_blocks = {}
 
 for node in tree.body:
  if isinstance(node, ast.FunctionDef):
   f_name = node.name
   s_idx = node.lineno - 1
   if node.decorator_list: s_idx = node.decorator_list[0].lineno - 1
   e_idx = node.end_lineno
   func_blocks[f_name] = "\n".join(lines[s_idx:e_idx])

 # Find the start of module-level render functions
 first_render_idx = len(lines)
 for node in tree.body:
  if isinstance(node, ast.FunctionDef) and node.name.startswith("render_"):
   first_render_idx = min(first_render_idx, node.lineno - 1)
   if node.decorator_list: first_render_idx = min(first_render_idx, node.decorator_list[0].lineno - 1)

 # Get prefix and strip existing region tags
 prefix_raw = lines[:first_render_idx]
 prefix = []
 for line in prefix_raw:
  if line.strip().startswith("#region:") or line.strip().startswith("#endregion:"):
   continue
  prefix.append(line)

 # Re-assemble with regions
 assembled = []
 assigned = set()
 for r_name, funcs in region_map.items():
  group = []
  for f in funcs:
   if f in func_blocks:
    group.append(func_blocks[f])
    assigned.add(f)
  if group:
   assembled.append(f"#region: {r_name}\n\n" + "\n\n".join(group) + f"\n\n#endregion: {r_name}")

 # Append unassigned render functions
 for f_name, block in func_blocks.items():
  if f_name.startswith("render_") and f_name not in assigned:
   assembled.append(block)

 final_content = "\n".join(prefix).strip() + "\n\n" + "\n\n".join(assembled) + "\n"
 final_content = re.sub(r"\n\n\n+", "\n\n", final_content)

 with open(file_path, "w", encoding="utf-8") as f:
  f.write(final_content)
 print("Comprehensive regions restored successfully.")

if __name__ == "__main__":
 restore()