Update app.py

app.py

@@ -93,18 +93,25 @@ def _clean_filename(filename_line_content):
     path_match = re.match(r'^([\w\-\.\s\/\\]+)', text) # Adjusted to be more general
     if path_match:
         # Further clean if it looks like "path/to/file (description)"
-        return re.split(r'\s*\(', path_match.group(1).strip(), 1).strip() # Corrected split index
+        # Corrected split index was already correct in the original code, just ensure it's applied
+        parts = re.split(r'\s*\(', path_match.group(1).strip(), 1)
+        return parts[0].strip() if parts else ""
 
     # Fallback for more complex lines, like "### File: `src/app.py` (main application)"
     backtick_match = re.search(r'`([^`]+)`', text)
     if backtick_match:
         potential_fn = backtick_match.group(1).strip()
-        cleaned_fn = re.split(r'\s*\(|\s{2,}', potential_fn, 1).strip() # Corrected split index
+        # Corrected split index was already correct
+        parts = re.split(r'\s*\(|\s{2,}', potential_fn, 1)
+        cleaned_fn = parts[0].strip() if parts else ""
+        cleaned_fn = cleaned_fn.strip('`\'":;,') # Clean common wrapping chars
         if cleaned_fn: return cleaned_fn
 
-    filename_candidate = re.split(r'\s*\(|\s{2,}', text, 1).strip() # Corrected split index
+    # Final fallback
+    parts = re.split(r'\s*\(|\s{2,}', text, 1)
+    filename_candidate = parts[0].strip() if parts else text.strip()
     filename_candidate = filename_candidate.strip('`\'":;,') # Clean common wrapping chars
-    return filename_candidate if filename_candidate else text # Fallback to original if all else fails
+    return filename_candidate if filename_candidate else text.strip()
 
 
 def _parse_chat_stream_logic(chat_json_string, existing_files_state=None):
@@ -120,27 +127,57 @@ def _parse_chat_stream_logic(chat_json_string, existing_files_state=None):
     except json.JSONDecodeError as e: results["error_message"] = f"JSON Parsing Error: {e}."; return results
     except ValueError as e: results["error_message"] = str(e); return results
 
+    # Ensure the bot's message is the last one and is the only one being parsed for new/updated files
+    # This prevents reprocessing the entire history on every new AI turn
+    message_obj = None
+    if ai_chat_history and isinstance(ai_chat_history[-1], dict) and ai_chat_history[-1].get("role", "").lower() == BOT_ROLE_NAME:
+         message_obj = ai_chat_history[-1]
+    
+    if not message_obj:
+         # If the last message isn't the bot's or is malformed, don't update state based on it
+         results["parsed_code_blocks"] = list(latest_blocks_dict.values()) # Return existing state
+         results["default_selected_filenames"] = [b["filename"] for b in results["parsed_code_blocks"] if not b.get("is_structure_block")]
+         # results["error_message"] = "No bot message found in the last entry for parsing." # Optional: add a debug message
+         return results
+
+
+    role, content = message_obj.get("role", "").lower(), message_obj.get("content", "")
+
     file_pattern = re.compile(r"### File:\s*(?P<filename_line>[^\n]+)\n(?:```(?P<lang>[\w\.\-\+]*)\n(?P<code>[\s\S]*?)\n```|(?P<binary_msg>\[Binary file(?: - [^\]]+)?\]))")
     structure_pattern = re.compile(r"## File Structure\n```(?:(?P<struct_lang>[\w.-]*)\n)?(?P<structure_code>[\s\S]*?)\n```")
 
-    for message_obj in ai_chat_history:
-        if not isinstance(message_obj, dict): continue
-        role, content = message_obj.get("role", "").lower(), message_obj.get("content", "")
-        if role == BOT_ROLE_NAME:
-            structure_match = structure_pattern.search(content)
-            if structure_match:
-                latest_blocks_dict["File Structure (original)"] = {"filename": "File Structure (original)", "language": structure_match.group("struct_lang") or "plaintext", "code": structure_match.group("structure_code").strip(), "is_binary": False, "is_structure_block": True}
-            for match in file_pattern.finditer(content):
-                filename = _clean_filename(match.group("filename_line"))
-                if not filename: continue
-                lang, code_block, binary_msg = match.group("lang"), match.group("code"), match.group("binary_msg")
-                item_data = {"filename": filename, "is_binary": False, "is_structure_block": False}
-                if code_block is not None:
-                    item_data["code"], item_data["language"] = code_block.strip(), (lang.strip().lower() if lang else _infer_lang_from_filename(filename))
-                elif binary_msg is not None:
-                    item_data["code"], item_data["language"], item_data["is_binary"] = binary_msg.strip(), "binary", True
-                else: continue
-                latest_blocks_dict[filename] = item_data
+    # Process only the *last* bot message for updates to file blocks
+    if role == BOT_ROLE_NAME:
+        structure_match = structure_pattern.search(content)
+        if structure_match:
+            # Overwrite or add the structure block from the latest response
+            latest_blocks_dict["File Structure (original)"] = {"filename": "File Structure (original)", "language": structure_match.group("struct_lang") or "plaintext", "code": structure_match.group("structure_code").strip(), "is_binary": False, "is_structure_block": True}
+        
+        # Find all file blocks in the latest message
+        current_message_file_blocks = {}
+        for match in file_pattern.finditer(content):
+            filename = _clean_filename(match.group("filename_line"))
+            if not filename: continue
+            lang, code_block, binary_msg = match.group("lang"), match.group("code"), match.group("binary_msg")
+            item_data = {"filename": filename, "is_binary": False, "is_structure_block": False}
+            if code_block is not None:
+                item_data["code"], item_data["language"] = code_block.strip(), (lang.strip().lower() if lang else _infer_lang_from_filename(filename))
+            elif binary_msg is not None:
+                item_data["code"], item_data["language"], item_data["is_binary"] = binary_msg.strip(), "binary", True
+            else: continue
+            current_message_file_blocks[filename] = item_data
+
+        # Update latest_blocks_dict with blocks from the current message
+        # Any file mentioned in the latest message replaces its old version
+        # Any file *not* mentioned in the latest message but present previously is kept
+        # If the AI response implies deletion (e.g., by omitting a file it previously listed),
+        # the user needs to signal this intention in a way the prompt covers (e.g., "omit it from your next full ### File: list").
+        # Our current logic keeps files unless explicitly removed from the *entire* chat history's final list of blocks.
+        # A more sophisticated approach might look at AI language indicating deletion, but simple omission
+        # in a single turn's output isn't a reliable delete signal with complex interaction.
+        # For this implementation, we'll just update based on the latest message's defined blocks.
+        latest_blocks_dict.update(current_message_file_blocks)
+
 
     current_parsed_blocks = list(latest_blocks_dict.values())
     current_parsed_blocks.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
@@ -152,9 +189,15 @@ def _parse_chat_stream_logic(chat_json_string, existing_files_state=None):
 def _export_selected_logic(selected_filenames, space_line_name_for_md, parsed_blocks_for_export):
     results = {"output_str": "", "error_message": None, "download_filepath": None}
     global parsed_code_blocks_state_cache
-    if not parsed_blocks_for_export: parsed_blocks_for_export = parsed_code_blocks_state_cache
+    # Use parsed_blocks_for_export if provided, otherwise use cache
+    if parsed_blocks_for_export is None:
+        parsed_blocks_for_export = parsed_code_blocks_state_cache
 
-    if not selected_filenames and not any(b.get("is_structure_block") for b in parsed_blocks_for_export):
+    # Filter out structure blocks and binary/error files for file listing/export
+    exportable_blocks = [b for b in parsed_blocks_for_export if not b.get("is_structure_block") and not b.get("is_binary") and not (b.get("code", "").startswith("[Error loading content:") or b.get("code", "").startswith("[Binary or Skipped file]"))]
+    binary_blocks = [b for b in parsed_blocks_for_export if b.get("is_binary") or b.get("code", "").startswith("[Binary or Skipped file]")]
+
+    if not exportable_blocks and not binary_blocks and not any(b.get("is_structure_block") for b in parsed_blocks_for_export):
         results["output_str"] = f"# Space: {space_line_name_for_md}\n## File Structure\n{bbb}\n📁 Root\n{bbb}\n\n*No files to list in structure or export.*"
         try:
             with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
@@ -162,27 +205,43 @@ def _export_selected_logic(selected_filenames, space_line_name_for_md, parsed_bl
         except Exception as e: print(f"Error creating temp file for empty export: {e}")
         return results
 
-    output_lines = [f"# Space: {space_line_name_for_md}", "## File Structure", bbb, "📁 Root"]
-    filenames_for_structure = []
+    output_lines = [f"# Space: {space_line_name_for_md}"]
+
+    # Add File Structure block if it exists in parsed blocks
+    structure_block = next((b for b in parsed_blocks_for_export if b.get("is_structure_block")), None)
+    if structure_block:
+        output_lines.extend(["## File Structure", bbb, structure_block["code"].strip(), bbb, ""])
+    else:
+        # If no structure block from AI, generate a simple one from detected files
+        output_lines.extend(["## File Structure", bbb, "📁 Root"])
+        filenames_for_structure_list = sorted(list(set(b["filename"] for b in exportable_blocks + binary_blocks))) # Include binary files in structure list
+        if filenames_for_structure_list:
+            for fname in filenames_for_structure_list: output_lines.append(f"  📄 {fname}")
+        output_lines.extend([bbb, ""])
+
+    output_lines.append("Below are the contents of all files in the space:\n")
+    exported_content = False
+
+    # Determine which files to export content for based on selection or default
+    files_to_export_content = []
     if selected_filenames:
-        filenames_for_structure = sorted(list(set(f for f in selected_filenames if not any(b["filename"] == f and b.get("is_structure_block") for b in parsed_blocks_for_export))))
+        files_to_export_content = [b for b in parsed_blocks_for_export if b["filename"] in selected_filenames and not b.get("is_structure_block")]
     else:
-        filenames_for_structure = sorted(list(set(b["filename"] for b in parsed_blocks_for_export if not b.get("is_structure_block") and not b.get("is_binary"))))
+        files_to_export_content = [b for b in parsed_blocks_for_export if not b.get("is_structure_block")]
 
-    for fname in filenames_for_structure: output_lines.append(f"📄 {fname}")
-    output_lines.extend([bbb, "\nBelow are the contents of all files in the space:\n"])
-    exported_content = False
-    for block in parsed_blocks_for_export:
-        is_selected_for_content = (selected_filenames and block["filename"] in selected_filenames) or \
-                                  (not selected_filenames and not block.get("is_structure_block") and not block.get("is_binary"))
-        if is_selected_for_content and not block.get("is_structure_block"):
-            output_lines.append(f"### File: {block['filename']}")
-            if block.get('is_binary'): output_lines.append(block['code'])
-            else: output_lines.extend([f"{bbb}{block.get('language', 'plaintext') or 'plaintext'}", block.get('code',''), bbb])
-            output_lines.append(""); exported_content = True
-
-    if not exported_content and not filenames_for_structure : output_lines.append("*No file content selected or available for export.*")
-    elif not exported_content and filenames_for_structure:  output_lines.append("*Selected files have no content blocks defined by AI or are binary.*")
+    # Sort blocks for consistent output order
+    files_to_export_content.sort(key=lambda b: (0, b["filename"]) if b.get("is_binary") else (1, b["filename"]))
+
+    for block in files_to_export_content:
+        output_lines.append(f"### File: {block['filename']}")
+        if block.get('is_binary') or block.get("code", "").startswith("["): # Treat errors/skipped as binary for output format
+            output_lines.append(block.get('code',''))
+        else:
+            output_lines.extend([f"{bbb}{block.get('language', 'plaintext') or 'plaintext'}", block.get('code',''), bbb])
+        output_lines.append(""); exported_content = True
+
+    if not exported_content and not filenames_for_structure_list: output_lines.append("*No file content selected or available for export.*")
+    elif not exported_content and filenames_for_structure_list:  output_lines.append("*Selected files have no content blocks defined by AI, are binary, or encountered loading errors.*")
     final_output_str = "\n".join(output_lines)
     results["output_str"] = final_output_str
     try:
@@ -195,20 +254,30 @@ def _convert_gr_history_to_api_messages(system_prompt, gr_history, current_user_
     messages = [{"role": "system", "content": system_prompt}] if system_prompt else []
     for user_msg, bot_msg in gr_history:
         if user_msg: messages.append({"role": "user", "content": user_msg})
-        if bot_msg: messages.append({"role": BOT_ROLE_NAME, "content": bot_msg})
+        # Ensure bot_msg is not None or empty before adding
+        if bot_msg and isinstance(bot_msg, str): messages.append({"role": BOT_ROLE_NAME, "content": bot_msg})
     if current_user_message: messages.append({"role": "user", "content": current_user_message})
     return messages
 
 def get_latest_bot_message_as_json(gr_history):
-    if not gr_history or not gr_history[-1] or len(gr_history[-1]) < 2 or not gr_history[-1]: # Check bot message at index 1
-        return json.dumps([])
-    return json.dumps([{"role": BOT_ROLE_NAME, "content": gr_history[-1]}], indent=2)
+    # Find the last non-empty bot message
+    latest_bot_msg_content = None
+    for user_msg, bot_msg in reversed(gr_history):
+        if bot_msg and isinstance(bot_msg, str):
+            latest_bot_msg_content = bot_msg
+            break # Found the latest, stop searching
+
+    if latest_bot_msg_content is None:
+        return json.dumps([]) # Return empty list if no bot message found
+
+    # Return JSON containing only the latest bot message for parsing
+    return json.dumps([{"role": BOT_ROLE_NAME, "content": latest_bot_msg_content}], indent=2)
 
 
 def _generate_ui_outputs_from_cache(owner, space_name):
     global parsed_code_blocks_state_cache
     preview_md_val = "*No files in cache to display.*"
-    formatted_md_val = f"# Space: {owner}/{space_name}\n## File Structure\n{bbb}\n📁 Root\n{bbb}\n\n*No files in cache.*"
+    formatted_md_val = f"# Space: {owner}/{space_name}\n## File Structure\n{bbb}\n📁 Root\n{bbb}\n\n*No files in cache.*" if owner or space_name else "*Load or define a Space to see its Markdown structure.*"
     download_file = None
 
     if parsed_code_blocks_state_cache:
@@ -217,211 +286,473 @@ def _generate_ui_outputs_from_cache(owner, space_name):
             preview_md_lines.append(f"\n----\n**File:** `{escape_html_for_markdown(block['filename'])}`")
             if block.get('is_structure_block'): preview_md_lines.append(f" (Original File Structure from AI)\n")
             elif block.get('is_binary'): preview_md_lines.append(f" (Binary File)\n")
-            elif block.get('language'): preview_md_lines.append(f" (Language: `{block['language']}`)\n")
+            elif block.get('language') and block.get('language') != 'binary': preview_md_lines.append(f" (Language: `{block['language']}`)\n")
             else: preview_md_lines.append("\n")
-            if block.get('is_binary'): preview_md_lines.append(f"\n`{escape_html_for_markdown(block.get('code',''))}`\n")
-            else: preview_md_lines.append(f"\n{bbb}{block.get('language', 'plaintext') or 'plaintext'}\n{block.get('code','')}\n{bbb}\n")
+
+            # Handle content display
+            content = block.get('code', '')
+            if block.get('is_binary') or content.startswith("[") : # Treat errors/skipped as binary for preview display
+                 preview_md_lines.append(f"\n`{escape_html_for_markdown(content)}`\n")
+            elif block.get('is_structure_block'):
+                 preview_md_lines.append(f"\n{bbb}{block.get('language', 'plaintext') or 'plaintext'}\n{content}\n{bbb}\n")
+            else:
+                 preview_md_lines.append(f"\n{bbb}{block.get('language', 'plaintext') or 'plaintext'}\n{content}\n{bbb}\n")
+
+
         preview_md_val = "\n".join(preview_md_lines)
         space_line_name = f"{owner}/{space_name}" if owner and space_name else (owner or space_name or "your-space")
-        valid_filenames_for_export = [b["filename"] for b in parsed_code_blocks_state_cache if not b.get("is_structure_block") and not b.get("is_binary") and not (b.get("code", "").startswith("[Error loading content:") or b.get("code", "").startswith("[Binary or Skipped file]"))]
-        export_result = _export_selected_logic(valid_filenames_for_export, space_line_name, parsed_code_blocks_state_cache)
+
+        # _export_selected_logic handles selecting which files to include in the export MD
+        export_result = _export_selected_logic(None, space_line_name, parsed_code_blocks_state_cache) # Passing None means export all non-structure/non-binary/non-error content
         formatted_md_val = export_result["output_str"]
         download_file = export_result["download_filepath"]
-    else:
-        preview_md_val = "No files loaded or detected yet."
-        formatted_md_val = f"# Space: {owner or 'owner'}/{space_name or 'space'}\n## File Structure\n{bbb}\n📁 Root\n{bbb}\n\n*No files found or defined for this Space.*" if owner and space_name else "*Load or define a Space to see its Markdown structure.*"
+
     return formatted_md_val, preview_md_val, gr.update(value=download_file, interactive=download_file is not None)
 
 def handle_groq_chat_submit(user_message, chat_history, api_key_input, model_select, system_prompt, hf_owner_name, hf_repo_name, _current_formatted_markdown):
     global parsed_code_blocks_state_cache
     _chat_msg_in, _chat_hist, _status = "", list(chat_history), "Initializing..."
     _detected_files_update, _formatted_output_update, _download_btn_update = gr.update(), gr.update(), gr.update(interactive=False, value=None)
+
+    # --- Before sending to AI: Parse existing files from the current formatted markdown ---
+    # This ensures the AI is aware of the *current* state including user edits
+    # or files loaded from HF, before it generates its response.
+    # Only do this on a new user message
+    if user_message and _current_formatted_markdown:
+        try:
+            parsed_from_md = build_logic_parse_markdown(_current_formatted_markdown)
+            # Update cache only with files parsed from the markdown, preserving structure block if any
+            structure_block = next((b for b in parsed_code_blocks_state_cache if b.get("is_structure_block")), None)
+            parsed_code_blocks_state_cache = []
+            if structure_block:
+                parsed_code_blocks_state_cache.append(structure_block)
+
+            for f_info in parsed_from_md.get("files", []):
+                 # Only add if it has content and isn't the structure block representation
+                 if f_info.get("path") and f_info.get("path") != "File Structure (original)":
+                      # Check if it's a binary representation string
+                      is_binary_repr = isinstance(f_info.get("content"), str) and (f_info["content"].startswith("[Binary file") or f_info["content"].startswith("[Error loading content:") or f_info["content"].startswith("[Binary or Skipped file]"))
+                      parsed_code_blocks_state_cache.append({
+                           "filename": f_info["path"],
+                           "code": f_info.get("content", ""),
+                           "language": "binary" if is_binary_repr else _infer_lang_from_filename(f_info["path"]),
+                           "is_binary": is_binary_repr,
+                           "is_structure_block": False
+                      })
+            parsed_code_blocks_state_cache.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
+            # print(f"Parsed {len(parsed_code_blocks_state_cache)} blocks from current Markdown.")
+            # print(parsed_code_blocks_state_cache)
+
+        except Exception as e:
+            # Log error but don't block chat submission
+            print(f"Error parsing formatted markdown before chat submit: {e}")
+            # Optionally update status: _status = f"Warning: Error parsing current files: {e}. Sending message anyway."
+
+    # --- End of pre-chat parsing ---
+
+
     if not user_message.strip():
         _status = "Cannot send an empty message."
         yield (user_message, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update); return
     _chat_hist.append((user_message, None)); _status = "Sending to AI..."
     yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
+
     effective_api_key = api_key_input or os.environ.get("GROQ_API_KEY")
     if not effective_api_key:
         _chat_hist[-1] = (user_message, "Error: Groq API Key not set."); _status = "API Key missing."
         yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update); return
+
     current_sys_prompt = system_prompt.strip() or DEFAULT_SYSTEM_PROMPT
-    api_msgs = _convert_gr_history_to_api_messages(current_sys_prompt, _chat_hist[:-1], user_message)
-    headers, payload = {"Authorization": f"Bearer {effective_api_key}", "Content-Type": "application/json"}, {"model": model_select, "messages": api_msgs}
+
+    # Include current file contents in the prompt as context
+    current_files_context = ""
+    if parsed_code_blocks_state_cache:
+        current_files_context = "\n\n## Current Files in Space\n"
+        for block in parsed_code_blocks_state_cache:
+            if block.get("is_structure_block"):
+                 current_files_context += f"### File: {block['filename']}\n{bbb}\n{block['code']}\n{bbb}\n"
+            else:
+                current_files_context += f"### File: {block['filename']}\n"
+                if block.get("is_binary"):
+                    current_files_context += f"{block['code']}\n" # e.g. [Binary file...]
+                else:
+                    current_files_context += f"{bbb}{block.get('language', 'plaintext') or 'plaintext'}\n{block.get('code','')}\n{bbb}\n"
+        current_files_context += "\n"
+
+    # Append current file context to the user message
+    user_message_with_context = user_message.strip()
+    if current_files_context.strip():
+         user_message_with_context = user_message_with_context + current_files_context + "Based on the current files above and our chat history, please provide updated file contents using the `### File: ...\n```...\n```\n` format for any files you are creating, modifying, or want to include in the final output. Omit files you want to delete from your response."
+
+
+    api_msgs = _convert_gr_history_to_api_messages(current_sys_prompt, _chat_hist[:-1], user_message_with_context)
+
+    headers = {"Authorization": f"Bearer {effective_api_key}", "Content-Type": "application/json"}
+    payload = {"model": model_select, "messages": api_msgs}
+
     try:
         _status = f"Waiting for {model_select}..."; yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
-        response = requests.post(GROQ_API_ENDPOINT, headers=headers, json=payload, timeout=180); response.raise_for_status()
+        response = requests.post(GROQ_API_ENDPOINT, headers=headers, json=payload, timeout=180)
+        response.raise_for_status()
         api_resp_json = response.json()
-        
+
         bot_response_actual = None
-        if api_resp_json.get("choices") and len(api_resp_json["choices"]) > 0:
-            message_obj = api_resp_json["choices"].get("message")
-            if message_obj and isinstance(message_obj, dict): # Check if message_obj is a dict
-                 bot_response_actual = message_obj.get("content")
+        # --- Corrected logic to safely access the message content ---
+        if api_resp_json and isinstance(api_resp_json, dict) and api_resp_json.get("choices") and isinstance(api_resp_json["choices"], list) and len(api_resp_json["choices"]) > 0:
+            first_choice = api_resp_json["choices"][0] # Access the first item in the 'choices' list
+            if isinstance(first_choice, dict) and first_choice.get("message"): # Ensure it's a dict and has 'message' key
+                message_obj = first_choice["message"] # Get the 'message' dictionary from the choice
+                if isinstance(message_obj, dict): # Ensure message_obj is a dictionary
+                    bot_response_actual = message_obj.get("content") # Get content from the message dictionary
+        # --- End corrected logic ---
 
         if bot_response_actual:
+            # If bot_response_actual was successfully extracted
             _chat_hist[-1] = (user_message, bot_response_actual); _status = "AI response received. Processing files..."
             yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
-            latest_bot_message_json = get_latest_bot_message_as_json(_chat_hist)
+
+            # Pass the *existing* state (which was updated from markdown at the start of the function)
+            # to the parsing logic, so it knows about files already present.
+            latest_bot_message_json = json.dumps([{"role": BOT_ROLE_NAME, "content": bot_response_actual}], indent=2)
             parsing_res = _parse_chat_stream_logic(latest_bot_message_json, existing_files_state=parsed_code_blocks_state_cache)
+
             if parsing_res["error_message"]:
                 _status = f"Parsing Error: {parsing_res['error_message']}"
+                # Append parsing error to the bot's response in chat? Or show separately?
+                # For now, update status and detected files area with error
                 _detected_files_update = gr.Markdown(f"## Parsing Error\n`{escape_html_for_markdown(parsing_res['error_message'])}`")
             else:
+                # parsed_code_blocks_state_cache is updated inside _parse_chat_stream_logic
                 _formatted_output_update, _detected_files_update, _download_btn_update = _generate_ui_outputs_from_cache(hf_owner_name, hf_repo_name)
                 _status = "Processing complete. Previews updated."
+
             yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update); return
         else:
+            # Handle cases where bot_response_actual is None (e.g., malformed API response)
             error_msg = f"API Error: Unexpected response structure or empty message content. Response: {json.dumps(api_resp_json)}"
-        
-        _chat_hist[-1] = (user_message, error_msg); _status = error_msg
+            # This error_msg will be handled by the final error block below
+
     except requests.exceptions.HTTPError as e: error_msg = f"API HTTP Error: {e} - {e.response.text if e.response else 'No details'}"
     except requests.exceptions.RequestException as e: error_msg = f"API Request Error: {e}"
     except Exception as e: error_msg = f"Unexpected error in chat submit: {e}"
-    
-    if _chat_hist and len(_chat_hist[-1]) == 2 and _chat_hist[-1] is None:
-        _chat_hist[-1] = (user_message, error_msg)
+
+    # If an error occurred (and bot_response_actual was not set or an exception happened),
+    # update chat history with the error message.
+    # Assumes the last entry is the one for the current turn, added as (user_message, None).
+    # It's safer to append a new message if the last one isn't the one we expected to update.
+    if _chat_hist and len(_chat_hist) > 0 and _chat_hist[-1][1] is None:
+         _chat_hist[-1] = (_chat_hist[-1][0], error_msg) # Keep user message, add error as bot message
+    else:
+         _chat_hist.append((user_message, error_msg)) # Append as a new user/bot turn if structure unexpected
+
     _status = error_msg
     yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
 
+
 def handle_load_existing_space(hf_api_key_ui, ui_owner_name, ui_space_name):
     global parsed_code_blocks_state_cache
     _formatted_md_val, _detected_preview_val, _status_val = "*Loading files...*", "*Loading files...*", f"Loading Space: {ui_owner_name}/{ui_space_name}..."
     _file_browser_update, _iframe_html_update, _download_btn_update = gr.update(visible=False, choices=[], value=None), gr.update(value=None, visible=False), gr.update(interactive=False, value=None)
     _build_status_clear, _edit_status_clear, _runtime_status_clear = "*Build status will appear here.*", "*Select a file to load or delete.*", "*Space runtime status will appear here after refresh.*"
+
+    # Yield initial state to update UI
     yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, gr.update(value=ui_owner_name), gr.update(value=ui_space_name), _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+
     owner_to_use, updated_owner_name_val = ui_owner_name, ui_owner_name
+    error_occurred = False
+
     if not owner_to_use:
         token, token_err = build_logic_get_api_token(hf_api_key_ui)
-        if token_err or not token: _status_val = f"Error: {token_err or 'Cannot determine owner.'}"
+        if token_err or not token:
+            _status_val = f"Error: {token_err or 'Cannot determine owner from token.'}"; error_occurred = True
         else:
             try:
                 user_info = build_logic_whoami(token=token)
-                if user_info and 'name' in user_info: owner_to_use, updated_owner_name_val = user_info['name'], user_info['name']; _status_val += f" (Auto-detected owner: {owner_to_use})"
-                else: _status_val = "Error: Could not auto-detect owner."
-            except Exception as e: _status_val = f"Error auto-detecting owner: {e}"
-        if "Error:" in _status_val: yield ("*Error*", "*Error*", _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear); return
+                if user_info and 'name' in user_info:
+                    owner_to_use, updated_owner_name_val = user_info['name'], user_info['name']; _status_val += f" (Auto-detected owner: {owner_to_use})"
+                else:
+                    _status_val = "Error: Could not auto-detect owner from token."; error_occurred = True
+            except Exception as e:
+                _status_val = f"Error auto-detecting owner: {e}"; error_occurred = True
+
     if not owner_to_use or not ui_space_name:
-        _status_val = "Error: Owner and Space Name are required."
-        yield ("*Error*", "*Error*", _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear); return
+        if not error_occurred: _status_val = "Error: Owner and Space Name are required."; error_occurred = True
+
+    if error_occurred:
+         # Yield error state
+         yield (f"*Error: {_status_val}*", f"*Error: {_status_val}*", _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+         return # Stop execution
 
     sdk_for_iframe, file_list, err_list_files = get_space_repository_info(hf_api_key_ui, ui_space_name, owner_to_use)
-    sub_owner, sub_repo = re.sub(r'[^a-z0-9\-]+', '-', owner_to_use.lower()).strip('-'), re.sub(r'[^a-z0-9\-]+', '-', ui_space_name.lower()).strip('-')
+
+    # Construct iframe URL early, even if file listing fails
+    sub_owner = re.sub(r'[^a-z0-9\-]+', '-', owner_to_use.lower()).strip('-') or 'owner' # Fallback owner
+    sub_repo = re.sub(r'[^a-z0-9\-]+', '-', ui_space_name.lower()).strip('-') or 'space' # Fallback repo
     iframe_url = f"https://{sub_owner}-{sub_repo}{'.static.hf.space' if sdk_for_iframe == 'static' else '.hf.space'}"
     _iframe_html_update = gr.update(value=f'<iframe src="{iframe_url}?__theme=light&embed=true" width="100%" height="500px" style="border:1px solid #eee; border-radius:8px;"></iframe>', visible=True)
+
+
     if err_list_files and not file_list:
         _status_val = f"File List Error: {err_list_files}"
-        yield (f"*Error: {err_list_files}*", "*Error*", _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear); return
+        parsed_code_blocks_state_cache = [] # Clear cache on error
+        _formatted_md_val, _detected_preview_val, _download_btn_update = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name)
+        _file_browser_update = gr.update(visible=True, choices=[], value="Error loading files") # Update file browser with error state
+        yield (f"*Error: {err_list_files}*", "*Error loading files*", _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+        return # Stop execution
+
     if not file_list:
-        _status_val = f"Loaded Space: {owner_to_use}/{ui_space_name}. No files found ({err_list_files or 'No files'})."
+        _status_val = f"Loaded Space: {owner_to_use}/{ui_space_name}. No files found ({err_list_files or 'Repository is empty'})."
         parsed_code_blocks_state_cache = []
         _formatted_md_val, _detected_preview_val, _download_btn_update = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name)
         _file_browser_update = gr.update(visible=True, choices=[], value="No files found")
-        yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear); return
+        yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+        return # Stop execution
+
 
     loaded_files_for_parse = []
-    _status_val = f"Loading {len(file_list)} files from {owner_to_use}/{ui_space_name} (SDK: {sdk_for_iframe or 'unknown'})..."; yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+    _status_val = f"Loading {len(file_list)} files from {owner_to_use}/{ui_space_name} (SDK: {sdk_for_iframe or 'unknown'})...";
+    # Yield intermediate status while loading files
+    yield (_formatted_md_val, _detected_preview_val, _status_val, gr.update(visible=True, choices=sorted(file_list or []), value=None), updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+
+
     for file_path in file_list:
+        # Skip files that are likely binary or not user-editable code/text
+        # Added more extensions and common non-code files like lock files
         _, ext = os.path.splitext(file_path)
-        if ext.lower() in [".png",".jpg",".jpeg",".gif",".ico",".svg",".pt",".bin",".safetensors",".onnx",".woff",".woff2",".ttf",".eot"] or ".git" in file_path or (file_path.startswith(".") and file_path not in [".env", ".gitignore", ".gitattributes"]):
+        if ext.lower() in [".png",".jpg",".jpeg",".gif",".ico",".svg",".pt",".bin",".safetensors",".onnx",".woff",".woff2",".ttf",".eot",".zip",".tar",".gz",". هفت",".pdf",".mp4",".avi",".mov",".mp3",".wav",".ogg"] or \
+           file_path.startswith(".git") or "/.git/" in file_path or \
+           file_path in ["requirements.txt", "environment.yml", "setup.py", "Pipfile", "pyproject.toml", "package.json", "yarn.lock", "pnpm-lock.yaml", "poetry.lock"] or \
+           file_path.endswith(".lock") or \
+           file_path.startswith("__pycache__/") or "/__pycache__/" in file_path or \
+           file_path.startswith("node_modules/") or "/node_modules/" in file_path or \
+           file_path.startswith("venv/") or "/venv/" in file_path or \
+           file_path.startswith(".venv/") or "/.venv/" in file_path:
             loaded_files_for_parse.append({"filename": file_path, "code": "[Binary or Skipped file]", "language": "binary", "is_binary": True, "is_structure_block": False}); continue
-        content, err_get = get_space_file_content(hf_api_key_ui, ui_space_name, owner_to_use, file_path)
-        if err_get: loaded_files_for_parse.append({"filename": file_path, "code": f"[Error loading content: {err_get}]", "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False}); print(f"Error loading {file_path}: {err_get}"); continue
-        loaded_files_for_parse.append({"filename": file_path, "code": content, "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
+
+        # Handle potential issues with reading large files or non-utf8 files
+        try:
+            content, err_get = get_space_file_content(hf_api_key_ui, ui_space_name, owner_to_use, file_path)
+            if err_get:
+                 # If there's an error getting content, record it but don't fail the whole load
+                 loaded_files_for_parse.append({"filename": file_path, "code": f"[Error loading content: {err_get}]", "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
+                 print(f"Error loading {file_path}: {err_get}");
+                 continue
+            # If content is successfully loaded
+            loaded_files_for_parse.append({"filename": file_path, "code": content, "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
+        except Exception as content_ex:
+            # Catch any other unexpected exceptions during file content fetching
+            loaded_files_for_parse.append({"filename": file_path, "code": f"[Unexpected error loading content: {content_ex}]", "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
+            print(f"Unexpected error loading {file_path}: {content_ex}")
+            continue
+
+
     parsed_code_blocks_state_cache = loaded_files_for_parse
     _formatted_md_val, _detected_preview_val, _download_btn_update = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name)
     _status_val = f"Successfully loaded Space: {owner_to_use}/{ui_space_name}. Markdown ready."
-    _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), value=None)
+    _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), value=None) # Use the full file list for the dropdown
+
     yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
 
+
 def handle_build_space_button(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, space_sdk_ui, formatted_markdown_content):
     _build_status, _iframe_html, _file_browser_update = "Starting space build process...", gr.update(value=None, visible=False), gr.update(visible=False, choices=[], value=None)
-    yield _build_status, _iframe_html, _file_browser_update
+    yield _build_status, _iframe_html, _file_browser_update # Yield initial status
     if not ui_space_name_part or "/" in ui_space_name_part: _build_status = f"Build Error: HF Space Name '{ui_space_name_part}' must be repo name only (no '/')."; yield _build_status, _iframe_html, _file_browser_update; return
     final_owner_for_build = ui_owner_name_part
     if not final_owner_for_build:
-        token_for_whoami, _ = build_logic_get_api_token(hf_api_key_ui)
+        token_for_whoami, token_err = build_logic_get_api_token(hf_api_key_ui)
+        if token_err: _build_status = f"Build Error: {token_err}"; yield _build_status, _iframe_html, _file_browser_update; return
         if token_for_whoami:
-            try: user_info = build_logic_whoami(token=token_for_whoami); final_owner_for_build = user_info['name'] if user_info and 'name' in user_info else final_owner_for_build
-            except Exception as e: _build_status += f"\n(Warn: Could not auto-detect owner for build: {e})"
-        else: _build_status += "\n(Warn: Owner not specified and no token to auto-detect for build.)"
-    if not final_owner_for_build: _build_status = "Build Error: HF Owner Name could not be determined."; yield _build_status, _iframe_html, _file_browser_update; return
+            try:
+                user_info = build_logic_whoami(token=token_for_whoami)
+                final_owner_for_build = user_info['name'] if user_info and 'name' in user_info else final_owner_for_build
+                if not final_owner_for_build: _build_status += "\n(Warning: Could not auto-detect owner from token for build. Please specify.)"
+            except Exception as e: _build_status += f"\n(Warning: Could not auto-detect owner for build: {e}. Please specify.)"
+        else: _build_status += "\n(Warning: Owner not specified and no token to auto-detect for build. Please specify owner or provide a token.)"
+
+    if not final_owner_for_build: _build_status = "Build Error: HF Owner Name could not be determined. Please specify it."; yield _build_status, _iframe_html, _file_browser_update; return
+
     result_message = build_logic_create_space(hf_api_key_ui, ui_space_name_part, final_owner_for_build, space_sdk_ui, formatted_markdown_content)
     _build_status = f"Build Process: {result_message}"
+
     if "Successfully" in result_message:
-        sub_owner, sub_repo = re.sub(r'[^a-z0-9\-]+', '-', final_owner_for_build.lower()).strip('-'), re.sub(r'[^a-z0-9\-]+', '-', ui_space_name_part.lower()).strip('-')
+        sub_owner = re.sub(r'[^a-z0-9\-]+', '-', final_owner_for_build.lower()).strip('-') or 'owner'
+        sub_repo = re.sub(r'[^a-z0-9\-]+', '-', ui_space_name_part.lower()).strip('-') or 'space'
         iframe_url = f"https://{sub_owner}-{sub_repo}{'.static.hf.space' if space_sdk_ui == 'static' else '.hf.space'}"
         _iframe_html = gr.update(value=f'<iframe src="{iframe_url}?__theme=light&embed=true" width="100%" height="700px" style="border:1px solid #eee; border-radius:8px;"></iframe>', visible=True)
-        _build_status += f"\nSpace live at: {iframe_url} (Repo: https://huggingface.co/spaces/{final_owner_for_build}/{ui_space_name_part})"
+        _build_status += f"\nSpace live at: [Link]({iframe_url}) (Repo: https://huggingface.co/spaces/{final_owner_for_build}/{ui_space_name_part})"
+
+        # Refresh file list after successful build
         file_list, err_list = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, final_owner_for_build)
         if err_list: _build_status += f"\nFile list refresh error after build: {err_list}"; _file_browser_update = gr.update(visible=True, choices=[], value="Error refreshing files")
         else: _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), value=None if file_list else "No files found")
+
     yield _build_status, _iframe_html, _file_browser_update
 
+
 def handle_load_file_for_editing(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, selected_file_path):
-    _file_content_val, _edit_status_val, _commit_msg_val, _lang_update = "", "Error: No file selected.", gr.update(), gr.update()
-    if not selected_file_path or selected_file_path in ["No files found", "Error refreshing files"]: yield _file_content_val, _edit_status_val, _commit_msg_val, _lang_update; return
+    _file_content_val, _edit_status_val, _commit_msg_val, _lang_update = "", "Error: No file selected.", gr.update(value=""), gr.update(language="plaintext") # Reset values
+    if not selected_file_path or selected_file_path in ["No files found", "Error loading files", "Error refreshing files"]:
+         yield _file_content_val, "Select a file from the dropdown.", _commit_msg_val, _lang_update # Clear editor and status
+         return
+
     owner_to_use = ui_owner_name_part
     if not owner_to_use:
-        token, _ = build_logic_get_api_token(hf_api_key_ui)
+        token, token_err = build_logic_get_api_token(hf_api_key_ui)
+        if token_err: _edit_status_val = f"Error: {token_err}"; yield (_file_content_val, _edit_status_val, _commit_msg_val, _lang_update); return
         if token:
-            try: user_info = build_logic_whoami(token=token); owner_to_use = user_info['name'] if user_info and 'name' in user_info else owner_to_use
-            except Exception as e: print(f"Owner auto-detect fail for editing file: {e}")
+            try:
+                 user_info = build_logic_whoami(token=token); owner_to_use = user_info['name'] if user_info and 'name' in user_info else owner_to_use
+                 if not owner_to_use: _edit_status_val = "Error: Could not auto-detect owner from token."; yield (_file_content_val, _edit_status_val, _commit_msg_val, _lang_update); return
+            except Exception as e: _edit_status_val = f"Error auto-detecting owner for editing file: {e}"; yield (_file_content_val, _edit_status_val, _commit_msg_val, _lang_update); return
         else: _edit_status_val = "Error: HF Owner Name not set and no token to auto-detect."; yield (_file_content_val, _edit_status_val, _commit_msg_val, _lang_update); return
+
     if not owner_to_use or not ui_space_name_part: _edit_status_val = "Error: HF Owner and/or Space Name is missing."; yield (_file_content_val, _edit_status_val, _commit_msg_val, _lang_update); return
+
+    _edit_status_val = f"Loading {selected_file_path}..."
+    yield gr.update(value=""), _edit_status_val, gr.update(value=""), gr.update(language="plaintext") # Yield loading state
+
     content, err = get_space_file_content(hf_api_key_ui, ui_space_name_part, owner_to_use, selected_file_path)
-    if err: _edit_status_val, _commit_msg_val = f"Error loading '{selected_file_path}': {err}", f"Error loading {selected_file_path}"; yield (f"Error loading {selected_file_path}:\n{err}", _edit_status_val, _commit_msg_val, _lang_update); return
-    _file_content_val, _edit_status_val, _commit_msg_val, _lang_update = content or "", f"Loaded {selected_file_path} for editing.", f"Update {selected_file_path} via AI Space Editor", gr.update(language=_infer_lang_from_filename(selected_file_path))
+
+    if err:
+        _edit_status_val = f"Error loading '{selected_file_path}': {err}"
+        _commit_msg_val = f"Error loading {selected_file_path}"
+        _file_content_val = f"Error loading {selected_file_path}:\n{err}"
+        _lang_update = gr.update(language="plaintext")
+        yield _file_content_val, _edit_status_val, _commit_msg_val, _lang_update
+        return
+
+    _file_content_val = content or ""
+    _edit_status_val = f"Loaded {selected_file_path} for editing."
+    _commit_msg_val = f"Update {selected_file_path} via AI Space Editor"
+    _lang_update = gr.update(language=_infer_lang_from_filename(selected_file_path))
+
     yield _file_content_val, _edit_status_val, _commit_msg_val, _lang_update
 
 def handle_commit_file_changes(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, file_to_edit_path, edited_content, commit_message):
     global parsed_code_blocks_state_cache
-    _edit_status_val, _file_browser_update_val, _formatted_md_out, _detected_preview_out, _download_btn_out = "Error: No file selected.", gr.update(), gr.update(), gr.update(), gr.update()
-    if not file_to_edit_path or file_to_edit_path == "No files found": yield _edit_status_val, _file_browser_update_val, _formatted_md_out, _detected_preview_out, _download_btn_out; return
+    _edit_status_val = "Processing commit..."
+    _file_browser_update_val = gr.update()
+    _formatted_md_out, _detected_preview_out, _download_btn_out = gr.update(), gr.update(), gr.update()
+
+    yield _edit_status_val, _file_browser_update_val, _formatted_md_out, _detected_preview_out, _download_btn_out # Yield initial status
+
+    if not file_to_edit_path or file_to_edit_path in ["No files found", "Error loading files", "Error refreshing files"]:
+        _edit_status_val = "Error: No valid file selected for commit.";
+        yield _edit_status_val, gr.update(), gr.update(), gr.update(), gr.update(); return
+
     owner_to_use = ui_owner_name_part
     if not owner_to_use:
-        token, _ = build_logic_get_api_token(hf_api_key_ui)
+        token, token_err = build_logic_get_api_token(hf_api_key_ui)
+        if token_err: _edit_status_val = f"Error: {token_err}"; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update()); return
         if token:
-            try: user_info = build_logic_whoami(token=token); owner_to_use = user_info['name'] if user_info and 'name' in user_info else owner_to_use
-            except Exception as e: print(f"Owner auto-detect fail for committing file: {e}")
-        else: _edit_status_val = "Error: HF Owner Name not set and no token to auto-detect."; yield (_edit_status_val, _file_browser_update_val, _formatted_md_out, _detected_preview_out, _download_btn_out); return
-    if not owner_to_use or not ui_space_name_part: _edit_status_val = "Error: HF Owner and/or Space Name is missing."; yield (_edit_status_val, _file_browser_update_val, _formatted_md_out, _detected_preview_out, _download_btn_out); return
+            try:
+                 user_info = build_logic_whoami(token=token); owner_to_use = user_info['name'] if user_info and 'name' in user_info else owner_to_use
+                 if not owner_to_use: _edit_status_val = "Error: Could not auto-detect owner from token."; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update()); return
+            except Exception as e: _edit_status_val = f"Error auto-detecting owner for committing file: {e}"; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update()); return
+        else: _edit_status_val = "Error: HF Owner Name not set and no token to auto-detect."; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update()); return
+
+    if not owner_to_use or not ui_space_name_part: _edit_status_val = "Error: HF Owner and/or Space Name is missing."; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update()); return
+
     status_msg = update_space_file(hf_api_key_ui, ui_space_name_part, owner_to_use, file_to_edit_path, edited_content, commit_message)
     _edit_status_val = status_msg
+
     if "Successfully updated" in status_msg:
+        # Update the cache with the new content
         found_in_cache = False
         for block in parsed_code_blocks_state_cache:
-            if block["filename"] == file_to_edit_path: block["code"], block["language"], block["is_binary"], found_in_cache = edited_content, _infer_lang_from_filename(file_to_edit_path), False, True; break
-        if not found_in_cache: parsed_code_blocks_state_cache.append({"filename": file_to_edit_path, "code": edited_content, "language": _infer_lang_from_filename(file_to_edit_path), "is_binary": False, "is_structure_block": False}); parsed_code_blocks_state_cache.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
+            if block["filename"] == file_to_edit_path:
+                block["code"] = edited_content
+                block["language"] = _infer_lang_from_filename(file_to_edit_path)
+                block["is_binary"] = False # Assume user edited text content
+                block["is_structure_block"] = False # Ensure it's not marked as structure
+                found_in_cache = True
+                break
+        if not found_in_cache:
+             # If file was added via editor or wasn't in initial load cache (e.g. binary/error), add it
+             parsed_code_blocks_state_cache.append({
+                 "filename": file_to_edit_path,
+                 "code": edited_content,
+                 "language": _infer_lang_from_filename(file_to_edit_path),
+                 "is_binary": False,
+                 "is_structure_block": False
+             })
+             # Re-sort the cache to maintain consistent order
+             parsed_code_blocks_state_cache.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
+
+        # Regenerate markdown and preview from the updated cache
         _formatted_md_out, _detected_preview_out, _download_btn_out = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name_part)
-        _file_browser_update_val = gr.update(value=file_to_edit_path)
+
+        # Refresh file list choices and keep the current file selected
+        new_file_list, err_list = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, owner_to_use)
+        if err_list:
+             _edit_status_val += f"\nFile list refresh error: {err_list}"
+             _file_browser_update_val = gr.update(choices=sorted(new_file_list or []), value="Error refreshing files")
+        else:
+             _file_browser_update_val = gr.update(choices=sorted(new_file_list or []), value=file_to_edit_path) # Keep current file selected
+
     yield _edit_status_val, _file_browser_update_val, _formatted_md_out, _detected_preview_out, _download_btn_out
 
+
 def handle_delete_file(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, file_to_delete_path):
     global parsed_code_blocks_state_cache
-    _edit_status_val, _file_browser_choices_update, _file_browser_value_update, _file_content_editor_update = "Processing deletion...", gr.update(), gr.update(), gr.update(value="")
+    _edit_status_val = "Processing deletion..."
+    _file_browser_choices_update = gr.update()
+    _file_browser_value_update = None # Clear selected file
+    _file_content_editor_update = gr.update(value="") # Clear editor content
+    _commit_msg_update = gr.update(value="") # Clear commit message
+    _lang_update = gr.update(language="plaintext") # Reset editor language
     _formatted_md_out, _detected_preview_out, _download_btn_out = gr.update(), gr.update(), gr.update()
-    yield (_edit_status_val, _file_browser_choices_update, _file_browser_value_update, _file_content_editor_update, _formatted_md_out, _detected_preview_out, _download_btn_out)
-    if not file_to_delete_path or file_to_delete_path in ["No files found", "Error refreshing files"]: _edit_status_val = "Error: No file selected for deletion."; yield (_edit_status_val, gr.update(), gr.update(), _file_content_editor_update, gr.update(), gr.update(), gr.update()); return
+
+    yield (_edit_status_val, _file_browser_choices_update, _file_browser_value_update, _file_content_editor_update, _commit_msg_update, _lang_update, _formatted_md_out, _detected_preview_out, _download_btn_out) # Yield initial status
+
+
+    if not file_to_delete_path or file_to_delete_path in ["No files found", "Error loading files", "Error refreshing files"]:
+        _edit_status_val = "Error: No valid file selected for deletion.";
+        yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()); return
+
     owner_to_use = ui_owner_name_part
     if not owner_to_use:
-        token, _ = build_logic_get_api_token(hf_api_key_ui)
+        token, token_err = build_logic_get_api_token(hf_api_key_ui)
+        if token_err: _edit_status_val = f"API Token Error: {token_err}"; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()); return
         if token:
-            try: user_info = build_logic_whoami(token=token); owner_to_use = user_info['name'] if user_info and 'name' in user_info else owner_to_use
-            except Exception as e: print(f"Owner auto-detect fail for deleting file: {e}"); _edit_status_val = "Error: Could not auto-detect owner."; yield (_edit_status_val, gr.update(), gr.update(), _file_content_editor_update, gr.update(), gr.update(), gr.update()); return
-        else: _edit_status_val = "Error: HF Token needed to auto-detect owner."; yield (_edit_status_val, gr.update(), gr.update(), _file_content_editor_update, gr.update(), gr.update(), gr.update()); return
-    if not owner_to_use or not ui_space_name_part: _edit_status_val = "Error: Owner and Space Name are required."; yield (_edit_status_val, gr.update(), gr.update(), _file_content_editor_update, gr.update(), gr.update(), gr.update()); return
+            try:
+                 user_info = build_logic_whoami(token=token); owner_to_use = user_info['name'] if user_info and 'name' in user_info else owner_to_use
+                 if not owner_to_use: _edit_status_val = "Error: Could not auto-detect owner from token."; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()); return
+            except Exception as e: _edit_status_val = f"Error auto-detecting owner for deleting file: {e}"; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()); return
+        else: _edit_status_val = "Error: HF Token needed to auto-detect owner."; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()); return
+
+
+    if not owner_to_use or not ui_space_name_part: _edit_status_val = "Error: Owner and Space Name are required."; yield (_edit_status_val, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()); return
+
     deletion_status_msg = build_logic_delete_space_file(hf_api_key_ui, ui_space_name_part, owner_to_use, file_to_delete_path)
     _edit_status_val = deletion_status_msg
+
     if "Successfully deleted" in deletion_status_msg:
+        # Remove the file from the cache
         parsed_code_blocks_state_cache = [b for b in parsed_code_blocks_state_cache if b["filename"] != file_to_delete_path]
+
+        # Regenerate markdown and preview from the updated cache
         _formatted_md_out, _detected_preview_out, _download_btn_out = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name_part)
+
+        # Refresh file list choices and clear selected file
         new_file_list, err_list = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, owner_to_use)
-        if err_list: _edit_status_val += f"\nFile list refresh error: {err_list}"; _file_browser_choices_update = gr.update(choices=[], value="Error refreshing files")
-        else: _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value=None)
-        _file_browser_value_update = None
-    else: new_file_list, _ = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, owner_to_use); _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value=file_to_delete_path)
-    yield (_edit_status_val, _file_browser_choices_update, _file_browser_value_update, _file_content_editor_update, _formatted_md_out, _detected_preview_out, _download_btn_out)
+        if err_list:
+             _edit_status_val += f"\nFile list refresh error: {err_list}"
+             _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value="Error refreshing files")
+        else:
+             _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value=None) # Clear selection
+
+        _file_browser_value_update = None # Explicitly set value to None to clear selection visual
+
+    else:
+        # If deletion failed, refresh the file list but keep the file selected
+        new_file_list, _ = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, owner_to_use)
+        _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value=file_to_delete_path)
+        _file_browser_value_update = file_to_delete_path # Keep the file selected
+
+    yield (_edit_status_val, _file_browser_choices_update, _file_browser_value_update, _file_content_editor_update, _commit_msg_update, _lang_update, _formatted_md_out, _detected_preview_out, _download_btn_out)
+
 
 def handle_refresh_space_status(hf_api_key_ui, ui_owner_name, ui_space_name):
     yield "*Fetching space status...*" # Initial feedback
@@ -433,25 +764,29 @@ def handle_refresh_space_status(hf_api_key_ui, ui_owner_name, ui_space_name):
         except Exception as e: yield f"**Error auto-detecting owner:** {e}"; return
     if not owner_to_use or not ui_space_name: yield "**Error:** Owner and Space Name are required."; return
     status_details, error_msg = get_space_runtime_status(hf_api_key_ui, ui_space_name, owner_to_use)
-    if error_msg: _status_display_md = f"**Error fetching status for {owner_to_use}/{ui_space_name}:**\n\n`{error_msg}`"
+    if error_msg: _status_display_md = f"**Error fetching status for {owner_to_use}/{ui_space_name}:**\n\n`{escape_html_for_markdown(error_msg)}`"
     elif status_details:
         stage, hardware, error, log_link = status_details.get('stage','N/A'), status_details.get('hardware','N/A'), status_details.get('error_message'), status_details.get('full_log_link','#')
         md_lines = [f"### Space Status: {owner_to_use}/{ui_space_name}", f"- **Stage:** `{stage}`", f"- **Current Hardware:** `{hardware}`"]
         if status_details.get('requested_hardware') and status_details.get('requested_hardware') != hardware: md_lines.append(f"- **Requested Hardware:** `{status_details.get('requested_hardware')}`")
         if error: md_lines.append(f"- **Error:** <span style='color:red;'>`{escape_html_for_markdown(error)}`</span>")
         md_lines.append(f"- [View Full Logs on Hugging Face]({log_link})")
-        # md_lines.append(f"\n<details><summary>Raw Status Data (JSON)</summary>\n```json\n{json.dumps(status_details.get('raw_data', {}), indent=2)}\n```\n</details>")
+        if status_details.get('raw_data'):
+             # Add raw data in a collapsible section for debugging
+             md_lines.append(f"\n<details><summary>Raw Status Data (JSON)</summary>\n\n```json\n{json.dumps(status_details.get('raw_data', {}), indent=2)}\n```\n</details>")
+
         _status_display_md = "\n".join(md_lines)
     else: _status_display_md = "Could not retrieve status details."
     yield _status_display_md
 
+
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="orange")) as demo:
     gr.Markdown("# 🤖 AI Code & Space Generator")
     gr.Markdown("Configure settings, chat with AI to generate/modify Hugging Face Spaces, then build, preview, and edit.")
     with gr.Row():
         with gr.Sidebar():
             gr.Markdown("## ⚙️ Configuration")
-            with gr.Group(): gr.Markdown("### API Keys & Tokens"); groq_api_key_input = gr.Textbox(label="Groq API Key", type="password", placeholder="gsk_..."); hf_api_key_input = gr.Textbox(label="Hugging Face Token", type="password", placeholder="hf_...")
+            with gr.Group(): gr.Markdown("### API Keys & Tokens"); groq_api_key_input = gr.Textbox(label="Groq API Key", type="password", placeholder="gsk_...", allow_paste=True); hf_api_key_input = gr.Textbox(label="Hugging Face Token", type="password", placeholder="hf_...", allow_paste=True)
             with gr.Group(): gr.Markdown("### Hugging Face Space"); owner_name_input = gr.Textbox(label="HF Owner Name", placeholder="e.g., your-username"); space_name_input = gr.Textbox(label="HF Space Name", value="my-ai-space", placeholder="e.g., my-cool-app"); space_sdk_select = gr.Dropdown(label="Space SDK", choices=["gradio", "streamlit", "docker", "static"], value="gradio", info="Used for new/build."); load_space_button = gr.Button("🔄 Load Existing Space", variant="secondary", size="sm")
             with gr.Group(): gr.Markdown("### AI Model Settings"); groq_model_select = gr.Dropdown(label="Groq Model", choices=["mixtral-8x7b-32768", "llama3-8b-8192", "llama3-70b-8192", "gemma-7b-it"], value="llama3-8b-8192"); groq_system_prompt_input = gr.Textbox(label="System Prompt", lines=8, value=DEFAULT_SYSTEM_PROMPT, interactive=True)
         with gr.Column(scale=3):
@@ -462,12 +797,12 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="orange"))
             gr.Markdown("---")
             with gr.Tabs():
                 with gr.TabItem("📝 Formatted Space Markdown"): gr.Markdown("Complete Markdown definition for your Space."); formatted_space_output_display = gr.Textbox(label="Current Space Definition", lines=15, interactive=True, show_copy_button=True, value="*Space definition...*"); download_button = gr.DownloadButton(label="Download .md", interactive=False, size="sm")
-                with gr.TabItem("🔍 Detected Files Preview"): gr.Markdown("Preview of files and content from AI/loaded Space."); detected_files_preview = gr.Markdown(value="*Files preview...*")
+                with gr.TabItem("🔍 Detected Files Preview"): gr.Markdown(value="*Files preview...*")
             gr.Markdown("---")
             with gr.Tabs():
                 with gr.TabItem("🚀 Build & Preview Space"):
                     with gr.Row(): build_space_button = gr.Button("Build / Update Space on HF", variant="primary", scale=2); refresh_status_button = gr.Button("🔄 Refresh Space Status", scale=1)
-                    build_status_display = gr.Textbox(label="Build Operation Status", interactive=False, lines=2, value="*Build status...*"); gr.Markdown("---"); space_runtime_status_display = gr.Markdown("*Space runtime status...*"); gr.Markdown("---"); space_iframe_display = gr.HTML(value="<!-- Space Iframe -->", visible=False)
+                    build_status_display = gr.Textbox(label="Build Operation Status", interactive=False, lines=2, value="*Build status will appear here.*"); gr.Markdown("---"); space_runtime_status_display = gr.Markdown("*Space runtime status will appear here after refresh.*"); gr.Markdown("---"); space_iframe_display = gr.HTML(value="<!-- Space Iframe -->", visible=False)
                 with gr.TabItem("✏️ Edit Space Files"):
                     gr.Markdown("Select a file to view, edit, or delete. Changes are committed to HF Hub.")
                     file_browser_dropdown = gr.Dropdown(label="Select File in Space", choices=[], interactive=True, visible=False, info="Load/build Space first.")
@@ -477,19 +812,29 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="orange"))
                     edit_status_display = gr.Textbox(label="File Edit/Delete Status", interactive=False, lines=2, value="*Select file...*")
 
     chat_outputs = [groq_chat_message_input, groq_chatbot_display, groq_status_output, detected_files_preview, formatted_space_output_display, download_button]
-    chat_inputs = [groq_chat_message_input, groq_chatbot_display, groq_api_key_input, groq_model_select, groq_system_prompt_input, owner_name_input, space_name_input, formatted_space_output_display]
+    chat_inputs = [groq_chat_message_input, groq_chatbot_display, groq_api_key_input, groq_model_select, groq_system_prompt_input, owner_name_input, space_name_input, formatted_space_output_display] # Added formatted_space_output_display to inputs
     groq_send_chat_button.click(fn=handle_groq_chat_submit, inputs=chat_inputs, outputs=chat_outputs)
     groq_chat_message_input.submit(fn=handle_groq_chat_submit, inputs=chat_inputs, outputs=chat_outputs)
+
+    # Added outputs for clearing/updating edit status when loading a space
     load_space_outputs = [formatted_space_output_display, detected_files_preview, groq_status_output, file_browser_dropdown, owner_name_input, space_name_input, space_iframe_display, download_button, build_status_display, edit_status_display, space_runtime_status_display]
     load_space_button.click(fn=handle_load_existing_space, inputs=[hf_api_key_input, owner_name_input, space_name_input], outputs=load_space_outputs)
+
     build_outputs = [build_status_display, space_iframe_display, file_browser_dropdown]
     build_space_button.click(fn=handle_build_space_button, inputs=[hf_api_key_input, space_name_input, owner_name_input, space_sdk_select, formatted_space_output_display], outputs=build_outputs)
+
+    # Added commit message and language update to file edit load outputs
     file_edit_load_outputs = [file_content_editor, edit_status_display, commit_message_input, file_content_editor]
     file_browser_dropdown.change(fn=handle_load_file_for_editing, inputs=[hf_api_key_input, space_name_input, owner_name_input, file_browser_dropdown], outputs=file_edit_load_outputs)
+
+    # Added formatted_space_output_display, detected_files_preview, download_button to commit outputs to refresh previews
     commit_file_outputs = [edit_status_display, file_browser_dropdown, formatted_space_output_display, detected_files_preview, download_button]
     update_file_button.click(fn=handle_commit_file_changes, inputs=[hf_api_key_input, space_name_input, owner_name_input, file_browser_dropdown, file_content_editor, commit_message_input], outputs=commit_file_outputs)
-    delete_file_outputs = [edit_status_display, file_browser_dropdown, file_browser_dropdown, file_content_editor, formatted_space_output_display, detected_files_preview, download_button]
+
+    # Added file_content_editor, commit_message_input, language update, formatted_space_output_display, detected_files_preview, download_button to delete outputs
+    delete_file_outputs = [edit_status_display, file_browser_dropdown, file_browser_dropdown, file_content_editor, commit_message_input, file_content_editor, formatted_space_output_display, detected_files_preview, download_button]
     delete_file_button.click(fn=handle_delete_file, inputs=[hf_api_key_input, space_name_input, owner_name_input, file_browser_dropdown], outputs=delete_file_outputs)
+
     refresh_status_button.click(fn=handle_refresh_space_status, inputs=[hf_api_key_input, owner_name_input, space_name_input], outputs=[space_runtime_status_display])
 
 if __name__ == "__main__":