Update app.py

app.py

@@ -1,380 +1,1128 @@
-# keylock/app.py
+# app.py
 import gradio as gr
-from PIL import Image, ImageFont
-import tempfile
-import os
+import re
 import json
-import logging
-import traceback
-import base64
-import io
-
-import core # Use relative import for core module
-import __version__ # Import version for footer
-
-app_logger = logging.getLogger("keylock_app")
-if not app_logger.hasHandlers(): # Basic logging setup if not configured elsewhere
-    handler = logging.StreamHandler()
-    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-    handler.setFormatter(formatter)
-    app_logger.addHandler(handler)
-    app_logger.setLevel(logging.INFO)
-
-# Theming
-try:
-    font_family = [gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"]
-except AttributeError: 
-    app_logger.warning("gr.themes.GoogleFont not found. Using fallback fonts. This might be due to Gradio version.")
-    font_family = ["ui-sans-serif", "system-ui", "sans-serif"]
-
-try:
-    blue_color = gr.themes.colors.blue
-    sky_color = gr.themes.colors.sky
-    slate_color = gr.themes.colors.slate
-    cyan_color = gr.themes.colors.cyan
-    neutral_color = gr.themes.colors.neutral
-except AttributeError:
-    app_logger.warning("gr.themes.colors not found. Using placeholder colors for themes. This might be due to Gradio version.")
-    class FallbackColors: # Basic fallback colors
-        blue = "blue"; sky = "skyblue"; slate = "slategray"; cyan = "cyan"; neutral = "gray"
-    blue_color = FallbackColors.blue
-    sky_color = FallbackColors.sky
-    slate_color = FallbackColors.slate
-    cyan_color = FallbackColors.cyan
-    neutral_color = FallbackColors.neutral
-
-ICON_EMBED = "➕"
-ICON_EXTRACT = "➖"
-
-def pil_to_base64_html(pil_image, max_width_px=None):
-    buffered = io.BytesIO(); pil_image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    style = f"max-width:{max_width_px}px; height:auto; border:1px solid #ccc; display:block; margin-left:auto; margin-right:auto;" if max_width_px else "border:1px solid #ccc; display:block; margin-left:auto; margin-right:auto;"
-    return f"<div style='text-align:center;'><img src='data:image/png;base64,{img_str}' alt='Stego Image' style='{style}'/></div>"
-
-def gradio_embed_data(kv_string: str, password: str, 
-                      input_image_pil: Image.Image, generate_carrier_flag: bool,
-                      show_keys_on_image_flag: bool, output_filename_base: str):
-    output_html_img_str, status_msg, dl_file_path = None, "An error occurred.", None
-    if not password: return None, "Error: Password cannot be empty.", None
-    if not kv_string or not kv_string.strip(): return None, "Error: Key-Value data cannot be empty.", None
-    try:
-        data_dict = core.parse_kv_string_to_dict(kv_string)
-        if not data_dict: return None, "Error: Parsed Key-Value data is empty.", None
-        
-        original_format_note = ""
-        if generate_carrier_flag or input_image_pil is None:
-            carrier_img = core.generate_keylock_carrier_image()
-        else:
-            carrier_img = input_image_pil.copy()
-            if hasattr(input_image_pil, 'format') and input_image_pil.format and input_image_pil.format.upper() != 'PNG':
-                original_format_note = (
-                    f"Input carrier image was format '{input_image_pil.format}'. "
-                    f"It will be processed and saved as PNG. "
-                )
-                app_logger.warning(
-                    f"{original_format_note}If original was lossy (e.g., JPEG), quality is preserved from upload; "
-                    f"if it had transparency (e.g., GIF), it will be lost during RGB conversion."
-                )
-        
-        carrier_img = carrier_img.convert("RGB") 
-        
-        keys_for_overlay = list(data_dict.keys()) if show_keys_on_image_flag else None
-        overlay_title = "KeyLock: Data Embedded"
-        
-        final_carrier_with_overlay = core.draw_key_list_dropdown_overlay(
-            carrier_img, 
-            keys=keys_for_overlay,
-            title=overlay_title
-        )
-        
-        serial_data = json.dumps(data_dict).encode('utf-8')
-        encrypted_data = core.encrypt_data(serial_data, password)
-        
-        stego_final_img = core.embed_data_in_image(final_carrier_with_overlay, encrypted_data)
-        stego_final_img = core.set_pil_image_format_to_png(stego_final_img) 
-        
-        fname_base = "".join(c if c.isalnum() or c in ('_','-') else '_' for c in output_filename_base.strip()) or "keylock_img"
-        temp_fp = None
-        with tempfile.NamedTemporaryFile(prefix=fname_base+"_", suffix=".png", delete=False) as tmp:
-            stego_final_img.save(tmp, format="PNG")
-            temp_fp = tmp.name
-        
-        output_html_img_str = pil_to_base64_html(stego_final_img, max_width_px=480)
-        status_msg = (f"Data embedded into '{os.path.basename(temp_fp)}'.\n"
-                      f"{original_format_note}"
-                      f"Image contains visual \"{overlay_title}\" overlay "
-                      f"{'(with key list)' if show_keys_on_image_flag and keys_for_overlay else ''} "
-                      f"and your LSB-encoded secret data.\n"
-                      f"Secrets: {len(serial_data)}B (raw), {len(encrypted_data)}B (encrypted).")
-        return output_html_img_str, status_msg, temp_fp 
-    except ValueError as e: return None, f"Error: {str(e)}", None
-    except Exception as e: app_logger.error(f"Embed Error: {e}", exc_info=True); return None, f"Unexpected Error: {str(e)}", None
-
-def gradio_extract_data(stego_image_pil: Image.Image, password: str):
-    if stego_image_pil is None: return "Error: No image provided.", "Error: No image."
-    if not password: return "Error: Password cannot be empty.", "Error: Password required."
-    try:
-        stego_image_rgb = stego_image_pil.convert("RGB")
-        if hasattr(stego_image_pil, 'format') and stego_image_pil.format and stego_image_pil.format.upper() != "PNG": 
-            app_logger.warning(f"Uploaded image for extraction is format '{stego_image_pil.format}', not PNG. LSB data may be compromised if not the original KeyLock file.")
-        
-        extracted_data = core.extract_data_from_image(stego_image_rgb) 
-        decrypted_bytes = core.decrypt_data(extracted_data, password)
-        try: 
-            data = json.loads(decrypted_bytes.decode('utf-8'))
-            txt, stat = json.dumps(data, indent=2), "Data extracted successfully (JSON)."
-        except (json.JSONDecodeError, UnicodeDecodeError):
-            try: 
-                txt = "Decrypted (UTF-8, not JSON):\n"+decrypted_bytes.decode('utf-8')
-                stat = "Warning: Decrypted as UTF-8 (not JSON)."
-            except UnicodeDecodeError: 
-                txt = "Decrypted (raw hex, not JSON/UTF-8):\n"+decrypted_bytes.hex()
-                stat = "Warning: Decrypted as raw hex."
-        return txt, stat
-    except ValueError as e: return f"Error: {str(e)}", f"Extraction Failed: {str(e)}"
-    except Exception as e: app_logger.error(f"Extract Error: {e}", exc_info=True); return f"Unexpected Error: {str(e)}", f"Error: {str(e)}"
-
-def build_interface():
-    custom_theme = gr.themes.Base(
-        primary_hue="teal",
-        secondary_hue="purple",
-        neutral_hue="zinc",
-        text_size="sm",
-        spacing_size="md",
-        radius_size="sm",
-        font=["System UI", "sans-serif"]
-    )
-    custom_css = """
-    body {
-      background: linear-gradient(to bottom right, #2c3e50, #34495e); 
-      color: #ecf0f1; 
-    }
-    .gradio-container {
-        background: transparent !important;
-    }
-    .gr-box, .gr-panel, .gr-pill {
-        background-color: rgba(44, 62, 80, 0.8) !important; 
-        border-color: rgba(189, 195, 199, 0.2) !important; 
-    }
-    .gr-textbox, .gr-dropdown, .gr-button, .gr-code, .gr-chat-message, .gr-image {
-        border-color: rgba(189, 195, 199, 0.3) !important;
-        background-color: rgba(52, 73, 94, 0.9) !important; 
-        color: #ecf0f1 !important; 
-    }
-    .gr-button.gr-button-primary {
-        background-color: #1abc9c !important; 
-        color: white !important;
-        border-color: #16a085 !important;
-    }
-    .gr-button.gr-button-secondary {
-         background-color: #9b59b6 !important; 
-         color: white !important;
-         border-color: #8e44ad !important;
-    }
-    .gr-button.gr-button-stop {
-        background-color: #e74c3c !important; 
-        color: white !important;
-        border-color: #c0392b !important;
-    }
-    .gr-markdown {
-        background-color: rgba(44, 62, 80, 0.7) !important; 
-        padding: 10px; 
-        border-radius: 5px; 
-    }
-    .gr-markdown h1, .gr-markdown h2, .gr-markdown h3, .gr-markdown h4, .gr-markdown h5, .gr-markdown h6 {
-        color: #ecf0f1 !important; 
-        border-bottom-color: rgba(189, 195, 199, 0.3) !important; 
-    }
-    .gr-markdown pre code {
-        background-color: rgba(52, 73, 94, 0.95) !important; 
-        border-color: rgba(189, 195, 199, 0.3) !important;
-    }
-    .gr-image div img { /* Style for image preview */
-        border: 1px solid #ccc;
-        background-color: rgba(52, 73, 94, 0.9) !important;
-    }
-    .gr-file div button { /* Style for file download button */
-        background-color: #1abc9c !important;
-        color: white !important;
-        border: 1px solid #16a085 !important;
-    }
+# Remove direct requests import, will use model_logic
+# import requests
+import os
+import tempfile
+
+# --- build_logic.py is now a hard requirement ---
+from build_logic import (
+    create_space as build_logic_create_space,
+    _get_api_token as build_logic_get_api_token, # Keep this for HF Hub token logic
+    whoami as build_logic_whoami, # Keep this for HF user info
+    list_space_files_for_browsing,
+    get_space_repository_info,
+    get_space_file_content,
+    update_space_file,
+    parse_markdown as build_logic_parse_markdown,
+    delete_space_file as build_logic_delete_space_file,
+    get_space_runtime_status
+)
+print("build_logic.py loaded successfully.")
+
+# --- model_logic.py is now a hard requirement ---
+from model_logic import (
+    get_available_providers,
+    get_models_for_provider,
+    get_default_model_for_provider,
+    get_model_id_from_display_name, # Might not be strictly needed in app.py, but good practice
+    generate_stream # This is the core function we'll use
+)
+print("model_logic.py loaded successfully.")
+# --- End imports ---
+
+
+bbb = chr(96) * 3
+# Declare the global variable at the module level where it's initialized
+parsed_code_blocks_state_cache = []
+BOT_ROLE_NAME = "assistant"
+# Removed GROQ_API_ENDPOINT as it's now in model_logic
+
+
+DEFAULT_SYSTEM_PROMPT = f"""You are an expert AI programmer. Your role is to generate code and file structures based on user requests, or to modify existing code provided by the user.
+When you provide NEW code for a file, or MODIFIED code for an existing file, use the following format exactly:
+### File: path/to/filename.ext
+(You can add a short, optional, parenthesized description after the filename on the SAME line)
+{bbb}language
+# Your full code here
+{bbb}
+If the file is binary, or you cannot show its content, use this format:
+### File: path/to/binaryfile.ext
+[Binary file - approximate_size bytes]
+When you provide a project file structure, use this format:
+## File Structure
+{bbb}
+📁 Root
+  📄 file1.py
+  📁 subfolder
+    📄 file2.js
+{bbb}
+The role name for your responses in the chat history must be '{BOT_ROLE_NAME}'.
+Adhere strictly to these formatting instructions.
+If you update a file, provide the FULL file content again under the same filename.
+Only the latest version of each file mentioned throughout the chat will be used for the final output.
+Filenames in the '### File:' line should be clean paths (e.g., 'src/app.py', 'Dockerfile') and should NOT include Markdown backticks around the filename itself.
+If the user provides existing code (e.g., by pasting a Markdown structure), and asks for modifications, ensure your response includes the complete, modified versions of ONLY the files that changed, using the ### File: format. Unchanged files do not need to be repeated by you. The system will merge your changes with the prior state.
+If the user asks to delete a file, simply omit it from your next full ### File: list.
+If no code is provided, assist the user with their tasks.
+"""
+
+# --- Core Utility, Parsing, Export functions (mostly unchanged) ---
+def escape_html_for_markdown(text):
+    if not isinstance(text, str): return ""
+    # Minimal escaping, expand if needed
+    return text.replace("&", "&").replace("<", "<").replace(">", ">")
+
+
+def _infer_lang_from_filename(filename):
+    if not filename: return "plaintext"
+    if '.' in filename:
+        ext = filename.split('.')[-1].lower()
+        mapping = {
+            'py': 'python', 'js': 'javascript', 'ts': 'typescript', 'jsx': 'javascript', 'tsx': 'typescript',
+            'html': 'html', 'htm': 'html', 'css': 'css', 'scss': 'scss', 'sass': 'sass', 'less': 'less',
+            'json': 'json', 'xml': 'xml', 'yaml': 'yaml', 'yml': 'yaml', 'toml': 'toml',
+            'md': 'markdown', 'rst': 'rst',
+            'sh': 'bash', 'bash': 'bash', 'zsh': 'bash', 'bat': 'batch', 'cmd': 'batch', 'ps1': 'powershell',
+            'c': 'c', 'h': 'c', 'cpp': 'cpp', 'hpp': 'cpp', 'cs': 'csharp', 'java': 'java',
+            'rb': 'ruby', 'php': 'php', 'go': 'go', 'rs': 'rust', 'swift': 'swift', 'kt': 'kotlin', 'kts': 'kotlin',
+            'sql': 'sql', 'dockerfile': 'docker', 'tf': 'terraform', 'hcl': 'terraform',
+            'txt': 'plaintext', 'log': 'plaintext', 'ini': 'ini', 'conf': 'plaintext', 'cfg': 'plaintext',
+            'csv': 'plaintext', 'tsv': 'plaintext', 'err': 'plaintext',
+            '.env': 'plaintext', '.gitignore': 'plaintext', '.npmrc': 'plaintext', '.gitattributes': 'plaintext',
+            'makefile': 'makefile',
+        }
+        return mapping.get(ext, "plaintext")
+    base_filename = os.path.basename(filename)
+    if base_filename == 'Dockerfile': return 'docker'
+    if base_filename == 'Makefile': return 'makefile'
+    if base_filename.startswith('.'): return 'plaintext'
+    return "plaintext"
+
+def _clean_filename(filename_line_content):
+    text = filename_line_content.strip()
+    text = re.sub(r'[`\*_]+', '', text) # Remove markdown emphasis characters
+    # Try to match a valid-looking path first (allow spaces in folder/file names if quoted or part of a segment)
+    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)"
+        # 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()
+        # 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
+
+    # 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.strip()
+
+
+def _parse_chat_stream_logic(latest_bot_message_content, existing_files_state=None):
     """
-    with gr.Blocks(theme=custom_theme, css=custom_css, title=f"KeyLock Steganography v{__version__}") as keylock_app_interface:
-        gr.Markdown(f"<div align='center' style='margin-bottom:15px;'><span style='font-size:2.5em;font-weight:bold;'>🔑 KeyLock v{__version__}</span><h2 style='font-size:1.2em;color:#bdc3c7;margin-top:5px;'>Portable API Key Wallet in a PNG</h2></div>")
-        gr.HTML("<div align='center' style='margin-bottom:10px;font-size:0.9em;color:#bdc3c7;'>Securely embed and extract API key-value pairs (or any text) within PNG images using LSB steganography and AES-256-GCM encryption.</div>")
-        gr.HTML("<div align='center' style='margin-bottom:15px;font-size:0.9em;'><span style='font-weight:bold;'>GitHub: <a href='https://github.com/broadfield-dev/KeyLock-API-Wallet' target='_blank' style='color:#1abc9c;'>KeyLock-API-Wallet</a> | Decoder Module: <a href='https://github.com/broadfield-dev/keylock-decode' target='_blank' style='color:#1abc9c;'>keylock-decode</a></span></div>")
-        gr.HTML("<hr style='border-color: rgba(189, 195, 199, 0.2); margin-bottom:25px;'>")
-
-        with gr.Tabs():
-            with gr.TabItem(f"{ICON_EMBED} Embed Data"):
-                with gr.Row():
-                    with gr.Column(scale=2):
-                        embed_kv_input = gr.Textbox(
-                            label="Secret Data (Key:Value Pairs, one per line)",
-                            placeholder="API_KEY_1: your_secret_value_1\nSERVICE_USER = '[email protected]'\n# Lines starting with # are ignored",
-                            lines=7,
-                            info="Enter secrets as Key:Value or Key=Value. Each pair on a new line."
-                        )
-                        embed_password_input = gr.Textbox(
-                            label="Encryption Password",
-                            type="password",
-                            placeholder="Enter a strong password",
-                            info="Required to encrypt data. Keep this safe!"
-                        )
-                        embed_output_filename_base = gr.Textbox(
-                            label="Base Name for Downloaded Stego Image",
-                            value="keylock_wallet",
-                            info="'.png' will be appended. e.g., 'my_project_secrets'"
-                        )
-                        with gr.Accordion("Carrier Image Options", open=False):
-                            embed_generate_carrier_checkbox = gr.Checkbox(
-                                label="Generate new KeyLock Wallet image",
-                                value=True,
-                                info="Uncheck to upload your own PNG carrier image."
-                            )
-                            embed_input_image_upload = gr.Image(
-                                label="Upload Your Own PNG Carrier (Optional)",
-                                type="pil",
-                                image_mode="RGB",
-                                sources=["upload"],
-                                visible=False, # Initially hidden
-                                show_download_button=False,
-                                interactive=True
-                            )
-                        embed_show_keys_checkbox = gr.Checkbox(
-                            label="Show list of key names on image overlay",
-                            value=True,
-                            info="Displays embedded key names (not values) on the image."
-                        )
-                        embed_button = gr.Button("Embed Secrets "+ICON_EMBED, variant="primary")
-
-                    with gr.Column(scale=3):
-                        gr.Markdown("### Output Image & Status")
-                        embed_output_status = gr.Textbox(
-                            label="Embedding Status",
-                            lines=4,
-                            interactive=False,
-                            placeholder="Status messages will appear here..."
-                        )
-                        embed_output_image_html = gr.HTML(
-                            label="Preview of Stego Image (Max 480px width)",
-                            value="<div style='text-align:center; color:#bdc3c7; padding:20px;'>Image preview will appear here after embedding.</div>"
-                        )
-                        embed_download_file = gr.File(
-                            label="Download Your KeyLock Image (PNG)",
-                            interactive=False,
-                            file_count="single"
-                        )
-                
-                def toggle_carrier_upload(generate_flag):
-                    return gr.update(visible=not generate_flag)
-
-                embed_generate_carrier_checkbox.change(
-                    fn=toggle_carrier_upload,
-                    inputs=[embed_generate_carrier_checkbox],
-                    outputs=[embed_input_image_upload]
-                )
-                embed_button.click(
-                    fn=gradio_embed_data,
-                    inputs=[
-                        embed_kv_input,
-                        embed_password_input,
-                        embed_input_image_upload,
-                        embed_generate_carrier_checkbox,
-                        embed_show_keys_checkbox,
-                        embed_output_filename_base
-                    ],
-                    outputs=[
-                        embed_output_image_html,
-                        embed_output_status,
-                        embed_download_file
-                    ]
-                )
-
-            with gr.TabItem(f"{ICON_EXTRACT} Extract Data"):
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        extract_stego_image_upload = gr.Image(
-                            label="Upload KeyLock PNG Image",
-                            type="pil",
-                            image_mode="RGB",
-                            sources=["upload"],
-                            show_download_button=False,
-                            interactive=True,
-                            info="Upload the PNG image containing KeyLock data."
-                        )
-                        extract_password_input = gr.Textbox(
-                            label="Decryption Password",
-                            type="password",
-                            placeholder="Enter the password used during embedding",
-                            info="Required to decrypt and extract data."
-                        )
-                        extract_button = gr.Button("Extract Secrets "+ICON_EXTRACT, variant="primary")
-                    
-                    with gr.Column(scale=2):
-                        gr.Markdown("### Extracted Data & Status")
-                        extract_output_status = gr.Textbox(
-                            label="Extraction Status",
-                            lines=2,
-                            interactive=False,
-                            placeholder="Status messages will appear here..."
-                        )
-                        extract_output_data = gr.Textbox(
-                            label="Extracted Secret Data",
-                            lines=10,
-                            interactive=False,
-                            placeholder="Extracted data (usually JSON) will appear here...",
-                            show_copy_button=True
-                        )
-                
-                extract_button.click(
-                    fn=gradio_extract_data,
-                    inputs=[
-                        extract_stego_image_upload,
-                        extract_password_input
-                    ],
-                    outputs=[
-                        extract_output_data,
-                        extract_output_status
-                    ]
-                )
-        
-        gr.Markdown("<hr style='border-color: rgba(189, 195, 199, 0.1); margin-top: 30px; margin-bottom:10px;'>")
-        gr.Markdown(f"<div style='text-align:center; font-size:0.8em; color:#7f8c8d;'>KeyLock-API-Wallet v{__version__}. Use responsibly.</div>")
-
-    return keylock_app_interface
-
-def main():
-    app_logger.info(f"Starting KeyLock Gradio Application v{__version__}...")
+    Parses a single bot message content string to find file blocks and updates the state.
+    Assumes existing_files_state is the current state *before* this message.
+    """
+    # This function takes state as an argument and returns new state, it doesn't need 'global'
+    latest_blocks_dict = {}
+    if existing_files_state:
+        # Copy existing blocks, except for potential structure blocks that might be overwritten
+        for block in existing_files_state:
+             if not block.get("is_structure_block"):
+                 latest_blocks_dict[block["filename"]] = block.copy()
+             # Keep existing structure block for now, it might be replaced below
+
+
+    results = {"parsed_code_blocks": [], "preview_md": "", "default_selected_filenames": [], "error_message": None}
+    content = latest_bot_message_content or ""
+
+    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```")
+
+    # Process the latest bot message for updates to file blocks
+    structure_match = structure_pattern.search(content)
+    if structure_match:
+        # Add/Overwrite 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}
+    else:
+         # If the latest message *doesn't* have a structure block, keep the previous one if it existed
+         existing_structure_block = next((b for b in (existing_files_state or []) if b.get("is_structure_block")), None)
+         if existing_structure_block:
+              latest_blocks_dict["File Structure (original)"] = existing_structure_block.copy()
+
+
+    # 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 # Should not happen with the regex
+        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
+    latest_blocks_dict.update(current_message_file_blocks)
+
+
+    # Convert dictionary values back to a list
+    current_parsed_blocks = list(latest_blocks_dict.values())
+    # Sort: structure block first, then files alphabetically
+    current_parsed_blocks.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
+
+    # Update the global cache outside this function if needed, or pass it back
+    # For now, let's return the new state and let the caller update the cache.
+    results["parsed_code_blocks"] = current_parsed_blocks
+    results["default_selected_filenames"] = [b["filename"] for b in current_parsed_blocks if not b.get("is_structure_block")]
+    return results
+
+def _export_selected_logic(selected_filenames, space_line_name_for_md, parsed_blocks_for_export):
+    # This function remains largely the same, using the provided parsed_blocks_for_export
+    # It takes state as an argument and doesn't need 'global'
+    results = {"output_str": "", "error_message": None, "download_filepath": None}
+    # Filter out structure blocks for file listing/export content
+    exportable_blocks_content = [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_content = [b for b in parsed_blocks_for_export if b.get("is_binary") or b.get("code", "").startswith("[Binary or Skipped file]")]
+
+    # Collect all filenames (including binary/error ones) for the structure list
+    all_filenames_in_state = sorted(list(set(b["filename"] for b in parsed_blocks_for_export if not b.get("is_structure_block"))))
+
+    if not all_filenames_in_state 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:
+                tmpfile.write(results["output_str"]); results["download_filepath"] = tmpfile.name
+        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}"]
+
+    # 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"])
+        if all_filenames_in_state:
+            for fname in all_filenames_in_state: 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
+    # Exportable content blocks only
+    files_to_export_content = []
+    if selected_filenames:
+        files_to_export_content = [b for b in exportable_blocks_content if b["filename"] in selected_filenames]
+    else:
+        files_to_export_content = exportable_blocks_content # Export all content blocks by default
+
+    # Add binary/error blocks if they were selected or if exporting all (and they exist)
+    # Binary/error blocks are listed in the structure, but their *content* is just the marker string
+    binary_error_blocks_to_export = []
+    if selected_filenames:
+         binary_error_blocks_to_export = [b for b in binary_blocks_content if b["filename"] in selected_filenames]
+    elif binary_blocks_content:
+         binary_error_blocks_to_export = binary_blocks_content # Include all binary/error if exporting all
+
+    # Combine and sort all blocks whose content/marker should be included
+    all_blocks_to_export_content = sorted(files_to_export_content + binary_error_blocks_to_export, key=lambda b: b["filename"])
+
+
+    for block in all_blocks_to_export_content:
+        output_lines.append(f"### File: {block['filename']}")
+        if block.get('is_binary') or block.get("code", "").startswith("[Binary file") or block.get("code", "").startswith("[Error loading content:") or block.get("code", "").startswith("[Binary or Skipped file]"):
+            # For binary/error placeholders, just add the marker line
+            output_lines.append(block.get('code','[Binary or Skipped file]'))
+        else:
+            # For actual code/text content
+            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 all_filenames_in_state: output_lines.append("*No files in state.*")
+    elif not exported_content: output_lines.append("*No files with editable content are in the state or selected.*") # Message updated
+
+    final_output_str = "\n".join(output_lines)
+    results["output_str"] = final_output_str
     try:
-        ImageFont.truetype("DejaVuSans.ttf", 10) 
-        app_logger.info("DejaVuSans font found, PIL font rendering should be good.")
-    except IOError:
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
+            tmpfile.write(final_output_str); results["download_filepath"] = tmpfile.name
+    except Exception as e: print(f"Error creating temp file: {e}"); results["error_message"] = "Could not prepare file for download."
+    return results
+
+
+def _convert_gr_history_to_api_messages(system_prompt, gr_history, current_user_message=None):
+    # This function is fine as is, it produces standard OpenAI format
+    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})
+        # 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 _generate_ui_outputs_from_cache(owner, space_name):
+    # Declare global at the top
+    global parsed_code_blocks_state_cache
+    # This function remains largely the same, generating UI previews and the export MD
+    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.*" 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:
+        preview_md_lines = ["## Detected/Updated Files & Content (Latest Versions):"]
+        for block in parsed_code_blocks_state_cache:
+            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') and block.get('language') != 'binary': preview_md_lines.append(f" (Language: `{block['language']}`)\n")
+            else: preview_md_lines.append("\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")
+
+        # _export_selected_logic handles selecting which files to include in the export MD
+        # Passing None means export all non-structure/non-binary/non-error content + list all files in structure
+        export_result = _export_selected_logic(None, space_line_name, parsed_code_blocks_state_cache)
+        formatted_md_val = export_result["output_str"]
+        download_file = export_result["download_filepath"]
+
+    return formatted_md_val, preview_md_val, gr.update(value=download_file, interactive=download_file is not None)
+
+
+# --- Refactored Chat Submit Handler ---
+def handle_chat_submit(user_message, chat_history, api_key_input, provider_select, model_select, system_prompt, hf_owner_name, hf_repo_name, _current_formatted_markdown):
+    # Declare global at the top
+    global parsed_code_blocks_state_cache
+    _chat_msg_in = ""
+    _chat_hist = list(chat_history)
+    _status = "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:
-            ImageFont.truetype("arial.ttf", 10) 
-            app_logger.info("Arial font found, PIL font rendering should be good.")
-        except IOError:
-            app_logger.warning("Common system fonts (DejaVuSans/Arial) not found. PIL might use basic bitmap font if other preferred fonts in core.py are also unavailable. Overlay text quality might be affected.")
-    
-    keylock_app_interface = build_interface()
-    
-    launch_args = {"allowed_paths": [tempfile.gettempdir()]}
-    
-    server_name = os.environ.get('GRADIO_SERVER_NAME')
-    server_port = os.environ.get('GRADIO_SERVER_PORT')
-
-    if server_name:
-        launch_args["server_name"] = server_name
-        app_logger.info(f"Using server_name from environment: {server_name}")
-    if server_port:
+            parsed_from_md = build_logic_parse_markdown(_current_formatted_markdown)
+            # Update cache with files parsed from the markdown.
+            # Structure block from AI is volatile, always prefer structure from latest AI message.
+            # Files from markdown overwrite any previous file blocks.
+            new_cache_state = []
+            # Add structure block from *current cache* if it exists, it will be replaced if the AI provides a new one
+            existing_structure_block = next((b for b in parsed_code_blocks_state_cache if b.get("is_structure_block")), None)
+            if existing_structure_block:
+                 new_cache_state.append(existing_structure_block.copy()) # Add copy
+
+
+            for f_info in parsed_from_md.get("files", []):
+                 # Only add if it has a path and isn't the structure block representation placeholder
+                 if f_info.get("path") and f_info["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]"))
+                      # Check if a block with this filename already exists in new_cache_state and replace it
+                      found_existing = False
+                      for i, block in enumerate(new_cache_state):
+                           if block["filename"] == f_info["path"] and not block.get("is_structure_block"): # Only replace non-structure blocks
+                                new_cache_state[i] = {
+                                     "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
+                                }
+                                found_existing = True
+                                break
+                      if not found_existing:
+                            new_cache_state.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
+                            })
+
+            # Sort the updated cache state
+            new_cache_state.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
+            parsed_code_blocks_state_cache = new_cache_state # Update global 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 (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update); return
+    _chat_hist.append((user_message, None)); _status = f"Sending to {model_select} via {provider_select}..."
+    yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
+
+    # Pass the API key from the UI directly to model_logic
+    api_key_override = api_key_input
+    # model_id = get_model_id_from_display_name(provider_select, model_select) # model_logic handles display name to ID
+
+
+    current_sys_prompt = system_prompt.strip() or DEFAULT_SYSTEM_PROMPT
+
+    # Include current file contents in the prompt as context for the AI
+    # This context is built from the *current cache state* (which was just updated from the formatted markdown)
+    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") or block.get("code", "").startswith("["): # Include binary/error markers
+                    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
+    # This combined message structure helps the model understand the context and the expected output format
+    user_message_with_context = user_message.strip()
+    if current_files_context.strip():
+         user_message_with_context = user_message_with_context + current_files_context + "\nBased 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. If you are providing a file structure, use the `## File Structure\n```\n...\n```\n` format. Omit files you want to delete from your response."
+
+
+    # Convert history to API messages, including the user message with context
+    api_msgs = _convert_gr_history_to_api_messages(current_sys_prompt, _chat_hist[:-1], user_message_with_context)
+
+    # --- Call the new model_logic streaming function ---
+    try:
+        _status = f"Waiting for {model_select} via {provider_select}..."; yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
+
+        # Accumulate the full response content for parsing *after* streaming
+        full_bot_response_content = ""
+        error_during_stream = None
+
+        # Generate stream from model_logic
+        for chunk in generate_stream(provider_select, model_select, api_key_override, api_msgs):
+             if chunk is None: continue # Skip None chunks if any
+             if isinstance(chunk, str) and (chunk.startswith("Error: ") or chunk.startswith("API HTTP Error")):
+                  # If an error chunk is received, treat it as the final output and stop
+                  full_bot_response_content = chunk
+                  error_during_stream = chunk
+                  break # Stop processing stream
+             else:
+                  # Accumulate response and update the last message in chat_hist
+                  full_bot_response_content += str(chunk) # Ensure chunk is string
+                  _chat_hist[-1] = (user_message, full_bot_response_content)
+                  _status = f"Streaming from {model_select}..."
+                  # Yield update immediately after receiving chunk
+                  yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
+
+        # After the stream finishes or breaks
+        if error_during_stream:
+             _status = error_during_stream # Set status to the error message
+        elif full_bot_response_content and not full_bot_response_content.startswith("Error: "): # Only parse if it's not an error message
+            _status = f"Streaming complete. Processing files from {model_select} response..."
+
+            # Pass the *current state* (updated from markdown at the start)
+            # and the *latest bot message content* to the parsing logic.
+            # _parse_chat_stream_logic will merge and update based on the bot's response.
+            parsing_res = _parse_chat_stream_logic(full_bot_response_content, 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 for visibility? Or just status?
+                # For now, update status and detected files area with error message
+                _detected_files_update = gr.Markdown(f"## Parsing Error\n`{escape_html_for_markdown(parsing_res['error_message'])}`")
+            else:
+                # Update the global cache with the new state returned by the parser
+                global parsed_code_blocks_state_cache
+                parsed_code_blocks_state_cache = parsing_res["parsed_code_blocks"]
+
+                # Regenerate UI outputs from the *updated* cache
+                _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."
+        else:
+            # Handle cases where the stream finished but yielded no content (e.g., filter) or only an error was yielded
+            if not error_during_stream:
+                 _status = "AI response complete, but returned no content."
+            # Keep existing previews/markdown if no content was generated to parse or if it was an error message
+            _formatted_output_update, _detected_files_update, _download_btn_update = _generate_ui_outputs_from_cache(hf_owner_name, hf_repo_name)
+
+
+    except Exception as e:
+        # Catch any errors that occurred *before* or *during* the stream setup/iteration
+        error_msg = f"An unexpected error occurred during AI generation: {e}"
+        print(f"Unexpected error in chat submit stream: {e}")
+        # Update the last chat message with the error
+        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
+        # Previews and markdown might not be affected by a generation error, keep existing state
+        _formatted_output_update, _detected_files_update, _download_btn_update = _generate_ui_outputs_from_cache(hf_owner_name, hf_repo_name)
+
+
+    # Final yield to update UI after all processing
+    yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
+
+
+# --- Handler to update model dropdown based on provider selection ---
+def update_models_dropdown(provider_select):
+    """Updates the model dropdown choices and selects the default model."""
+    if not provider_select:
+        return gr.update(choices=[], value=None)
+    models = get_models_for_provider(provider_select)
+    default_model = get_default_model_for_provider(provider_select)
+    # Ensure default is in choices, or pick first, or None
+    if default_model and default_model in models:
+        selected_value = default_model
+    elif models:
+        selected_value = models[0]
+    else:
+        selected_value = None
+
+    return gr.update(choices=models, value=selected_value)
+
+
+# --- Existing handlers for Load, Build, Edit, Delete, Status (Mostly unchanged, just global placement) ---
+
+def handle_load_existing_space(hf_api_key_ui, ui_owner_name, ui_space_name):
+    # Declare global at the top
+    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.*"
+    _chat_history_clear = [] # Clear chat history on loading a new space
+
+    # 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, _chat_history_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 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 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:
+        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, _chat_history_clear)
+         parsed_code_blocks_state_cache = [] # Clear cache on error
+         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)
+
+    # 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}"
+        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, _chat_history_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 '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, _chat_history_clear)
+        return # Stop execution
+
+
+    loaded_files_for_cache = [] # Build a list to become the new cache state
+    _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, _chat_history_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",".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 or \
+           file_path == "README.md" or file_path == "LICENSE": # Optionally skip common non-code files like README/LICENSE
+            loaded_files_for_cache.append({"filename": file_path, "code": "[Binary or Skipped file]", "language": "binary", "is_binary": True, "is_structure_block": False}); continue
+
+        # Handle potential issues with reading large files or non-utf8 files
         try:
-            launch_args["server_port"] = int(server_port)
-            app_logger.info(f"Using server_port from environment: {server_port}")
-        except ValueError:
-            app_logger.warning(f"Invalid GRADIO_SERVER_PORT: {server_port}. Using default.")
-            
-    keylock_app_interface.launch(**launch_args)
+            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_cache.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_cache.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_cache.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
+
+    # Add a placeholder structure block if none was loaded (AI will generate one later if needed)
+    # This ensures the cache isn't empty except for files
+    # structure_block = next((b for b in loaded_files_for_cache if b.get("is_structure_block")), None)
+    # if not structure_block:
+    #      loaded_files_for_cache.insert(0, {"filename": "File Structure (original)", "code": "📁 Root\n  ...\n", "language": "plaintext", "is_binary": False, "is_structure_block": True})
+
+    parsed_code_blocks_state_cache = loaded_files_for_cache
+    _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. {len(file_list)} files listed."
+    _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), value=None) # Use the full file list for the dropdown
+
+    # Final yield with updated state
+    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, _chat_history_clear)
+
+
+def handle_build_space_button(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, space_sdk_ui, formatted_markdown_content):
+    # Declare global at the top
+    global parsed_code_blocks_state_cache
+    # ... (this function calls build_logic_create_space and refreshes file list)
+    _build_status, _iframe_html, _file_browser_update = "Starting space build process...", gr.update(value=None, visible=False), gr.update(visible=False, choices=[], value=None)
+    # Include outputs for owner/space name textboxes in the initial yield
+    yield _build_status, _iframe_html, _file_browser_update, gr.update(value=ui_owner_name_part), gr.update(value=ui_space_name_part) # 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, gr.update(), gr.update(); return
+    final_owner_for_build = ui_owner_name_part
+    if not final_owner_for_build:
+        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, gr.update(), gr.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
+                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, gr.update(), gr.update(); return
+
+    # Before building, parse the markdown to ensure the cache reflects exactly what's being built
+    # This prevents inconsistencies if the user manually edited the markdown output
+    try:
+        parsed_from_md_for_build = build_logic_parse_markdown(formatted_markdown_content)
+        # Replace the global cache state with the state derived from the markdown being built
+        parsed_code_blocks_state_cache = []
+        if parsed_from_md_for_build.get("owner_md"): # Update UI owner/space name if present in MD
+             ui_owner_name_part = parsed_from_md_for_build["owner_md"] # Use this updated value later
+        if parsed_from_md_for_build.get("repo_name_md"):
+             ui_space_name_part = parsed_from_md_for_build["repo_name_md"] # Use this updated value later
+
+        # Rebuild cache from parsed markdown files + structure block
+        structure_block_md = next((f for f in parsed_from_md_for_build.get("files", []) if f.get("path") == "File Structure (original)"), None)
+        if structure_block_md:
+             parsed_code_blocks_state_cache.append({
+                 "filename": structure_block_md["path"],
+                 "code": structure_block_md["content"],
+                 "language": "plaintext", # Markdown parser doesn't detect lang for structure block ```
+                 "is_binary": False,
+                 "is_structure_block": True
+             })
+
+        for f_info in parsed_from_md_for_build.get("files", []):
+             if f_info.get("path") and f_info["path"] != "File Structure (original)"):
+                  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"]))
+
+    except Exception as e:
+         _build_status = f"Build Error: Failed to parse Markdown structure before building: {e}";
+         # Yield error status, including keeping current owner/space name in textboxes
+         yield _build_status, _iframe_html, _file_browser_update, gr.update(value=ui_owner_name_part), gr.update(value=ui_space_name_part); return # Stop build on parse error
+
+
+    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}"
+
+    # Update UI with owner/space names extracted from markdown if present
+    owner_name_output = gr.update(value=ui_owner_name_part)
+    space_name_output = gr.update(value=ui_space_name_part)
+
+    if "Successfully" in result_message:
+        # Use potentially updated owner/space name from markdown parsing
+        sub_owner = re.sub(r'[^a-z0-9\-]+', '-', ui_owner_name_part.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: [Link]({iframe_url}) (Repo: https://huggingface.co/spaces/{ui_owner_name_part}/{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, ui_owner_name_part)
+        if err_list: _build_status += f"\nFile list refresh error after build: {err_list}"; _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), 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")
+
+    # Final yield including potential updates to owner/space name textboxes
+    yield _build_status, _iframe_html, _file_browser_update, owner_name_output, space_name_output
+
+
+# File editing handlers are okay, just need to ensure they update the cache properly after commit/delete
+def handle_load_file_for_editing(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, selected_file_path):
+    # Declare global at the top (even though it's not modified here, it's good practice if the function *might* interact with it in the future, or just for consistency)
+    # In this specific function, `global` isn't strictly needed as it only *reads* indirectly via _generate_ui_outputs_from_cache which handles its own global
+    # Keeping it here for consistency as per the error symptoms observed in similar functions.
+    global parsed_code_blocks_state_cache # Added global here
+
+    _file_content_val, _edit_status_val, _commit_msg_val, _lang_update = "", "Error: No file selected.", gr.update(value=""), gr.update(language="python") # 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, 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
+                 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="python") # 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 = 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="python") # Default language for error display
+        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):
+    # Declare global at the top
+    global parsed_code_blocks_state_cache
+    _edit_status_val = "Processing commit..."
+    # Initialize updates for components that might change
+    _file_browser_update_val = gr.update() # Will update choices or value
+    _formatted_md_out = gr.update() # Will update markdown
+    _detected_preview_out = gr.update() # Will update markdown preview
+    _download_btn_out = gr.update() # Will update download button
+
+    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, 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
+                 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"] = 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/edited via editor and wasn't in initial load cache (e.g. binary/error placeholder), add/replace it
+             # First remove any existing placeholder for this file
+             parsed_code_blocks_state_cache = [b for b in parsed_code_blocks_state_cache if b["filename"] != file_to_edit_path]
+             # Then add the new text content block
+             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)
+
+        # 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):
+    # Declare global at the top
+    global parsed_code_blocks_state_cache
+    _edit_status_val = "Processing deletion..."
+    # Initialize updates for components that might change/clear
+    _file_browser_choices_update = gr.update() # Update choices
+    _file_browser_value_update = None # Clear selected file value
+    _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 = gr.update() # Update markdown
+    _detected_preview_out = gr.update() # Update markdown preview
+    _download_btn_out = gr.update() # Update download button
+
+
+    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, 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
+                 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
+
+    # Always refresh the file list dropdown choices after a delete attempt, successful or not
+    new_file_list, err_list = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, owner_to_use)
+
+    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)
+
+
+        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") # Set value to error state
+        else:
+             _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value=None) # Clear selection visually and internally
+
+        _file_browser_value_update = None # Explicitly set value to None to clear selection visual
+
+
+    else: # If deletion failed
+        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")
+            _file_browser_value_update = "Error refreshing files" # Keep error state in value if list failed
+        else:
+             # If list refresh succeeded but delete failed, refresh choices and keep the *failed-to-delete* file selected
+             _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 visually
+
+        # Markdown and preview are not changed if deletion failed, keep current updates as gr.update()
+        # Regenerate previews to show they are unchanged
+        _formatted_md_out, _detected_preview_out, _download_btn_out = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name_part)
+
+
+    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)
+
+
+# Space status handler is okay
+def handle_refresh_space_status(hf_api_key_ui, ui_owner_name, ui_space_name):
+    # This function doesn't modify the global cache, so no global declaration needed.
+    # ... (rest of this function is the same)
+    yield "*Fetching space status...*" # Initial feedback
+    owner_to_use = ui_owner_name
+    if not owner_to_use:
+        token, token_err = build_logic_get_api_token(hf_api_key_ui)
+        if token_err or not token: yield f"**Error:** {token_err or 'Cannot determine owner.'}"; 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
+        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`{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})")
+        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
+
+
+# Define a custom theme with a dark background and contrasting colors
+# And add custom CSS for a background gradient and component styling
+custom_theme = gr.themes.Base(
+    primary_hue="teal", # Teal for primary actions
+    secondary_hue="purple", # Purple for secondary elements
+    neutral_hue="zinc", # Zinc for neutral/backgrounds (dark gray)
+    text_size="sm", # Smaller text size for a denser, professional look
+    spacing_size="md", # Medium spacing
+    radius_size="sm", # Small border radius
+    font=["System UI", "sans-serif"] # Use system font
+)
+
+custom_css = """
+body {
+  background: linear-gradient(to bottom right, #2c3e50, #34495e); /* Dark blue-gray gradient */
+  color: #ecf0f1; /* Light text color for dark background */
+}
+/* Adjust main Gradio container background to be transparent to see body gradient */
+.gradio-container {
+    background: transparent !important;
+}
+/* Adjust component backgrounds for contrast against the body gradient */
+.gr-box, .gr-panel, .gr-pill {
+    background-color: rgba(44, 62, 80, 0.8) !important; /* Slightly lighter transparent dark blue-gray */
+    border-color: rgba(189, 195, 199, 0.2) !important; /* Light border for contrast */
+}
+/* Adjust inputs, dropdowns, buttons etc. for visibility */
+.gr-textbox, .gr-dropdown, .gr-button, .gr-code, .gr-chat-message {
+    border-color: rgba(189, 195, 199, 0.3) !important;
+    background-color: rgba(52, 73, 94, 0.9) !important; /* Slightly different dark blue-gray */
+    color: #ecf0f1 !important; /* Ensure text is light */
+}
+.gr-button.gr-button-primary {
+    background-color: #1abc9c !important; /* Teal from primary_hue */
+    color: white !important;
+    border-color: #16a085 !important;
+}
+.gr-button.gr-button-secondary {
+     background-color: #9b59b6 !important; /* Purple from secondary_hue */
+     color: white !important;
+     border-color: #8e44ad !important;
+}
+.gr-button.gr-button-stop {
+    background-color: #e74c3c !important; /* Red for stop/delete */
+    color: white !important;
+    border-color: #c0392b !important;
+}
+/* Adjust markdown backgrounds */
+.gr-markdown {
+    background-color: rgba(44, 62, 80, 0.7) !important; /* Transparent dark background */
+    padding: 10px; /* Add some padding */
+    border-radius: 5 al;
+}
+/* Style markdown headers for better contrast */
+.gr-markdown h1, .gr-markdown h2, .gr-markdown h3, .gr-markdown h4, .gr-markdown h5, .gr-markdown h6 {
+    color: #ecf0f1 !important; /* Ensure headers are light */
+    border-bottom-color: rgba(189, 195, 199, 0.3) !important; /* Light separator */
+}
+/* Style code blocks within markdown */
+.gr-markdown pre code {
+    background-color: rgba(52, 73, 94, 0.95) !important; /* Darker code background */
+    border-color: rgba(189, 195, 199, 0.3) !important;
+}
+/* Chatbot specific styling */
+.gr-chatbot {
+    background-color: rgba(44, 62, 80, 0.7) !important;
+    border-color: rgba(189, 195, 199, 0.2) !important;
+}
+.gr-chatbot .message {
+    background-color: rgba(52, 73, 94, 0.9) !important; /* Dark background for messages */
+    color: #ecf0f1 !important;
+    border-color: rgba(189, 195, 199, 0.3) !important;
+}
+.gr-chatbot .message.user {
+    background-color: rgba(46, 204, 113, 0.9) !important; /* Greenish background for user messages */
+    color: black !important; /* Dark text for green background */
+}
+"""
+
+
+# Get initial providers and models for UI setup
+available_providers = get_available_providers()
+default_provider = available_providers[0] if available_providers else None
+initial_models = get_models_for_provider(default_provider) if default_provider else []
+initial_default_model = get_default_model_for_provider(default_provider) if default_provider else None
+# Ensure initial_default_model is in the initial_models list, fallback if not
+if initial_default_model not in initial_models and initial_models:
+    initial_default_model = initial_models[0]
+
+
+with gr.Blocks(theme=custom_theme, css=custom_css) 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");
+            # Single API key input, model_logic decides which env var to check or uses this override
+            api_key_input = gr.Textbox(label="AI Provider API Key (Optional Override)", type="password", placeholder="Paste key here or set env var (e.g., GROQ_API_KEY)");
+            hf_api_key_input = gr.Textbox(label="Hugging Face Token (for building/loading)", type="password", placeholder="hf_...")
+            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");
+            provider_select = gr.Dropdown(label="AI Provider", choices=available_providers, value=default_provider, info="Select an AI model provider.");
+            model_select = gr.Dropdown(label="AI Model", choices=initial_models, value=initial_default_model, info="Select a model.");
+            system_prompt_input = gr.Textbox(label="System Prompt", lines=8, value=DEFAULT_SYSTEM_PROMPT, interactive=True)
+        with gr.Column(scale=3):
+            gr.Markdown("## 💬 AI Chat & Code Generation")
+            # Updated chatbot avatar
+            chatbot_display = gr.Chatbot(label="AI Chat", height=400, bubble_full_width=False, avatar_images=(None, "https://huggingface.co/datasets/huggingface/badges/resolve/main/huggingface-bot-avatar.svg"))
+            with gr.Row(): chat_message_input = gr.Textbox(show_label=False, placeholder="Your Message...", scale=7); send_chat_button = gr.Button("Send", variant="primary", scale=1, size="lg")
+            status_output = gr.Textbox(label="Chat/Process Status", interactive=False, lines=1, value="Ready.")
+            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"):
+                    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 outputs also include updating owner/space names in the textboxes
+                    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.")
+                    file_content_editor = gr.Code(label="File Content Editor", language="python", lines=15, interactive=True)
+                    commit_message_input = gr.Textbox(label="Commit Message", placeholder="e.g., Updated app.py", value="Update via AI Space Editor")
+                    with gr.Row(): update_file_button = gr.Button("Commit Changes", variant="primary", scale=2); delete_file_button = gr.Button("🗑️ Delete Selected File", variant="stop", scale=1)
+                    edit_status_display = gr.Textbox(label="File Edit/Delete Status", interactive=False, lines=2, value="*Select file...*")
+
+    # --- Event Handlers ---
+
+    # Provider dropdown change event to update model dropdown
+    provider_select.change(
+        fn=update_models_dropdown,
+        inputs=provider_select,
+        outputs=model_select
+    )
+
+    # Chat submit handler outputs
+    chat_outputs = [chat_message_input, chatbot_display, status_output, detected_files_preview, formatted_space_output_display, download_button]
+    # Chat submit handler inputs
+    chat_inputs = [chat_message_input, chatbot_display, api_key_input, provider_select, model_select, system_prompt_input, owner_name_input, space_name_input, formatted_space_output_display] # Pass current formatted markdown as context
+
+    # Wire chat buttons
+    send_chat_button.click(
+        fn=handle_chat_submit,
+        inputs=chat_inputs,
+        outputs=chat_outputs
+    )
+    chat_message_input.submit( # Allow submitting with Enter key
+        fn=handle_chat_submit,
+        inputs=chat_inputs,
+        outputs=chat_outputs
+    )
+
+    # Load space outputs include clearing chat history
+    load_space_outputs = [formatted_space_output_display, detected_files_preview, 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, chatbot_display] # Added chatbot_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 now include updating owner/space name textboxes
+    build_outputs = [build_status_display, space_iframe_display, file_browser_dropdown, owner_name_input, space_name_input] # Added owner_name_input, space_name_input
+    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)
+
+    # File edit load outputs include clearing/setting commit message and language
+    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)
+
+    # Commit file outputs include refreshing previews and file browser state
+    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 include refreshing previews, file browser state, and clearing editor
+    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] # Two file_browser_dropdown outputs: choices and value
+    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 handler is okay
+    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__":
-    main()
+    demo.launch(debug=True, share=False)