Update app_logic.py

app_logic.py

@@ -11,9 +11,8 @@ from huggingface_hub import (
 )
 import logging
 from pathlib import Path
-from PIL import Image # For type hinting from gr.Image(type="pil")
+from PIL import Image 
 
-# Attempt to import keylock_decode
 try:
     from keylock_decode import decode_from_image_pil
     KEYLOCK_DECODE_AVAILABLE = True
@@ -28,8 +27,7 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-
-# --- Helper Function to Get API Token (Unchanged from previous version) ---
+# --- Helper Function to Get API Token (Unchanged) ---
 def _get_api_token(ui_token_from_textbox=None):
     env_token = os.getenv('HF_TOKEN')
     if env_token:
@@ -41,185 +39,147 @@ def _get_api_token(ui_token_from_textbox=None):
     logger.warning("HF API token not found in environment or UI textbox.")
     return None, "Error: Hugging Face API token not provided. Please enter it in the textbox or load it from an image."
 
-# --- Updated Function for KeyLock-Decode Integration with Debugging ---
+# --- `load_token_from_image_and_set_env` (Unchanged from previous debug version) ---
 def load_token_from_image_and_set_env(image_pil_object: Image.Image, password: str):
     if not KEYLOCK_DECODE_AVAILABLE:
         return "Error: KeyLock-Decode library is not installed. This feature is disabled."
-
-    if image_pil_object is None:
-        return "Error: No image provided for decoding."
-    if not password:
-        return "Error: Password cannot be empty for image decoding."
-
+    if image_pil_object is None: return "Error: No image provided for decoding."
+    if not password: return "Error: Password cannot be empty for image decoding."
     status_messages_display = []
-    
-    # --- DEBUGGING: Check image properties and save it ---
     debug_image_path_str = "Not saved (error or no image)."
     try:
         if image_pil_object:
-            # Log original properties
             original_mode = getattr(image_pil_object, 'mode', 'N/A')
-            original_format = getattr(image_pil_object, 'format', 'N/A') # Format might be None for PIL object not read from file
+            original_format = getattr(image_pil_object, 'format', 'N/A')
             original_size = getattr(image_pil_object, 'size', 'N/A')
             logger.info(f"Received PIL Image from Gradio. Mode: {original_mode}, Original Format: {original_format}, Size: {original_size}")
             status_messages_display.append(f"[DEBUG] Gradio provided PIL Image - Mode: {original_mode}, Original Format: {original_format}, Size: {original_size}")
-
-            # Save the image received from Gradio to a temporary file for external checking
-            # Ensure the temp directory is writable in the Hugging Face Space environment
-            # /tmp/ is usually writable.
             temp_dir = Path(tempfile.gettempdir()) / "gradio_image_debug"
             temp_dir.mkdir(parents=True, exist_ok=True)
             debug_image_path = temp_dir / "image_from_gradio.png"
-            
-            # Save the image. If it's RGBA, try to keep it RGBA.
-            # PNG is lossless, so this should be okay if the mode is preserved.
             image_pil_object.save(debug_image_path, "PNG") 
             debug_image_path_str = str(debug_image_path)
-            status_messages_display.append(f"**[DEBUG ACTION REQUIRED]** Image received from Gradio has been saved to: `{debug_image_path_str}` in the Space's temporary filesystem. Please try decoding this *specific file* using the standalone `keylock-decode` CLI to verify its integrity. E.g., `keylock-decode {debug_image_path_str} \"YOUR_PASSWORD\"`")
+            status_messages_display.append(f"**[DEBUG ACTION REQUIRED]** Image from Gradio saved to: `{debug_image_path_str}`. Try decoding this file with the standalone `keylock-decode` CLI.")
             logger.info(f"Debug image from Gradio saved to: {debug_image_path_str}")
         else:
-            logger.warning("image_pil_object is None when trying to debug-save.")
-            status_messages_display.append("[DEBUG] No PIL image object received to save for debugging.")
-
+            status_messages_display.append("[DEBUG] No PIL image object received for debug-saving.")
     except Exception as save_exc:
         logger.exception("Error during debug-saving of image from Gradio:")
-        status_messages_display.append(f"[DEBUG] Error saving image for debugging: {str(save_exc)}. Path was: {debug_image_path_str}")
-    # --- END DEBUGGING ---
-
+        status_messages_display.append(f"[DEBUG] Error saving image for debugging: {str(save_exc)}. Path: {debug_image_path_str}")
     try:
-        logger.info(f"Attempting to decode from image using KeyLock-Decode...")
-        # Pass the original PIL object
-        decoded_data, status_msgs_from_lib = decode_from_image_pil(
-            image_pil_object, # Use the PIL object directly
-            password,
-            set_environment_variables=True
-        )
-
+        decoded_data, status_msgs_from_lib = decode_from_image_pil(image_pil_object, password, set_environment_variables=True)
         status_messages_display.extend(status_msgs_from_lib)
-
         if decoded_data:
             status_messages_display.append("\n**Decoded Data Summary (sensitive values masked):**")
             for key, value in decoded_data.items():
                 display_value = '********' if any(k_word in key.upper() for k_word in ['TOKEN', 'KEY', 'SECRET', 'PASS']) else value
                 status_messages_display.append(f"- {key}: {display_value}")
-            
             if os.getenv('HF_TOKEN'):
-                status_messages_display.append(f"\n**SUCCESS: HF_TOKEN was found and has been set in the current process environment.** It will be prioritized for operations.")
+                status_messages_display.append(f"\n**SUCCESS: HF_TOKEN was found and set in environment.**")
             elif 'HF_TOKEN' in decoded_data:
-                 status_messages_display.append(f"\nWarning: HF_TOKEN was in decoded data but os.getenv('HF_TOKEN') is not picking it up. This is unexpected if 'set_environment_variables=True' worked.")
+                 status_messages_display.append(f"\nWarning: HF_TOKEN decoded but os.getenv('HF_TOKEN') not found.")
             else:
-                status_messages_display.append("\nNote: HF_TOKEN was not specifically found in the decoded data from the image.")
-        # else:
-        #     if not any("No payload data found" in msg for msg in status_msgs_from_lib) and not any("Failed to decrypt" in msg for msg in status_msgs_from_lib): # Avoid redundant messages
-        #          status_messages_display.append("No data was decoded, or the decoded data was empty (and no specific error from library).")
-
-
+                status_messages_display.append("\nNote: HF_TOKEN not specifically found in decoded image data.")
     except ValueError as e:
-        logger.error(f"KeyLock-Decode ValueError: {e}")
         status_messages_display.append(f"**Decoding Error (e.g., bad password, corrupted data):** {e}")
     except Exception as e:
-        logger.exception("An unexpected error occurred during image decoding with KeyLock-Decode:")
         status_messages_display.append(f"**An unexpected error occurred during decoding:** {str(e)}")
-    
     return "\n".join(status_messages_display)
 
 
-# --- `parse_markdown` (Unchanged) ---
+# --- Updated `parse_markdown` function ---
 def parse_markdown(markdown_input):
-    """Parse markdown input to extract space details and file structure."""
+    """
+    Parse markdown input to extract space details and file structure.
+    Correctly handles code blocks by excluding the ``` fences.
+    """
     space_info = {"repo_name_md": "", "owner_md": "", "files": []}
-    current_file = None
-    file_content = []
-    in_file_content = False 
-    in_code_block = False 
+    current_file_path = None
+    current_file_content_lines = []
+    
+    # State machine variables
+    # in_file_definition: True when we are after "### File: " and before the next "### File: " or end of input.
+    # in_code_block: True when we are between ``` and ```.
+    in_file_definition = False
+    in_code_block = False
 
     lines = markdown_input.strip().split("\n")
-    for line_idx, line_content_orig in enumerate(lines):
-        line_content_stripped = line_content_orig.strip()
 
-        if in_file_content:
-            if line_content_stripped.startswith("```"):
-                if in_code_block: 
-                    file_content.append(line_content_orig) 
-                    in_code_block = False
-                else: 
-                    in_code_block = True
-                    file_content.append(line_content_orig)
-            elif in_code_block: 
-                file_content.append(line_content_orig)
-            elif not in_code_block: 
-                if line_content_stripped.startswith("### File:") or line_content_stripped.startswith("## File Structure") or line_content_stripped.startswith("# Space:"):
-                    if current_file and file_content:
-                        space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
-                    current_file = None 
-                    file_content = []
-                    in_file_content = False 
-                else: 
-                     file_content.append(line_content_orig)
+    for line_content_orig in lines:
+        line_content_stripped = line_content_orig.strip()
 
-        if line_content_stripped.startswith("# Space:"):
-            if current_file and file_content: 
-                space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
-            full_space_name_md = line_content_stripped.replace("# Space:", "").strip()
-            if "/" in full_space_name_md:
-                space_info["owner_md"], space_info["repo_name_md"] = full_space_name_md.split("/", 1)
+        # Check for new file definition first
+        if line_content_stripped.startswith("### File:"):
+            # Save previous file's content if any
+            if current_file_path and in_file_definition:
+                space_info["files"].append({
+                    "path": current_file_path,
+                    "content": "\n".join(current_file_content_lines)
+                })
+                current_file_content_lines = [] # Reset for the new file
+
+            current_file_path = line_content_stripped.replace("### File:", "").strip()
+            in_file_definition = True
+            in_code_block = False # Reset code block state for a new file
+            continue # Move to next line, don't process "### File:" as content
+
+        # If we are not in any file definition, skip lines until we find one
+        if not in_file_definition:
+            if line_content_stripped.startswith("# Space:"): # Parse space name from markdown
+                full_space_name_md = line_content_stripped.replace("# Space:", "").strip()
+                if "/" in full_space_name_md:
+                    space_info["owner_md"], space_info["repo_name_md"] = full_space_name_md.split("/", 1)
+                else:
+                    space_info["repo_name_md"] = full_space_name_md
+            # Other top-level constructs like "## File Structure" can be ignored or parsed if needed
+            continue
+
+        # Now, we are definitely within a "### File: ..." block's content area
+        
+        # Handle code block fences ```
+        if line_content_stripped.startswith("```"):
+            if in_code_block:
+                in_code_block = False # Exiting a code block
             else:
-                space_info["repo_name_md"] = full_space_name_md
-            current_file = None
-            file_content = []
-            in_file_content = False
-            in_code_block = False
-
-        elif line_content_stripped.startswith("## File Structure"):
-            if current_file and file_content: 
-                space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
-            current_file = None
-            file_content = []
-            in_file_content = False
-            in_code_block = False
-            continue 
-
-        elif line_content_stripped.startswith("### File:"):
-            if current_file and file_content: 
-                space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
-            
-            current_file = line_content_stripped.replace("### File:", "").strip()
-            file_content = [] 
-            in_file_content = True 
-            in_code_block = False 
-            
-    if current_file and file_content:
-        space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
-    
+                in_code_block = True  # Entering a code block
+            # In either case, we *don't* append the ``` line itself to content
+            continue # Move to next line
+
+        # If we are inside a code block OR not (i.e., plain text lines for the file),
+        # append the original line (preserving leading/trailing whitespace for that line).
+        current_file_content_lines.append(line_content_orig)
+
+    # After loop, append the last file's content if any
+    if current_file_path and in_file_definition: # Ensure we were actually defining a file
+        space_info["files"].append({
+            "path": current_file_path,
+            "content": "\n".join(current_file_content_lines)
+        })
+
+    # Filter out any entries that might have been added with no path (shouldn't happen with current logic)
     space_info["files"] = [f for f in space_info["files"] if f.get("path")]
     return space_info
 
+
 # --- `_determine_repo_id` (Unchanged) ---
 def _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui):
-    if not space_name_ui:
-        return None, "Error: Space Name cannot be empty."
-    if "/" in space_name_ui:
-        return None, "Error: Space Name should not contain '/'. Please use the Owner field for the namespace."
-
+    if not space_name_ui: return None, "Error: Space Name cannot be empty."
+    if "/" in space_name_ui: return None, "Error: Space Name should not contain '/'. Use Owner field."
     final_owner = owner_ui; error_message = None
     if not final_owner:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
         if token_err: return None, token_err 
-        if not resolved_api_token: return None, "Error: API token required for auto owner determination (internal check failed)."
+        if not resolved_api_token: return None, "Error: API token required for auto owner determination."
         try:
             user_info = whoami(token=resolved_api_token)
-            if user_info and 'name' in user_info:
-                final_owner = user_info['name']
-                logger.info(f"Determined owner: {final_owner} using API token.")
-            else:
-                error_message = "Error: Could not retrieve username from API token. Check permissions or specify Owner."
-        except Exception as e:
-            error_message = f"Error retrieving username from API token: {str(e)}. Specify Owner manually."
+            if user_info and 'name' in user_info: final_owner = user_info['name']
+            else: error_message = "Error: Could not retrieve username. Check token/permissions or specify Owner."
+        except Exception as e: error_message = f"Error retrieving username: {str(e)}. Specify Owner."
         if error_message: return None, error_message
     if not final_owner: return None, "Error: Owner could not be determined."
     return f"{final_owner}/{space_name_ui}", None
 
-# --- Core Functions: `create_space`, `view_space_files`, `update_space_file` (Unchanged from previous correct version) ---
+# --- Core Functions: `create_space`, `view_space_files`, `update_space_file` (Unchanged) ---
 def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, markdown_input):
     repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
@@ -228,7 +188,7 @@ def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, mar
         repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
         if err: return err
         repo_id_for_error_logging = repo_id
-        space_info = parse_markdown(markdown_input)
+        space_info = parse_markdown(markdown_input) # Uses the updated parser
         if not space_info["files"]: return "Error: No files found in markdown. Use '### File: path/to/file.ext'."
         with tempfile.TemporaryDirectory() as temp_dir:
             repo_local_path = Path(temp_dir) / "repo_upload_content"
@@ -237,7 +197,7 @@ def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, mar
                 if not file_info.get("path"): continue
                 file_path_abs = repo_local_path / file_info["path"]
                 file_path_abs.parent.mkdir(parents=True, exist_ok=True)
-                with open(file_path_abs, "w", encoding="utf-8") as f: f.write(file_info["content"])
+                with open(file_path_abs, "w", encoding="utf-8") as f: f.write(file_info["content"]) # Content is now clean
             try:
                 create_repo(repo_id=repo_id, token=resolved_api_token, repo_type="space", space_sdk=sdk_ui, private=False)
             except Exception as e: