Update app_logic.py

app_logic.py

@@ -1,7 +1,8 @@
 import os
 import re
 import tempfile
-import git
+import shutil # For rmtree
+import git # Used by Repository indirectly
 from huggingface_hub import (
     create_repo,
     upload_folder,
@@ -13,13 +14,14 @@ import logging
 from pathlib import Path
 from PIL import Image 
 
+# Attempt to import keylock_decode
 try:
     from keylock_decode import decode_from_image_pil
     KEYLOCK_DECODE_AVAILABLE = True
 except ImportError:
     KEYLOCK_DECODE_AVAILABLE = False
     decode_from_image_pil = None
-    logging.warning("keylock-decode library not found. Image decoding feature will be disabled.")
+    logging.warning("keylock-decode library not found. KeyLock Wallet image feature will be disabled.")
 
 logging.basicConfig(
     level=logging.INFO, 
@@ -37,36 +39,21 @@ def _get_api_token(ui_token_from_textbox=None):
         logger.info("Using API token from UI textbox.")
         return 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."
+    return None, "Error: Hugging Face API token not provided. Please enter it or load from a KeyLock Wallet image."
 
-# --- `load_token_from_image_and_set_env` (Unchanged from previous debug version) ---
+# --- `load_token_from_image_and_set_env` (Unchanged from previous debug version, ensure debug lines are suitable) ---
 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 image_pil_object is None: return "Error: No KeyLock Wallet image provided for decoding."
     if not password: return "Error: Password cannot be empty for image decoding."
     status_messages_display = []
-    debug_image_path_str = "Not saved (error or no image)."
-    try:
-        if image_pil_object:
-            original_mode = getattr(image_pil_object, 'mode', 'N/A')
-            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}")
-            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"
-            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 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:
-            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: {debug_image_path_str}")
+    # Optional: Keep debug saving if still needed, otherwise remove for cleaner output
+    # debug_image_path_str = "Not saved." 
+    # try: ... debug save logic ...
+    # except Exception as save_exc: ...
     try:
+        logger.info(f"Attempting to decode from KeyLock Wallet image...")
         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:
@@ -75,92 +62,48 @@ def load_token_from_image_and_set_env(image_pil_object: Image.Image, password: s
                 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 set in environment.**")
+                status_messages_display.append(f"\n**SUCCESS: HF_TOKEN was found and set in environment from KeyLock Wallet image.**")
             elif 'HF_TOKEN' in decoded_data:
-                 status_messages_display.append(f"\nWarning: HF_TOKEN decoded but os.getenv('HF_TOKEN') not found.")
+                 status_messages_display.append(f"\nWarning: HF_TOKEN decoded but os.getenv('HF_TOKEN') not found (unexpected).")
             else:
-                status_messages_display.append("\nNote: HF_TOKEN not specifically found in decoded image data.")
-    except ValueError as e:
+                status_messages_display.append("\nNote: HF_TOKEN not specifically found in decoded KeyLock Wallet image data.")
+    except ValueError as e: # Specific errors from keylock-decode
         status_messages_display.append(f"**Decoding Error (e.g., bad password, corrupted data):** {e}")
     except Exception as e:
+        logger.exception("Unexpected error during KeyLock Wallet image decoding:")
         status_messages_display.append(f"**An unexpected error occurred during decoding:** {str(e)}")
     return "\n".join(status_messages_display)
 
-
-# --- Updated `parse_markdown` function ---
+# --- `parse_markdown` (Unchanged from previous corrected version) ---
 def parse_markdown(markdown_input):
-    """
-    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_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
-
+    current_file_path = None; current_file_content_lines = []
+    in_file_definition = False; in_code_block = False
     lines = markdown_input.strip().split("\n")
-
     for line_content_orig in lines:
         line_content_stripped = line_content_orig.strip()
-
-        # 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
-
+                space_info["files"].append({"path": current_file_path, "content": "\n".join(current_file_content_lines)})
             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
+            current_file_content_lines = []
+            in_file_definition = True; in_code_block = False
+            continue
         if not in_file_definition:
-            if line_content_stripped.startswith("# Space:"): # Parse space name from markdown
+            if line_content_stripped.startswith("# Space:"):
                 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
+                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
             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:
-                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).
+            in_code_block = not in_code_block
+            continue
         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)
+    if current_file_path and in_file_definition:
+        space_info["files"].append({"path": current_file_path, "content": "\n".join(current_file_content_lines)})
     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."
@@ -179,8 +122,80 @@ def _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui):
     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) ---
+# --- New/Modified Functions for File Browsing and Editing ---
+CLONE_BASE_DIR = Path(tempfile.gettempdir()) / "space_builder_active_clones"
+CLONE_BASE_DIR.mkdir(parents=True, exist_ok=True)
+
+def _cleanup_old_clones(repo_id_slug_to_keep=None):
+    """Clean up old clone directories, optionally keeping one specific repo's clone."""
+    try:
+        for item in CLONE_BASE_DIR.iterdir():
+            if item.is_dir():
+                if repo_id_slug_to_keep and item.name == repo_id_slug_to_keep:
+                    continue
+                logger.info(f"Cleaning up old clone: {item}")
+                shutil.rmtree(item)
+    except Exception as e:
+        logger.error(f"Error during old clone cleanup: {e}")
+
+def get_space_local_clone_path(ui_api_token_from_textbox, space_name_ui, owner_ui, force_refresh=False):
+    """Clones a Space locally, manages cleanup, and returns the path."""
+    repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
+    try:
+        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
+        if token_err: return None, token_err
+
+        repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
+        if err: return None, err
+        repo_id_for_error_logging = repo_id
+
+        repo_id_slug = repo_id.replace("/", "_") # Sanitize for directory name
+        _cleanup_old_clones(repo_id_slug_to_keep=repo_id_slug)
+        local_clone_dir = CLONE_BASE_DIR / repo_id_slug
+
+        if force_refresh and local_clone_dir.exists():
+            logger.info(f"Force refresh: Removing existing clone at {local_clone_dir}")
+            shutil.rmtree(local_clone_dir)
+        
+        if not local_clone_dir.exists():
+            local_clone_dir.mkdir(parents=True, exist_ok=True) # Ensure parent exists before Repository tries to use it
+            logger.info(f"Cloning Space {repo_id} to {local_clone_dir}...")
+            try:
+                Repository(
+                    local_dir=str(local_clone_dir),
+                    clone_from=f"https://huggingface.co/spaces/{repo_id}",
+                    repo_type="space", use_auth_token=resolved_api_token,
+                )
+                logger.info(f"Successfully cloned {repo_id} to {local_clone_dir}")
+            except Exception as clone_exc:
+                logger.exception(f"Error cloning {repo_id}:")
+                if local_clone_dir.exists(): shutil.rmtree(local_clone_dir) # Cleanup partial clone
+                return None, f"Error cloning Space '{repo_id}': {str(clone_exc)}"
+        else:
+            logger.info(f"Using existing clone for {repo_id} at {local_clone_dir}")
+            # For a true refresh of an existing clone, you'd `git pull` here.
+            # `force_refresh=True` handles this by re-cloning.
+
+        return str(local_clone_dir), None
+    except Exception as e:
+        logger.exception(f"Error in get_space_local_clone_path for {repo_id_for_error_logging}:")
+        return None, f"Error preparing local clone: {str(e)}"
+
+def read_file_from_local_path(absolute_local_file_path_str: str):
+    """Reads content of a file given its absolute local path."""
+    try:
+        file_path = Path(absolute_local_file_path_str)
+        if not file_path.is_file():
+            return None, f"Error: Path is not a file or does not exist: {absolute_local_file_path_str}"
+        content = file_path.read_text(encoding="utf-8")
+        return content, None
+    except Exception as e:
+        logger.exception(f"Error reading local file {absolute_local_file_path_str}:")
+        return None, f"Error reading file content: {str(e)}"
+
+# --- Core Functions: `create_space`, `update_space_file` (view_space_files can be deprecated) ---
 def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, markdown_input):
+    # (Unchanged from previous correct version)
     repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
@@ -188,61 +203,72 @@ 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) # Uses the updated parser
+        space_info = parse_markdown(markdown_input)
         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"
-            repo_local_path.mkdir(exist_ok=True)
+        with tempfile.TemporaryDirectory() as temp_dir: # Temp dir for staging files from markdown
+            repo_staging_path = Path(temp_dir) / "repo_staging_content"
+            repo_staging_path.mkdir(exist_ok=True)
             for file_info in space_info["files"]:
                 if not file_info.get("path"): continue
-                file_path_abs = repo_local_path / file_info["path"]
+                file_path_abs = repo_staging_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"]) # Content is now clean
+                with open(file_path_abs, "w", encoding="utf-8") as f: f.write(file_info["content"])
             try:
                 create_repo(repo_id=repo_id, token=resolved_api_token, repo_type="space", space_sdk=sdk_ui, private=False)
             except Exception as e:
                 err_str = str(e).lower()
                 if not ("already exists" in err_str or "you already created this repo" in err_str or "exists" in err_str):
                     return f"Error creating Space '{repo_id}': {str(e)}"
-            upload_folder(repo_id=repo_id, folder_path=str(repo_local_path), path_in_repo=".", token=resolved_api_token, repo_type="space", commit_message=f"Initial Space setup of {repo_id} via Builder")
+                logger.info(f"Space {repo_id} already exists or creation confirmed, proceeding with upload.")
+            upload_folder(repo_id=repo_id, folder_path=str(repo_staging_path), path_in_repo=".", token=resolved_api_token, repo_type="space", commit_message=f"Initial Space setup of {repo_id} via Builder")
             return f"Successfully created/updated Space: [{repo_id}](https://huggingface.co/spaces/{repo_id})"
     except Exception as e:
         logger.exception(f"Error in create_space for {repo_id_for_error_logging}:")
         return f"Error during Space creation/update: {str(e)}"
 
-def view_space_files(ui_api_token_from_textbox, space_name_ui, owner_ui):
-    repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
-    try:
-        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
-        if token_err: return token_err
-        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
-        files = list_repo_files(repo_id=repo_id, token=resolved_api_token, repo_type="space")
-        if files: return f"Files in `{repo_id}`:\n\n" + "\n".join([f"- `{f}`" for f in files])
-        return f"No files found in Space `{repo_id}`."
-    except Exception as e:
-        logger.exception(f"Error in view_space_files for {repo_id_for_error_logging}:")
-        return f"Error listing files for `{repo_id_for_error_logging}`: {str(e)}"
+# view_space_files can be removed or kept as a simple alternative list view
+# def view_space_files(...)
 
 def update_space_file(ui_api_token_from_textbox, space_name_ui, owner_ui, file_path_in_repo, file_content, commit_message_ui):
+    # This function is now simpler as it expects file_path_in_repo to be the correct relative path.
+    # It still does its own clone for transactional integrity.
     repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
         if token_err: return token_err
+
         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
-        if not file_path_in_repo: return "Error: File Path cannot be empty."
+        
+        if not file_path_in_repo: return "Error: File Path to update cannot be empty."
+        # Sanitize file_path_in_repo: remove leading slashes, ensure forward slashes
+        file_path_in_repo = file_path_in_repo.lstrip('/').replace(os.sep, '/')
+
         commit_message_ui = commit_message_ui or f"Update {file_path_in_repo} via Space Builder"
-        with tempfile.TemporaryDirectory() as temp_dir:
-            repo_local_clone_path = Path(temp_dir) / "cloned_space_repo_for_update"
-            cloned_repo = Repository(local_dir=str(repo_local_clone_path), clone_from=f"https://huggingface.co/spaces/{repo_id}", repo_type="space", use_auth_token=resolved_api_token, git_user="Space Builder Bot", git_email="[email protected]")
+
+        with tempfile.TemporaryDirectory() as temp_dir_for_update: # Fresh temp dir for this update operation
+            repo_local_clone_path = Path(temp_dir_for_update) / "update_clone"
+            
+            cloned_repo = Repository(
+                local_dir=str(repo_local_clone_path),
+                clone_from=f"https://huggingface.co/spaces/{repo_id}",
+                repo_type="space", use_auth_token=resolved_api_token,
+                git_user="Space Builder Bot", git_email="[email protected]" # Optional
+            )
+            logger.info(f"Cloned Space {repo_id} to {repo_local_clone_path} for update operation.")
+
+            # The file_path_in_repo is relative to the cloned_repo.local_dir
             full_local_file_path = Path(cloned_repo.local_dir) / file_path_in_repo
-            full_local_file_path.parent.mkdir(parents=True, exist_ok=True)
-            with open(full_local_file_path, "w", encoding="utf-8") as f: f.write(file_content)
+            full_local_file_path.parent.mkdir(parents=True, exist_ok=True) # Create parent dirs if needed
+            with open(full_local_file_path, "w", encoding="utf-8") as f:
+                f.write(file_content)
+            logger.info(f"Wrote updated content to {full_local_file_path} for commit.")
+
             cloned_repo.push_to_hub(commit_message=commit_message_ui)
+            logger.info(f"Pushed update for {file_path_in_repo} to {repo_id}")
+            
             return f"Successfully updated `{file_path_in_repo}` in Space [{repo_id}](https://huggingface.co/spaces/{repo_id})"
     except Exception as e:
-        logger.exception(f"Error in update_space_file for {repo_id_for_error_logging}:")
+        logger.exception(f"Error in update_space_file for {repo_id_for_error_logging}, file {file_path_in_repo}:")
         return f"Error updating file for `{repo_id_for_error_logging}`: {str(e)}"