Update app.py

app.py

@@ -1,1127 +1,380 @@
-# app.py
+# keylock/app.py
 import gradio as gr
-import re
-import json
-# Remove direct requests import, will use model_logic
-# import requests
-import os
+from PIL import Image, ImageFont
 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):
-    """
-    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:
-        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:
-            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 ---
+import os
+import json
+import logging
+import traceback
+import base64
+import io
+
+from . import core # Use relative import for core module
+from . 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:
-        _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
-                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."
+        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:
-            # 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:
-            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
+            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:
-        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
+        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"]
     )
-    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])
-
+    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;
+    }
+    """
+    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__}...")
+    try:
+        ImageFont.truetype("DejaVuSans.ttf", 10) 
+        app_logger.info("DejaVuSans font found, PIL font rendering should be good.")
+    except IOError:
+        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:
+        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)
 
 if __name__ == "__main__":
-    demo.launch(debug=True, share=False)
+    main()