Update app.py

app.py

@@ -14,6 +14,7 @@ from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_st
 # --- Helper functions for Hugging Face integration ---
 
 def show_profile(profile: gr.OAuthProfile | None) -> str:
+    """Displays the logged-in Hugging Face profile username."""
     if profile is None:
         return "*Not logged in.*"
     return f"✅ Logged in as **{profile.username}**"
@@ -22,6 +23,7 @@ def list_private_models(
     profile: gr.OAuthProfile | None,
     oauth_token: gr.OAuthToken | None
 ) -> str:
+    """Lists private models for the logged-in user (not used in the main workflow, but kept)."""
     if profile is None or oauth_token is None:
         return "Please log in to see your models."
     try:
@@ -31,15 +33,19 @@ def list_private_models(
         ]
         return "No models found." if not models else "Models:\n\n" + "\n - ".join(models)
     except Exception as e:
+        # Catching generic exception is acceptable for helper functions
         return f"Error listing models: {e}"
 
 def create_space_action(repo_name: str, sdk: str, profile: gr.OAuthProfile, token: gr.OAuthToken):
+    """Creates a new Hugging Face Space repository."""
+    if not profile or not token:
+        raise ValueError("Hugging Face profile or token is missing.")
     repo_id = f"{profile.username}/{repo_name}"
     try:
         create_repo(
             repo_id=repo_id,
             token=token.token,
-            exist_ok=True,
+            exist_ok=True, # Allow creating if it already exists
             repo_type="space",
             space_sdk=sdk
         )
@@ -50,12 +56,13 @@ def create_space_action(repo_name: str, sdk: str, profile: gr.OAuthProfile, toke
         raise RuntimeError(f"Failed to create Space {repo_id}: {e}")
 
 def upload_file_to_space_action(
-    file_obj,
+    file_obj: io.StringIO, # Specify type hint for clarity
     path_in_repo: str,
     repo_id: str,
     profile: gr.OAuthProfile,
     token: gr.OAuthToken
 ) -> None:
+    """Uploads a file to a Hugging Face Space repository."""
     if not (profile and token and repo_id):
         raise ValueError("Hugging Face profile, token, or repo_id is missing.")
     try:
@@ -70,22 +77,24 @@ def upload_file_to_space_action(
         raise RuntimeError(f"Failed to upload `{path_in_repo}` to {repo_id}: {e}")
 
 def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
+    """Fetches build or run logs for a Space."""
     if not repo_id or not token:
          return f"Cannot fetch {level} logs: repo_id or token missing."
     jwt_url  = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
     try:
         r        = get_session().get(jwt_url, headers=build_hf_headers(token=token))
-        hf_raise_for_status(r)
+        hf_raise_for_status(r) # Raise HTTPError for bad responses (4xx or 5xx)
         jwt      = r.json()["token"]
         logs_url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
         lines, count = [], 0
+        # Using stream=True is good for potentially large logs
         with get_session().get(logs_url, headers=build_hf_headers(token=jwt), stream=True, timeout=30) as resp:
             hf_raise_for_status(resp)
             for raw in resp.iter_lines():
-                if count >= 200:
+                if count >= 200: # Limit output lines to prevent UI overload
                     lines.append("... truncated ...")
                     break
-                if not raw.startswith(b"data: "):
+                if not raw.startswith(b"data: "): # EventStream protocol expected from HF logs API
                     continue
                 payload = raw[len(b"data: "):]
                 try:
@@ -96,21 +105,27 @@ def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
                         lines.append(f"[{ts}] {txt}")
                         count += 1
                 except json.JSONDecodeError:
+                    # Skip lines that aren't valid JSON events
                     continue
         return "\n".join(lines) if lines else f"No {level} logs found."
     except Exception as e:
+        # Catching generic exception is acceptable for helper functions
         return f"Error fetching {level} logs: {e}"
 
 
 def get_build_logs_action(repo_id, profile, token):
+    """Action to fetch build logs with a small delay."""
     if not (repo_id and profile and token):
         return "⚠️ Cannot fetch build logs: log in and create a Space first."
+    # Small delay to allow build process to potentially start on HF side
     time.sleep(5)
     return _fetch_space_logs_level(repo_id, "build", token.token)
 
 def get_container_logs_action(repo_id, profile, token):
+    """Action to fetch container logs with a delay."""
     if not (repo_id and profile and token):
         return "⚠️ Cannot fetch container logs: log in and create a Space first."
+    # Longer delay to allow container to start after build completes
     time.sleep(10)
     return _fetch_space_logs_level(repo_id, "run", token.token)
 
@@ -118,32 +133,38 @@ def get_container_logs_action(repo_id, profile, token):
 # --- Google Gemini integration with model selection ---
 
 def configure_gemini(api_key: str | None, model_name: str | None) -> str:
+    """Configures the Gemini API and checks if the model is accessible."""
     if not api_key:
         return "Gemini API key is not set."
     if not model_name:
         return "Please select a Gemini model."
     try:
         genai.configure(api_key=api_key)
+        # Attempt a simple call to verify credentials and model availability
+        # This will raise an exception if the key is invalid or model not found
         genai.GenerativeModel(model_name).generate_content("ping", stream=False)
         return f"Gemini configured successfully with **{model_name}**."
     except Exception as e:
         return f"Error configuring Gemini: {e}"
 
 def call_gemini(prompt: str, api_key: str, model_name: str) -> str:
+    """Calls the Gemini API with a given prompt."""
     if not api_key or not model_name:
         raise ValueError("Gemini API key or model not set.")
     try:
         genai.configure(api_key=api_key)
         model = genai.GenerativeModel(model_name)
+        # Using generate_content and stream=False for simplicity here
         response = model.generate_content(prompt, stream=False)
-        return response.text or ""
+        return response.text or "" # Return empty string if no text
     except Exception as e:
+        # Re-raising as RuntimeError for the workflow to catch and manage
         raise RuntimeError(f"Gemini API call failed: {e}")
 
 
 # --- AI workflow logic (State Machine) ---
 
-# Define States
+# Define States for the workflow
 STATE_IDLE = "idle"
 STATE_AWAITING_REPO_NAME = "awaiting_repo_name"
 STATE_CREATING_SPACE = "creating_space"
@@ -159,13 +180,14 @@ STATE_DEBUGGING_CODE = "debugging_code"
 STATE_UPLOADING_FIXED_APP_PY = "uploading_fixed_app_py"
 STATE_COMPLETE = "complete"
 
-MAX_DEBUG_ATTEMPTS = 3
+MAX_DEBUG_ATTEMPTS = 3 # Limit the number of automatic debug attempts
 
 def add_bot_message(history: list[dict], bot_message: str) -> list[dict]:
-    """Helper to add a new assistant message to the history."""
+    """Helper to add a new assistant message to the chatbot history."""
     history.append({"role": "assistant", "content": bot_message})
     return history
 
+# This is the main generator function for the workflow, triggered by the 'Send' button
 def ai_workflow_chat(
     message: str,
     history: list[dict],
@@ -176,28 +198,33 @@ def ai_workflow_chat(
     repo_id_state:  str | None,
     workflow_state: str,
     space_sdk:      str,
-    preview_html:   str,
-    container_logs: str,
-    build_logs:     str,
+    preview_html:   str, # Passed in to maintain its value in State
+    container_logs: str, # Passed in to maintain its value in State
+    build_logs:     str, # Passed in to maintain its value in State
     debug_attempts_state: int,
     app_description_state: str | None,
     repo_name_state: str | None,
     generated_code_state: str | None,
+    # Absorb potential extra args passed by Gradio event listeners (e.g. old value, event data)
     *args,
     **kwargs
 ) -> tuple[
-    list[dict],
-    str | None,
-    str,
-    str,
-    str,
-    str,
-    int,
-    str | None,
-    str | None,
-    str | None,
+    list[dict],       # 0: Updated chat history
+    str | None,       # 1: Updated repo_id
+    str,              # 2: Updated workflow state
+    str,              # 3: Updated iframe HTML
+    str,              # 4: Updated container logs
+    str,              # 5: Updated build logs
+    int,              # 6: Updated debug attempts count
+    str | None,       # 7: Updated app description
+    str | None,       # 8: Updated repo name
+    str | None,       # 9: Updated generated code (for temporary storage)
 ]:
-    # Unpack state variables
+    """
+    Generator function to handle the AI workflow state machine.
+    Each 'yield' pauses execution and sends values to update Gradio outputs/state.
+    """
+    # Unpack state variables from Gradio State components passed as inputs
     repo_id = repo_id_state
     state = workflow_state
     attempts = debug_attempts_state
@@ -205,305 +232,462 @@ def ai_workflow_chat(
     repo_name = repo_name_state
     generated_code = generated_code_state
 
-
+    # Keep copies of potentially updated UI elements passed as inputs to update them later
     updated_preview = preview_html
     updated_build = build_logs
     updated_run = container_logs
 
-    # Add user message to history in the new format
+    # Add the user's message to the chat history immediately
     user_message_entry = {"role": "user", "content": message}
+    # Add username if logged in (optional, but nice)
     if hf_profile and hf_profile.username:
          user_message_entry["name"] = hf_profile.username
     history.append(user_message_entry)
 
-    # Yield immediately to show user message
-    yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
+    # Yield immediately to update the chat UI with the user's message
+    # This provides immediate feedback to the user while the AI processes
+    yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+           attempts, app_desc, repo_name, generated_code)
 
     try:
-        # --- State Machine Logic ---
+        # --- State Machine Logic based on the current 'state' variable ---
 
         if state == STATE_IDLE:
-            # Check prerequisites within the state logic as well
+            # Check prerequisites before starting any workflow actions
             if not (hf_profile and hf_token):
                 history = add_bot_message(history, "Workflow paused: Please log in to Hugging Face first.")
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-                return
+                # Yield updated history and current state, then exit for this click
+                yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                       attempts, app_desc, repo_name, generated_code)
+                return # Exit the generator for this click
+
             if not (gemini_api_key and gemini_model):
                  history = add_bot_message(history, "Workflow paused: Please enter your API key and select a Gemini model.")
-                 yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-                 return
+                 # Yield updated history and current state, then exit for this click
+                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                        attempts, app_desc, repo_name, generated_code)
+                 return # Exit the generator for this click
 
-            # Look for commands
+            # Look for specific commands in the user's message
             reset_match = "reset" in message.lower()
-            generate_match = re.search(r'generate (?:me )?(?:a|an) \w+ app called (\w+)', message, re.I)
+            # Capture app description AND repo name using regex
+            generate_match = re.search(r'generate (?:me )?(?:a|an) (.+) app called (\w+)', message, re.I)
+            # Capture repo name for a simple 'create space' command
             create_match = re.search(r'create (?:a|an)? space called (\w+)', message, re.I)
 
             if reset_match:
+                # Reset the workflow state and associated variables
                 history = add_bot_message(history, "Workflow reset.")
-                yield history, None, STATE_IDLE, "<p>No Space created yet.</p>", "", "", 0, None, None, None
-                return
+                # Yield updated history and reset state variables to their initial values
+                yield (history, None, STATE_IDLE, "<p>No Space created yet.</p>", "", "", 0,
+                       None, None, None)
+                # No return needed after yield in this generator pattern; execution for this click ends here.
 
             elif generate_match:
-                new_repo_name = generate_match.group(1)
-                new_app_desc = message
-                history = add_bot_message(history, f"Acknowledged: '{message}'. Starting workflow to create Space `{hf_profile.username}/{new_repo_name}`.")
+                # User requested generation with description and name
+                new_app_desc = generate_match.group(1).strip() # Capture description part
+                new_repo_name = generate_match.group(2).strip() # Capture name part
+                history = add_bot_message(history, f"Acknowledged: '{message}'. Starting workflow to create Space `{hf_profile.username}/{new_repo_name}` for a '{new_app_desc}' app.")
+                # Update state variables for the next step (creation)
                 state = STATE_CREATING_SPACE
                 repo_name = new_repo_name
                 app_desc = new_app_desc
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
+                # Yield updated history and state variables
+                yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                       attempts, app_desc, repo_name, generated_code)
+                # No return needed
 
             elif create_match:
-                 new_repo_name = create_match.group(1)
+                 # User requested simple space creation with a name
+                 new_repo_name = create_match.group(1).strip()
                  history = add_bot_message(history, f"Acknowledged: '{message}'. Starting workflow to create Space `{hf_profile.username}/{new_repo_name}`.")
+                 # Update state variables for the next step (creation)
                  state = STATE_CREATING_SPACE
                  repo_name = new_repo_name
-                 yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                 # Yield updated history and state variables
+                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                        attempts, app_desc, repo_name, generated_code)
+                 # No return needed
 
             elif "create" in message.lower() and not repo_id:
+                # User wants to create but didn't specify a name yet
                 history = add_bot_message(history, "Okay, what should the Space be called? (e.g., `my-awesome-app`)")
-                state = STATE_AWAITING_REPO_NAME
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                state = STATE_AWAITING_REPO_NAME # Transition to the state where we wait for the name
+                # Yield updated history and state
+                yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                       attempts, app_desc, repo_name, generated_code)
+                # No return needed
 
             else:
+                # Command not recognized in IDLE state
                 history = add_bot_message(history, "Command not recognized. Try 'generate me a gradio app called myapp', or 'reset'.")
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-                return
+                # Yield updated history and current state
+                yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                       attempts, app_desc, repo_name, generated_code)
+                # No return needed
 
 
-        if state == STATE_AWAITING_REPO_NAME:
-            new_repo_name = message.strip()
-            if not new_repo_name or re.search(r'[^a-zA-Z0-9_-]', new_repo_name):
+        elif state == STATE_AWAITING_REPO_NAME:
+             # User's message is expected to be the repo name
+             new_repo_name = message.strip()
+             # Basic validation for Hugging Face repo name format
+             if not new_repo_name or re.search(r'[^a-zA-Z0-9_-]', new_repo_name):
                  history = add_bot_message(history, "Invalid name. Please provide a single word/slug for the Space name (letters, numbers, underscores, hyphens only).")
-                 yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-            else:
-                history = add_bot_message(history, f"Using Space name `{new_repo_name}`. Creating Space `{hf_profile.username}/{new_repo_name}`...")
-                state = STATE_CREATING_SPACE
-                repo_name = new_repo_name
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                 # Stay in AWAITING_REPO_NAME state and yield message
+                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                        attempts, app_desc, repo_name, generated_code)
+                 # No return needed
 
+             else:
+                 history = add_bot_message(history, f"Using Space name `{new_repo_name}`. Creating Space `{hf_profile.username}/{new_repo_name}`...")
+                 state = STATE_CREATING_SPACE # Transition state to creation
+                 repo_name = new_repo_name # Store the validated repo name
+                 # Yield updated history, state, and repo name.
+                 # The next click will proceed from the STATE_CREATING_SPACE block.
+                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                        attempts, app_desc, repo_name, generated_code)
+                 # No return needed
+
+        # Note: Each 'elif' block below represents a distinct step in the workflow triggered
+        # when the 'state' variable matches its condition on a button click.
 
         elif state == STATE_CREATING_SPACE:
-            if not repo_name:
-                history = add_bot_message(history, "Internal error: Repo name missing for creation. Resetting.")
-                yield history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0, None, None, None
-                return
-
-            try:
-                new_repo_id, iframe_html = create_space_action(repo_name, space_sdk, hf_profile, hf_token)
-                updated_preview = iframe_html
-                repo_id = new_repo_id
-                history = add_bot_message(history, f"✅ Space `{repo_id}` created. Click 'Send' to generate and upload code.")
-                state = STATE_GENERATING_CODE
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+             # Ensure repo_name is available (it should have been set in the previous step)
+             if not repo_name:
+                 history = add_bot_message(history, "Internal error: Repo name missing for creation. Resetting.")
+                 yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
+                       None, None, None)
+                 # No return needed
 
-            except Exception as e:
-                history = add_bot_message(history, f"❌ Error creating space: {e}. Click 'reset'.")
-                yield history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0, None, None, None
+             else:
+                 try:
+                     # Perform the action to create the Space on Hugging Face
+                     new_repo_id, iframe_html = create_space_action(repo_name, space_sdk, hf_profile, hf_token)
+                     updated_preview = iframe_html # Update the iframe content to show the new space
+                     repo_id = new_repo_id # Store the official repo_id
+                     history = add_bot_message(history, f"✅ Space `{repo_id}` created. Click 'Send' to generate and upload code.")
+                     state = STATE_GENERATING_CODE # Transition to the next state
+                     # Yield updated state variables and history
+                     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                            attempts, app_desc, repo_name, generated_code)
+                     # No return needed
+
+                 except Exception as e:
+                     history = add_bot_message(history, f"❌ Error creating space: {e}. Click 'reset'.")
+                     # Yield error message and reset state on failure
+                     yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
+                           None, None, None)
+                     # No return needed
 
 
         elif state == STATE_GENERATING_CODE:
-            prompt_desc = app_desc if app_desc else 'a Gradio image-blur test app with upload and slider controls'
-            prompt = f"""
+             # Define the prompt for Gemini based on the app description or a default
+             prompt_desc = app_desc if app_desc else 'a Gradio image-blur test app with upload and slider controls'
+             prompt = f"""
 You are an AI assistant specializing in Hugging Face Spaces using the {space_sdk} SDK.
 Generate a full, single-file Python app based on:
 '{prompt_desc}'
 Return **only** the python code block for app.py. Do not include any extra text, explanations, or markdown outside the code block.
 """
-            try:
-                history = add_bot_message(history, "🧠 Generating `app.py` code with Gemini...")
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
-                code = call_gemini(prompt, gemini_api_key, gemini_model)
-                code = code.strip()
-                if code.startswith("```python"):
-                    code = code[len("```python"):].strip()
-                if code.endswith("```"):
-                    code = code[:-len("```")].strip()
-
-                if not code:
-                     raise ValueError("Gemini returned empty code.")
-
-                history = add_bot_message(history, "✅ `app.py` code generated. Click 'Send' to upload.")
-                state = STATE_UPLOADING_APP_PY
-                generated_code = code
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+             try:
+                 history = add_bot_message(history, "🧠 Generating `app.py` code with Gemini...")
+                 # Yield to show message before the potentially time-consuming API call
+                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                        attempts, app_desc, repo_name, generated_code)
+
+                 # Perform the Gemini API call to generate code
+                 code = call_gemini(prompt, gemini_api_key, gemini_model)
+                 code = code.strip()
+                 # Clean up common markdown code block formatting if present
+                 if code.startswith("```python"):
+                     code = code[len("```python"):].strip()
+                 if code.endswith("```"):
+                     code = code[:-len("```")].strip()
+
+                 if not code:
+                      raise ValueError("Gemini returned empty code.")
+
+                 history = add_bot_message(history, "✅ `app.py` code generated. Click 'Send' to upload.")
+                 state = STATE_UPLOADING_APP_PY # Transition to the upload state
+                 generated_code = code # Store the generated code in the state variable for the next step
+                 # Yield updated state variables and history
+                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                        attempts, app_desc, repo_name, generated_code)
+                 # No return needed
 
-            except Exception as e:
-                history = add_bot_message(history, f"❌ Error generating code: {e}. Click 'reset'.")
-                yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
+             except Exception as e:
+                 history = add_bot_message(history, f"❌ Error generating code: {e}. Click 'reset'.")
+                 # Yield error message and reset state on failure
+                 yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                       None, None, None)
+                 # No return needed
 
 
         elif state == STATE_UPLOADING_APP_PY:
-            if not generated_code:
-                 history = add_bot_message(history, "Internal error: No code to upload. Resetting.")
-                 yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
-                 return
-
-            history = add_bot_message(history, "☁️ Uploading `app.py`...")
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+             # Retrieve the generated code from the state variable
+             code_to_upload = generated_code
+             if not code_to_upload:
+                  history = add_bot_message(history, "Internal error: No code to upload. Resetting.")
+                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                        None, None, None)
+                  # No return needed
 
-            try:
-                upload_file_to_space_action(io.StringIO(generated_code), "app.py", repo_id, hf_profile, hf_token)
-                history = add_bot_message(history, "✅ Uploaded `app.py`. Click 'Send' to generate requirements.")
-                state = STATE_GENERATING_REQUIREMENTS
-                generated_code = None
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
-            except Exception as e:
-                history = add_bot_message(history, f"❌ Error uploading app.py: {e}. Click 'reset'.")
-                yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
+             else:
+                 history = add_bot_message(history, "☁️ Uploading `app.py`...")
+                 # Yield to show message before the upload action
+                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                        attempts, app_desc, repo_name, generated_code)
+
+                 try:
+                     # Perform the file upload action
+                     upload_file_to_space_action(io.StringIO(code_to_upload), "app.py", repo_id, hf_profile, hf_token)
+                     history = add_bot_message(history, "✅ Uploaded `app.py`. Click 'Send' to generate requirements.")
+                     state = STATE_GENERATING_REQUIREMENTS # Transition state
+                     generated_code = None # Clear the stored code after use to free memory/state space
+                     # Yield updated state variables and history
+                     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                            attempts, app_desc, repo_name, generated_code)
+                     # No return needed
+
+                 except Exception as e:
+                     history = add_bot_message(history, f"❌ Error uploading app.py: {e}. Click 'reset'.")
+                     # Yield error message and reset state on failure
+                     yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                           None, None, None)
+                     # No return needed
 
 
         elif state == STATE_GENERATING_REQUIREMENTS:
-            history = add_bot_message(history, "📄 Generating `requirements.txt`...")
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
-            reqs_list = ["gradio"] if space_sdk == "gradio" else ["streamlit"]
-            if "google.generativeai" in str(app_desc).lower() or "gemini" in str(app_desc).lower() or gemini_api_key:
-                 reqs_list.append("google-generativeai")
-            if "requests" in str(app_desc).lower():
-                 reqs_list.append("requests")
-            reqs_list.append("huggingface_hub")
-            if "image" in str(app_desc).lower() or "upload" in str(app_desc).lower() or "blur" in str(app_desc).lower() or "vision" in str(app_desc).lower():
-                reqs_list.append("Pillow")
-            if "numpy" in str(app_desc).lower(): reqs_list.append("numpy")
-            if "pandas" in str(app_desc).lower(): reqs_list.append("pandas")
-            if "scikit-image" in str(app_desc).lower() or "skimage" in str(app_desc).lower() or "cv2" in str(app_desc).lower() or "opencv-python" in str(app_desc).lower():
-                 reqs_list.append("scikit-image")
-                 reqs_list.append("opencv-python")
-
-            reqs_list = list(dict.fromkeys(reqs_list))
-
-            reqs_content = "\n".join(reqs_list) + "\n"
-
-            history = add_bot_message(history, "✅ `requirements.txt` generated. Click 'Send' to upload.")
-            state = STATE_UPLOADING_REQUIREMENTS
-            generated_code = reqs_content
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+             history = add_bot_message(history, "📄 Generating `requirements.txt`...")
+             # Yield to show message before generating requirements
+             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                    attempts, app_desc, repo_name, generated_code)
+
+             # Logic to determine required packages based on SDK and keywords in the app description
+             reqs_list = ["gradio"] if space_sdk == "gradio" else ["streamlit"]
+             if app_desc: # Check app_desc for keywords only if it's not None
+                 app_desc_lower = app_desc.lower()
+                 if "google.generativeai" in app_desc_lower or "gemini" in app_desc_lower or gemini_api_key:
+                      reqs_list.append("google-generativeai")
+                 if "requests" in app_desc_lower:
+                      reqs_list.append("requests")
+                 # Add common libraries if description suggests they might be needed
+                 if "image" in app_desc_lower or "upload" in app_desc_lower or "blur" in app_desc_lower or "vision" in app_desc_lower:
+                     reqs_list.append("Pillow") # Pillow is a common image processing library
+                 if "numpy" in app_desc_lower: reqs_list.append("numpy")
+                 if "pandas" in app_desc_lower: reqs_list.append("pandas")
+                 # Add scikit-image and opencv if image processing is heavily implied
+                 if any(lib in app_desc_lower for lib in ["scikit-image", "skimage", "cv2", "opencv"]):
+                      reqs_list.extend(["scikit-image", "opencv-python"]) # Note: opencv-python for pip
+
+             # Add essential libraries regardless of description keywords
+             reqs_list.append("huggingface_hub") # Needed for interaction helpers if used in app
+
+             # Use dict.fromkeys to get unique items while preserving insertion order (Python 3.7+)
+             reqs_list = list(dict.fromkeys(reqs_list))
+
+             reqs_content = "\n".join(reqs_list) + "\n"
+
+             history = add_bot_message(history, "✅ `requirements.txt` generated. Click 'Send' to upload.")
+             state = STATE_UPLOADING_REQUIREMENTS # Transition state
+             generated_code = reqs_content # Store requirements content
+             # Yield updated state variables and history
+             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                    attempts, app_desc, repo_name, generated_code)
+             # No return needed
 
 
         elif state == STATE_UPLOADING_REQUIREMENTS:
+            # Retrieve requirements content from state variable
             reqs_content_to_upload = generated_code
             if not reqs_content_to_upload:
                  history = add_bot_message(history, "Internal error: No requirements content to upload. Resetting.")
-                 yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
-                 return
-
-            history = add_bot_message(history, "☁️ Uploading `requirements.txt`...")
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
-            try:
-                upload_file_to_space_action(io.StringIO(reqs_content_to_upload), "requirements.txt", repo_id, hf_profile, hf_token)
-                history = add_bot_message(history, "✅ Uploaded `requirements.txt`. Click 'Send' to generate README.")
-                state = STATE_GENERATING_README
-                generated_code = None
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                 yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                        None, None, None)
+                 # No return needed
 
-
-            except Exception as e:
-                history = add_bot_message(history, f"❌ Error uploading requirements.txt: {e}. Click 'reset'.")
-                yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
+            else:
+                 history = add_bot_message(history, "☁️ Uploading `requirements.txt`...")
+                 # Yield to show message before upload
+                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                        attempts, app_desc, repo_name, generated_code)
+
+                 try:
+                     # Perform requirements file upload
+                     upload_file_to_space_action(io.StringIO(reqs_content_to_upload), "requirements.txt", repo_id, hf_profile, hf_token)
+                     history = add_bot_message(history, "✅ Uploaded `requirements.txt`. Click 'Send' to generate README.")
+                     state = STATE_GENERATING_README # Transition state
+                     generated_code = None # Clear content after use
+                     # Yield updated state variables and history
+                     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                            attempts, app_desc, repo_name, generated_code)
+                     # No return needed
+
+                 except Exception as e:
+                     history = add_bot_message(history, f"❌ Error uploading requirements.txt: {e}. Click 'reset'.")
+                     # Yield error message and reset state on failure
+                     yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                           None, None, None)
+                     # No return needed
 
         elif state == STATE_GENERATING_README:
             history = add_bot_message(history, "📝 Generating `README.md`...")
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+            # Yield message before generating README
+            yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                   attempts, app_desc, repo_name, generated_code)
 
+            # Generate simple README content with Space metadata header
             readme_title = repo_name if repo_name else "My Awesome Space"
             readme_description = app_desc if app_desc else f"This Hugging Face Space hosts an AI-generated {space_sdk} application."
 
-            readme_content = f"# {readme_title}\n\n{readme_description}\n\n" \
-                             "This Space was automatically generated by an AI workflow.\n\n" \
-                             f"Built with the {space_sdk} SDK.\n"
+            readme_content = f"""---
+title: {readme_title}
+emoji: 🚀
+colorFrom: blue
+colorTo: yellow
+sdk: {space_sdk}
+app_file: app.py
+pinned: false
+---
+
+# {readme_title}
+
+{readme_description}
+
+This Space was automatically generated by an AI workflow using Google Gemini and Gradio.
+""" # Added Space metadata header and slightly improved content
 
             history = add_bot_message(history, "✅ `README.md` generated. Click 'Send' to upload.")
-            state = STATE_UPLOADING_README
-            generated_code = readme_content
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+            state = STATE_UPLOADING_README # Transition state
+            generated_code = readme_content # Store README content
+            # Yield updated state variables and history
+            yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                   attempts, app_desc, repo_name, generated_code)
+            # No return needed
 
 
         elif state == STATE_UPLOADING_README:
+            # Retrieve README content from state variable
             readme_content_to_upload = generated_code
             if not readme_content_to_upload:
                  history = add_bot_message(history, "Internal error: No README content to upload. Resetting.")
-                 yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
-                 return
-
-            history = add_bot_message(history, "☁️ Uploading `README.md`...")
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
-            try:
-                upload_file_to_space_action(io.StringIO(readme_content_to_upload), "README.md", repo_id, hf_profile, hf_token)
-                history = add_bot_message(history, "✅ Uploaded `README.md`. All files uploaded. Space is now building. Click 'Send' to check build logs.")
-                state = STATE_CHECKING_LOGS_BUILD
-                generated_code = None
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
-
-            except Exception as e:
-                history = add_bot_message(history, f"❌ Error uploading README.md: {e}. Click 'reset'.")
-                yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
+                 yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                        None, None, None)
+                 # No return needed
 
+            else:
+                 history = add_bot_message(history, "☁️ Uploading `README.md`...")
+                 # Yield message before upload
+                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                        attempts, app_desc, repo_name, generated_code)
+
+                 try:
+                     # Perform README file upload
+                     upload_file_to_space_action(io.StringIO(readme_content_to_upload), "README.md", repo_id, hf_profile, hf_token)
+                     history = add_bot_message(history, "✅ Uploaded `README.md`. All files uploaded. Space is now building. Click 'Send' to check build logs.")
+                     state = STATE_CHECKING_LOGS_BUILD # Transition to checking build logs
+                     generated_code = None # Clear content after use
+                     # Yield updated state variables and history
+                     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                            attempts, app_desc, repo_name, generated_code)
+                     # No return needed
+
+                 except Exception as e:
+                     history = add_bot_message(history, f"❌ Error uploading README.md: {e}. Click 'reset'.")
+                     # Yield error message and reset state on failure
+                     yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                           None, None, None)
+                     # No return needed
 
         elif state == STATE_CHECKING_LOGS_BUILD:
              history = add_bot_message(history, "🔍 Fetching build logs...")
-             yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+             # Yield message before fetching logs (which includes a delay)
+             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                    attempts, app_desc, repo_name, generated_code)
 
+             # Fetch build logs from HF Space
              build_logs_text = get_build_logs_action(repo_id, hf_profile, hf_token)
-             updated_build = build_logs_text
+             updated_build = build_logs_text # Update the logs display variable
 
-             if "Error" in updated_build or "Exception" in updated_build or "Build failed" in updated_build:
-                  history = add_bot_message(history, "⚠️ Build logs indicate potential issues. Please inspect above. Click 'Send' to check container logs (app might still start).")
-                  state = STATE_CHECKING_LOGS_RUN
-                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+             # Simple check for common error indicators in logs (case-insensitive)
+             if "error" in updated_build.lower() or "exception" in updated_build.lower() or "build failed" in updated_build.lower():
+                  history = add_bot_message(history, "⚠️ Build logs indicate potential issues. Please inspect above. Click 'Send' to check container logs (app might still start despite build warnings).")
+                  state = STATE_CHECKING_LOGS_RUN # Transition even on build error, to see if container starts
+                  # Yield updated state, logs, and variables
+                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                         attempts, app_desc, repo_name, generated_code)
+                  # No return needed
 
              else:
                   history = add_bot_message(history, "✅ Build logs fetched. Click 'Send' to check container logs.")
-                  state = STATE_CHECKING_LOGS_RUN
-                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                  state = STATE_CHECKING_LOGS_RUN # Transition to next log check
+                  # Yield updated state, logs, and variables
+                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                         attempts, app_desc, repo_name, generated_code)
+                  # No return needed
 
 
         elif state == STATE_CHECKING_LOGS_RUN:
              history = add_bot_message(history, "🔍 Fetching container logs...")
-             yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+             # Yield message before fetching logs (includes a delay)
+             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                    attempts, app_desc, repo_name, generated_code)
 
+             # Fetch container logs from HF Space
              container_logs_text = get_container_logs_action(repo_id, hf_profile, hf_token)
-             updated_run = container_logs_text
+             updated_run = container_logs_text # Update the logs display variable
 
-             if ("Error" in updated_run or "Exception" in updated_run) and attempts < MAX_DEBUG_ATTEMPTS:
-                  attempts += 1
+             # Check for errors in run logs and if we have debug attempts left
+             if ("error" in updated_run.lower() or "exception" in updated_run.lower()) and attempts < MAX_DEBUG_ATTEMPTS:
+                  attempts += 1 # Increment debug attempts counter
                   history = add_bot_message(history, f"❌ Errors detected in container logs. Attempting debug fix #{attempts}/{MAX_DEBUG_ATTEMPTS}. Click 'Send' to proceed.")
-                  state = STATE_DEBUGGING_CODE
-                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
-             elif ("Error" in updated_run or "Exception" in updated_run) and attempts >= MAX_DEBUG_ATTEMPTS:
+                  state = STATE_DEBUGGING_CODE # Transition to the debugging state
+                  # Yield updated state, logs, attempts, and variables
+                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                         attempts, app_desc, repo_name, generated_code)
+                  # No return needed
+
+             elif ("error" in updated_run.lower() or "exception" in updated_run.lower()) and attempts >= MAX_DEBUG_ATTEMPTS:
+                  # Max debug attempts reached
                   history = add_bot_message(history, f"❌ Errors detected in container logs. Max debug attempts ({MAX_DEBUG_ATTEMPTS}) reached. Please inspect logs manually or click 'reset'.")
-                  state = STATE_COMPLETE
-                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                  state = STATE_COMPLETE # Workflow ends on failure after attempts
+                  # Yield updated state, logs, attempts, and variables
+                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                         attempts, app_desc, repo_name, generated_code)
+                  # No return needed
 
              else:
+                  # No significant errors found in logs, assume success
                   history = add_bot_message(history, "✅ App appears to be running successfully! Check the iframe above. Click 'reset' to start a new project.")
-                  state = STATE_COMPLETE
-                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                  state = STATE_COMPLETE # Workflow ends on success
+                  # Yield updated state, logs, attempts, and variables
+                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                         attempts, app_desc, repo_name, generated_code)
+                  # No return needed
 
 
         elif state == STATE_DEBUGGING_CODE:
              history = add_bot_message(history, f"🧠 Calling Gemini to generate fix based on logs...")
-             yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+             # Yield message before Gemini API call
+             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                    attempts, app_desc, repo_name, generated_code)
 
+             # Construct prompt for Gemini including the container logs
              debug_prompt = f"""
 You are debugging a {space_sdk} Space. The goal is to fix the code in `app.py` based on the container logs provided.
 
 Here are the container logs:
+Use code with caution.
+Python
 {updated_run}
-
 Generate the *complete, fixed* content for `app.py` based on these logs.
 Return **only** the python code block for app.py. Do not include any extra text, explanations, or markdown outside the code block.
 """
              try:
+                  # Call Gemini to generate the corrected code
                   fix_code = call_gemini(debug_prompt, gemini_api_key, gemini_model)
                   fix_code = fix_code.strip()
+                  # Clean up potential markdown formatting
                   if fix_code.startswith("```python"):
                       fix_code = fix_code[len("```python"):].strip()
                   if fix_code.endswith("```"):
@@ -513,94 +697,125 @@ Return **only** the python code block for app.py. Do not include any extra text,
                      raise ValueError("Gemini returned empty fix code.")
 
                   history = add_bot_message(history, "✅ Fix code generated. Click 'Send' to upload.")
-                  state = STATE_UPLOADING_FIXED_APP_PY
-                  generated_code = fix_code
-                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
+                  state = STATE_UPLOADING_FIXED_APP_PY # Transition to the upload state for the fix
+                  generated_code = fix_code # Store the generated fix code
+                  # Yield updated state, code, and variables
+                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                         attempts, app_desc, repo_name, generated_code)
+                  # No return needed
 
              except Exception as e:
                   history = add_bot_message(history, f"❌ Error generating debug code: {e}. Click 'reset'.")
-                  yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
+                  # Yield error message and reset state on failure
+                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                         None, None, None)
+                  # No return needed
 
         elif state == STATE_UPLOADING_FIXED_APP_PY:
+             # Retrieve the fixed code from the state variable
              fixed_code_to_upload = generated_code
              if not fixed_code_to_upload:
                   history = add_bot_message(history, "Internal error: No fixed code available to upload. Resetting.")
-                  yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
-                 return
-
-             history = add_bot_message(history, "☁️ Uploading fixed `app.py`...")
-             yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
-
-             try:
-                  upload_file_to_space_action(io.StringIO(fixed_code_to_upload), "app.py", repo_id, hf_profile, hf_token)
-                  history = add_bot_message(history, "✅ Fixed `app.py` uploaded. Space will rebuild. Click 'Send' to check logs again.")
-                  state = STATE_CHECKING_LOGS_RUN
-                  generated_code = None
-                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                        None, None, None)
+                  # No return needed
 
-
-             except Exception as e:
-                  history = add_bot_message(history, f"❌ Error uploading fixed app.py: {e}. Click 'reset'.")
-                  yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
+             else:
+                  history = add_bot_message(history, "☁️ Uploading fixed `app.py`...")
+                  # Yield message before upload
+                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                         attempts, app_desc, repo_name, generated_code)
+
+                  try:
+                      # Perform the upload of the fixed app.py
+                      upload_file_to_space_action(io.StringIO(fixed_code_to_upload), "app.py", repo_id, hf_profile, hf_token)
+                      history = add_bot_message(history, "✅ Fixed `app.py` uploaded. Space will rebuild. Click 'Send' to check logs again.")
+                      state = STATE_CHECKING_LOGS_RUN # Go back to checking run logs after uploading the fix
+                      generated_code = None # Clear code after use
+                      # Yield updated state, code, and variables
+                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                             attempts, app_desc, repo_name, generated_code)
+                      # No return needed
+
+                  except Exception as e:
+                      history = add_bot_message(history, f"❌ Error uploading fixed app.py: {e}. Click 'reset'.")
+                      # Yield error message and reset state on failure
+                      yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+                            None, None, None)
+                      # No return needed
 
         elif state == STATE_COMPLETE:
-             pass # No further action needed in this state
+             # If in the complete state, the workflow is finished for this project.
+             # Subsequent clicks just add user messages; we simply yield the current state.
+             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
+                    attempts, app_desc, repo_name, generated_code)
+             # No return needed
 
 
     except Exception as e:
+        # This catches any unexpected errors that occur within any state's logic
         error_message = f"Workflow step failed unexpectedly ({state}): {e}. Click 'Send' to re-attempt this step or 'reset'."
         history = add_bot_message(history, error_message)
-        print(f"Critical Error in state {state}: {e}")
-        yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
-        # Removed the 'return' here
+        print(f"Critical Error in state {state}: {e}") # Log the error for debugging purposes
+        # Yield an error state and reset essential workflow variables on critical failure
+        yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
+               None, None, None)
+        # No return needed after yield
 
 
 # --- Build the Gradio UI ---
 
 with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
-    # State variables
+    # Gradio State variables - these persist their values across user interactions (clicks)
     hf_profile   = gr.State(None)
     hf_token     = gr.State(None)
     gemini_key   = gr.State(None)
-    gemini_model = gr.State("gemini-1.5-flash")
-    repo_id      = gr.State(None)
-    workflow     = gr.State(STATE_IDLE)
-    sdk_state    = gr.State("gradio")
-    debug_attempts = gr.State(0)
-    app_description = gr.State(None)
-    repo_name_state = gr.State(None)
-    generated_code_state = gr.State(None)
+    gemini_model = gr.State("gemini-1.5-flash") # Default selected model
+    repo_id      = gr.State(None) # Stores the ID of the created Space
+    workflow     = gr.State(STATE_IDLE) # Stores the current state of the AI workflow
+    sdk_state    = gr.State("gradio") # Stores the selected Space SDK (Gradio or Streamlit)
+    debug_attempts = gr.State(0) # Counter for how many debugging attempts have been made
+    app_description = gr.State(None) # Stores the user's initial description of the desired app
+    repo_name_state = gr.State(None) # Stores the chosen repository name for the Space
+    generated_code_state = gr.State(None) # Temporary storage for generated file content (app.py, reqs, README)
 
     with gr.Row():
-        # Sidebar
+        # Sidebar column for inputs and status displays
         with gr.Column(scale=1, min_width=300):
             gr.Markdown("## Hugging Face Login")
             login_status = gr.Markdown("*Not logged in.*")
+            # Hugging Face Login Button
             login_btn    = gr.LoginButton(variant="huggingface")
 
-            # Initial load to check login status
+            # Initial load event to check login status (if cached)
             ai_builder_tab.load(show_profile, outputs=login_status)
-            # Update status on login click
+            # Update status display when login button reports success
             login_btn.click(show_profile, outputs=login_status)
 
             gr.Markdown("## Google AI Studio API Key")
-            gemini_input  = gr.Textbox(label="API Key", type="password", interactive=True)
-            gemini_status = gr.Markdown("")
+            # Textbox for Gemini API key. Read from environment variable if available.
+            gemini_input  = gr.Textbox(
+                label="API Key",
+                type="password", # Hides input for security
+                interactive=True,
+                value=os.environ.get("GOOGLE_API_KEY") # Pre-fill if GOOGLE_API_KEY env var is set
+            )
+            gemini_status = gr.Markdown("") # Display Gemini configuration status
 
             gr.Markdown("## Gemini Model")
+            # Radio buttons to select the Gemini model
             model_selector = gr.Radio(
                 choices=[
                     ("Gemini 1.5 Flash", "gemini-1.5-flash"),
                     ("Gemini 1.5 Pro", "gemini-1.5-pro"),
                     ("Gemini 1.0 Pro", "gemini-1.0-pro"),
                 ],
-                value="gemini-1.5-flash",
+                value="gemini-1.5-flash", # Default selection
                 label="Select model",
                 interactive=True
             )
 
-            # Configure Gemini status on load
+            # Configure Gemini status on initial load (if API key env var is set)
             ai_builder_tab.load(
                 configure_gemini,
                 inputs=[gemini_key, gemini_model],
@@ -608,145 +823,165 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
             )
 
             gr.Markdown("## Space SDK")
+            # Radio buttons to select the Space SDK (Gradio or Streamlit)
             sdk_selector = gr.Radio(choices=["gradio","streamlit"], value="gradio", label="Template SDK", interactive=True)
+            # Update the sdk_state state variable when the selection changes
             sdk_selector.change(lambda s: s, inputs=sdk_selector, outputs=sdk_state)
 
             gr.Markdown("## Workflow Status")
+            # Textboxes to display the current workflow state and Space ID
             status_text = gr.Textbox(label="Current State", value=STATE_IDLE, interactive=False)
             repo_id_text = gr.Textbox(label="Current Space ID", value="None", interactive=False)
 
-            # --- Debugging Indicator ---
-            prereq_status = gr.Markdown("Checking...")
+            # --- Prerequisite Status Indicator ---
+            # Markdown to show if prerequisites (HF login, Gemini key) are met
+            prereq_status = gr.Markdown("Checking prerequisites...")
 
 
-        # Main content
+        # Main content area column
         with gr.Column(scale=3):
+            # Chatbot to display the conversation and workflow messages
             chatbot    = gr.Chatbot(type='messages', label="AI Workflow Chat")
+            # Textbox for user input messages
             user_input = gr.Textbox(placeholder="Type your message…", interactive=True)
-            send_btn   = gr.Button("Send", interactive=False)
-
-            # Logic to enable send button only when logged in and API key is set
-            def update_send_button_state(profile: gr.OAuthProfile | None, token: gr.OAuthToken | None, key: str | None, model: str | None, *args, **kwargs):
-                """
-                Gradio sometimes passes extra positional/keyword args (old_value, event, etc.),
-                so we absorb them with *args/**kwargs.
-                Returns Gradio.update for the button's interactive state and a status string.
-                """
+            # Button to send the user message and trigger the workflow step
+            send_btn   = gr.Button("Send", interactive=False) # Starts disabled until prereqs are met
+
+            # Helper function to control send button interactivity and prerequisite status text
+            # This function is triggered by changes in login status and Gemini configuration
+            def update_send_button_state(
+                profile: gr.OAuthProfile | None,
+                token: gr.OAuthToken | None,
+                key: str | None,
+                model: str | None,
+                # Absorb potential extra args Gradio passes to event handlers
+                *args,
+                **kwargs
+            ):
+                """Determines if the send button should be active and updates status text."""
                 is_logged_in    = profile is not None and token is not None
-                is_gemini_ready = key     is not None and model is not None
-
-                # Print states for debugging
-                print(f"update_send_button_state called:")
-                print(f"  Profile: {profile.username if profile else None}")
-                print(f"  Token: {'Set' if token else 'None'}")
-                print(f"  API Key: {'Set' if key else 'None'}")
-                print(f"  Model: {model}")
-                print(f"  Logged in: {is_logged_in}, Gemini Ready: {is_gemini_ready}")
-
-                status_str = ""
-                if is_logged_in and is_gemini_ready:
+                is_gemini_ready = key     is not None and model is not None # Check if key and model are set
+
+                status_parts = []
+                if not is_logged_in:
+                    status_parts.append("⚠️ Not logged in to Hugging Face.")
+                if not key:
+                    status_parts.append("⚠️ Gemini API key not set.")
+                if not model:
+                     status_parts.append("⚠️ Gemini model not selected.")
+
+                is_ready = is_logged_in and is_gemini_ready
+
+                if is_ready:
                     status_str = "✅ Ready to send commands."
                 else:
-                     status_parts = []
-                     if not is_logged_in:
-                         status_parts.append("⚠️ Not logged in to Hugging Face.")
-                     if not key:
-                         status_parts.append("⚠️ Gemini API key not set.")
-                     if not model:
-                          status_parts.append("⚠️ Gemini model not selected.")
                      status_str = " ".join(status_parts)
-                     if not status_str:
-                          status_str = "✅ Ready to send commands."
+                     if not status_str: # Fallback, should not be needed if not is_ready
+                          status_str = "Checking prerequisites..."
 
 
-                return gr.update(interactive=is_logged_in and is_gemini_ready), status_str
+                # gr.update is used to dynamically change a component's properties
+                return gr.update(interactive=is_ready), status_str
 
-            # --- Implement Chained Events ---
+            # --- Implement Chained Events for Prerequisites ---
+            # Gradio's `.then()` allows chaining events: Action A happens, then Action B happens.
 
-            # 1. Login Button: Update profile/token state, THEN update send button state
+            # 1. Login Button: When clicked and successful, update profile/token state,
+            # THEN update send button state based on all prereqs.
             login_btn.click(
-                 lambda x: (x[0], x[1]), # Extract profile and token from LoginButton output
+                 # The LoginButton outputs a tuple (OAuthProfile, OAuthToken) on success
+                 lambda x: (x[0], x[1]),
                  inputs=[login_btn],
-                 outputs=[hf_profile, hf_token]
+                 outputs=[hf_profile, hf_token] # Update these State variables
             ).then( # Chain the next action after state is updated
                  update_send_button_state,
                  inputs=[hf_profile, hf_token, gemini_key, gemini_model],
-                 outputs=[send_btn, prereq_status]
+                 outputs=[send_btn, prereq_status] # Update button interactivity and status text
             )
 
-
-            # 2. Gemini Key Input: Update key state, THEN configure Gemini, THEN update send button state
+            # 2. Gemini Key Input: When text changes, update key state,
+            # THEN configure Gemini status, THEN update send button state.
+            # The Textbox 'change' event passes the new value as its input
             gemini_input.change(
-                lambda k: k, # Update gemini_key state
+                lambda k: k, # Simple function to pass the new value to the state variable
                 inputs=[gemini_input],
-                outputs=[gemini_key]
-            ).then( # Chain configure_gemini after key is updated
+                outputs=[gemini_key] # Update gemini_key state variable
+            ).then( # Chain configure_gemini after key state is updated
                 configure_gemini,
                 inputs=[gemini_key, gemini_model],
-                outputs=[gemini_status]
+                outputs=[gemini_status] # Update Gemini status text
             ).then( # Chain update_send_button_state after config status is updated
                 update_send_button_state,
                 inputs=[hf_profile, hf_token, gemini_key, gemini_model],
-                outputs=[send_btn, prereq_status]
+                outputs=[send_btn, prereq_status] # Update button interactivity and status text
             )
 
-
-            # 3. Gemini Model Selector: Update model state, THEN configure Gemini, THEN update send button state
+            # 3. Gemini Model Selector: When selection changes, update model state,
+            # THEN configure Gemini status, THEN update send button state.
+            # The Radio 'change' event passes the new value as its input
             model_selector.change(
-                lambda m: m, # Update gemini_model state
+                lambda m: m, # Simple function to pass the new value to the state variable
                 inputs=[model_selector],
-                outputs=[gemini_model]
-            ).then( # Chain configure_gemini after model is updated
+                outputs=[gemini_model] # Update gemini_model state variable
+            ).then( # Chain configure_gemini after model state is updated
                 configure_gemini,
                 inputs=[gemini_key, gemini_model],
-                outputs=[gemini_status]
+                outputs=[gemini_status] # Update Gemini status text
             ).then( # Chain update_send_button_state after config status is updated
                 update_send_button_state,
                 inputs=[hf_profile, hf_token, gemini_key, gemini_model],
-                outputs=[send_btn, prereq_status]
+                outputs=[send_btn, prereq_status] # Update button interactivity and status text
             )
 
-            # 4. Initial Load: Update send button state based on initial (or cached) state
+            # 4. Initial Load: On page load, check prereqs and update send button/status.
+            # This accounts for cached logins or environment variables set before launch.
             ai_builder_tab.load(
                 update_send_button_state,
                 inputs=[hf_profile, hf_token, gemini_key, gemini_model],
-                outputs=[send_btn, prereq_status]
+                outputs=[send_btn, prereq_status] # Update button interactivity and status text
             )
 
-
-            iframe    = gr.HTML("<p>No Space created yet.</p>")
-            build_txt = gr.Textbox(label="Build Logs", lines=10, interactive=False, value="")
-            run_txt   = gr.Textbox(label="Container Logs", lines=10, interactive=False, value="")
+            # UI elements to display the Space preview iframe and build/run logs
+            iframe    = gr.HTML("<p>No Space created yet.</p>") # HTML element for the Space iframe
+            # Textboxes for logs, interactive=False means user can't type here
+            build_txt = gr.Textbox(label="Build Logs", lines=10, interactive=False, value="", max_lines=20) # Set max_lines for scrollability
+            run_txt   = gr.Textbox(label="Container Logs", lines=10, interactive=False, value="", max_lines=20) # Set max_lines for scrollability
 
             # The main event handler for the Send button
+            # This .click() event triggers the ai_workflow_chat generator function
             send_btn.click(
-                ai_workflow_chat,
+                ai_workflow_chat, # The generator function to run
+                # Inputs are read from UI components and State variables
                 inputs=[
-                    user_input, chatbot,
-                    hf_profile, hf_token,
-                    gemini_key, gemini_model,
-                    repo_id, workflow, sdk_state,
-                    iframe, run_txt, build_txt,
-                    debug_attempts, app_description, repo_name_state, generated_code_state
+                    user_input, chatbot, # UI inputs (message, current chat history)
+                    hf_profile, hf_token, # HF State variables
+                    gemini_key, gemini_model, # Gemini State variables
+                    repo_id, workflow, sdk_state, # Workflow State variables
+                    iframe, run_txt, build_txt, # UI outputs whose current values are needed by the generator
+                    debug_attempts, app_description, repo_name_state, generated_code_state # Other State variables
                 ],
+                # Outputs are updated by the values yielded from the generator
                 outputs=[
-                    chatbot,
-                    repo_id, workflow,
-                    iframe, run_txt, build_txt,
-                    debug_attempts, app_description, repo_name_state, generated_code_state
+                    chatbot, # Update Chatbot with new messages
+                    repo_id, workflow, # Update workflow State variables
+                    iframe, run_txt, build_txt, # Update UI outputs
+                    debug_attempts, app_description, repo_name_state, generated_code_state # Update other State variables
                 ]
-            ).success(
+            ).success( # Chain a .success() event to run *after* the .click() handler completes without error
+                 # Clear the user input textbox after the message is sent and processed
                  lambda: gr.update(value=""),
                  inputs=None,
-                 outputs=user_input
+                 outputs=user_input # Update the user input textbox
             )
 
-            # Link state variables to UI status displays (reactive updates)
+            # Link State variables' changes to UI status displays (reactive updates)
+            # When the 'workflow' state variable changes, update the text in status_text
             workflow.change(lambda s: s, inputs=workflow, outputs=status_text)
+            # When the 'repo_id' state variable changes, update the text in repo_id_text
             repo_id.change(lambda r: r if r else "None", inputs=repo_id, outputs=repo_id_text)
 
 
-    # Add an initial message to the chatbot on load
+    # Add an initial welcome message to the chatbot when the UI loads
     def greet():
         return [{"role": "assistant", "content": "Welcome! Please log in to Hugging Face and provide your Google AI Studio API key to start building Spaces. Once ready, type 'generate me a gradio app called myapp' or 'create' to begin."}]
 
@@ -754,12 +989,21 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
 
 
 if __name__ == "__main__":
-    # Optional: Configure retries for huggingface_hub requests
+    # Optional: Configure retries for huggingface_hub requests to make them more robust
     # from requests.adapters import HTTPAdapter
     # from urllib3.util.retry import Retry
-    # retry_strategy = Retry(total=5, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504])
+    # retry_strategy = Retry(total=5, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504]) # Define retry strategy for specific HTTP codes
     # adapter = HTTPAdapter(max_retries=retry_strategy)
-    # get_session().mount("http://", adapter)
-    # get_session().mount("https://", adapter)
-
-    ai_builder_tab.launch()
+    # session = get_session() # Get the session object used internally by huggingface_hub
+    # session.mount("http://", adapter)
+    # session.mount("https://", adapter)
+
+    # Optional: Configure Gradio settings using environment variables
+    # Set max upload size (e.g., 100MB) for files like app.py
+    os.environ["GRADIO_MAX_FILE_SIZE"] = "100MB"
+    # Optional: Set a local temporary directory for Gradio uploads
+    os.environ["GRADIO_TEMP_DIR"] = "./tmp"
+    os.makedirs(os.environ["GRADIO_TEMP_DIR"], exist_ok=True) # Ensure the directory exists
+
+    # Launch the Gradio UI
+    ai_builder_tab.launch()