Update app.py

app.py

@@ -5,15 +5,12 @@ import io
 import json
 import os
 from pathlib import Path
+import time # Added for simulating mock delay
 
 # --- Configuration ---
 GEMINI_MODEL_NAME = "gemini-2.5-pro-preview-03-25"
-# Maximum estimated tokens to try fitting into a single prompt
-# Adjust based on typical file sizes and Gemini limits/performance
-# 1M tokens is roughly 4MB-5MB of text, but structure matters. Start lower.
-MAX_PROMPT_TOKENS_ESTIMATE = 800000 # Be conservative initially
+MAX_PROMPT_TOKENS_ESTIMATE = 800000 # Adjust as needed
 
-# Define the types of analysis available
 AVAILABLE_ANALYSES = {
     "generate_docs": "Generate Missing Docstrings/Comments",
     "find_bugs": "Identify Potential Bugs & Anti-patterns",
@@ -22,26 +19,45 @@ AVAILABLE_ANALYSES = {
     "suggest_refactoring": "Suggest Refactoring Opportunities"
 }
 
-# Define common code file extensions to include
 CODE_EXTENSIONS = {'.py', '.js', '.java', '.c', '.cpp', '.h', '.cs', '.go', '.rb', '.php', '.swift', '.kt', '.ts', '.html', '.css', '.scss', '.sql'}
 
+# --- Session State Initialization ---
+# Initialize session state for mock mode toggle if it doesn't exist
+if 'mock_api_call' not in st.session_state:
+    st.session_state.mock_api_call = False # Default to using the real API
+
 # --- Gemini API Setup ---
-try:
-    if 'GEMINI_API_KEY' not in st.secrets:
-        st.error("🚨 Gemini API Key not found. Add it to `.streamlit/secrets.toml`.")
-        st.stop()
-    genai.configure(api_key=st.secrets["GEMINI_API_KEY"])
-    model = genai.GenerativeModel(GEMINI_MODEL_NAME)
-    print("Gemini Model Initialized.")
-except Exception as e:
-    st.error(f"🚨 Error initializing Gemini SDK: {e}")
-    st.stop()
+# Defer full initialization until needed if mock mode might be used first
+model = None
+def initialize_gemini_model():
+    global model
+    if model is None and not st.session_state.mock_api_call:
+        try:
+            if 'GEMINI_API_KEY' not in st.secrets:
+                st.error("🚨 Gemini API Key not found. Add it to `.streamlit/secrets.toml`.")
+                st.stop()
+            genai.configure(api_key=st.secrets["GEMINI_API_KEY"])
+            model = genai.GenerativeModel(GEMINI_MODEL_NAME)
+            print("Gemini Model Initialized.")
+            return True
+        except Exception as e:
+            st.error(f"🚨 Error initializing Gemini SDK: {e}")
+            st.stop()
+            return False
+    elif st.session_state.mock_api_call:
+        print("Running in Mock Mode. Skipping Gemini initialization.")
+        return True # Allow proceeding in mock mode
+    elif model is not None:
+         print("Gemini Model already initialized.")
+         return True
+    return False
+
 
 # --- Helper Functions ---
 
 def estimate_token_count(text):
-    """Roughly estimate token count (4 chars per token is a common rule of thumb)."""
-    return len(text) // 3 # Be generous here
+    """Roughly estimate token count (3-4 chars per token)."""
+    return len(text) // 3
 
 def process_zip_file(uploaded_file):
     """Extracts code files and their content from the uploaded zip file."""
@@ -53,16 +69,13 @@ def process_zip_file(uploaded_file):
     try:
         with zipfile.ZipFile(io.BytesIO(uploaded_file.getvalue()), 'r') as zip_ref:
             for member in zip_ref.infolist():
-                # Skip directories and files in hidden folders like .git, __pycache__
-                if member.is_dir() or member.filename.startswith('.') or '__' in member.filename:
+                if member.is_dir() or any(part.startswith('.') for part in Path(member.filename).parts) or '__' in member.filename:
                     continue
 
                 file_path = Path(member.filename)
-                # Check if the file extension is in our allowed list
                 if file_path.suffix.lower() in CODE_EXTENSIONS:
                     try:
                         with zip_ref.open(member) as file:
-                            # Decode defensively, try common encodings
                             try:
                                 content = file.read().decode('utf-8')
                             except UnicodeDecodeError:
@@ -70,7 +83,7 @@ def process_zip_file(uploaded_file):
                                     content = file.read().decode('latin-1')
                                 except Exception as decode_err:
                                      ignored_files.append(f"{member.filename} (Decode Error: {decode_err})")
-                                     continue # Skip if undecodable
+                                     continue
 
                             code_files[member.filename] = content
                             total_chars += len(content)
@@ -78,7 +91,9 @@ def process_zip_file(uploaded_file):
                     except Exception as read_err:
                          ignored_files.append(f"{member.filename} (Read Error: {read_err})")
                 else:
-                    ignored_files.append(f"{member.filename} (Skipped Extension: {file_path.suffix})")
+                    # Only add to ignored if it's not explicitly ignored by path rules above
+                     if not (any(part.startswith('.') for part in Path(member.filename).parts) or '__' in member.filename):
+                        ignored_files.append(f"{member.filename} (Skipped Extension: {file_path.suffix})")
 
     except zipfile.BadZipFile:
         st.error("🚨 Invalid or corrupted ZIP file.")
@@ -93,16 +108,14 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
     """Constructs the prompt for Gemini, including code content and JSON structure request."""
     prompt_content = "Analyze the following codebase provided as a collection of file paths and their content.\n\n"
     current_token_estimate = estimate_token_count(prompt_content)
-
-    # Concatenate file content with markers
     included_files = []
     concatenated_code = ""
+
     for filename, content in code_files_dict.items():
         file_marker = f"--- START FILE: {filename} ---\n"
         file_content = f"{content}\n"
         file_end_marker = f"--- END FILE: {filename} ---\n\n"
         segment = file_marker + file_content + file_end_marker
-
         segment_token_estimate = estimate_token_count(segment)
 
         if current_token_estimate + segment_token_estimate <= MAX_PROMPT_TOKENS_ESTIMATE:
@@ -110,8 +123,8 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
             current_token_estimate += segment_token_estimate
             included_files.append(filename)
         else:
-            st.warning(f"⚠️ Codebase likely exceeds context window estimate ({MAX_PROMPT_TOKENS_ESTIMATE} tokens). Analysis will be performed only on the first {len(included_files)} files ({current_token_estimate} tokens). Consider analyzing smaller parts separately.")
-            break # Stop adding files if limit reached
+            st.warning(f"⚠️ Codebase may exceed context window estimate (~{MAX_PROMPT_TOKENS_ESTIMATE} tokens). Analysis performed only on the first {len(included_files)} files ({current_token_estimate} tokens).")
+            break
 
     if not included_files:
         st.error("🚨 No code files could be included within the estimated token limit.")
@@ -119,25 +132,22 @@ def construct_analysis_prompt(code_files_dict, requested_analyses):
 
     prompt_content += concatenated_code
 
-    # Define the requested JSON structure based on selections
     json_structure_description = "{\n"
+    # Dynamically build the JSON structure based on selection
+    structure_parts = []
     if "generate_docs" in requested_analyses:
-        json_structure_description += '    "documentation_suggestions": [{"file": "path/to/file", "line": number, "suggestion": "Suggested docstring/comment"}],\n'
+        structure_parts.append('    "documentation_suggestions": [{"file": "path/to/file", "line": number, "suggestion": "Suggested docstring/comment"}]')
     if "find_bugs" in requested_analyses:
-        json_structure_description += '    "potential_bugs": [{"file": "path/to/file", "line": number, "description": "Description of potential bug/anti-pattern", "severity": "High/Medium/Low"}],\n'
+        structure_parts.append('    "potential_bugs": [{"file": "path/to/file", "line": number, "description": "Description of potential bug/anti-pattern", "severity": "High/Medium/Low"}]')
     if "check_style" in requested_analyses:
-        json_structure_description += '    "style_issues": [{"file": "path/to/file", "line": number, "description": "Description of style deviation"}],\n'
+        structure_parts.append('    "style_issues": [{"file": "path/to/file", "line": number, "description": "Description of style deviation"}]')
     if "summarize_modules" in requested_analyses:
-        json_structure_description += '    "module_summaries": [{"file": "path/to/file", "summary": "One-paragraph summary of the file purpose/functionality"}],\n'
+        structure_parts.append('    "module_summaries": [{"file": "path/to/file", "summary": "One-paragraph summary of the file purpose/functionality"}]')
     if "suggest_refactoring" in requested_analyses:
-        json_structure_description += '    "refactoring_suggestions": [{"file": "path/to/file", "line": number, "area": "e.g., function name, class name", "suggestion": "Description of refactoring suggestion"}],\n'
-
-    # Remove trailing comma and add closing brace
-    if json_structure_description.endswith(',\n'):
-        json_structure_description = json_structure_description[:-2] + "\n}"
-    else:
-         json_structure_description += "}" # Handle case where no sections selected (though UI should prevent)
+        structure_parts.append('    "refactoring_suggestions": [{"file": "path/to/file", "line": number, "area": "e.g., function name, class name", "suggestion": "Description of refactoring suggestion"}]')
 
+    json_structure_description += ",\n".join(structure_parts)
+    json_structure_description += "\n}"
 
     prompt_footer = f"""
 **Analysis Task:**
@@ -151,72 +161,136 @@ Respond ONLY with a single, valid JSON object adhering strictly to the following
 **JSON Output Only:**
 """
     full_prompt = prompt_content + prompt_footer
-    # print(f"--- PROMPT (First 500 chars): ---\n{full_prompt[:500]}\n--------------------------") # Debug: Print start of prompt
+    # print(f"--- PROMPT (First 500 chars): ---\n{full_prompt[:500]}\n--------------------------")
+    # print(f"--- PROMPT (Last 500 chars): ---\n{full_prompt[-500:]}\n--------------------------")
     return full_prompt, included_files
 
 
 def call_gemini_api(prompt):
-    """Calls the Gemini API and attempts to parse the JSON response."""
+    """Calls the Gemini API or returns mock data based on session state."""
     if not prompt:
         return None, "Prompt generation failed."
 
-    try:
-        st.write(f"📡 Sending request to {GEMINI_MODEL_NAME}...") # Progress update
-        response = model.generate_content(
-            prompt,
-            generation_config=genai.types.GenerationConfig(
-                # candidate_count=1, # Default is 1
-                # stop_sequences=['...'], # Optional stop sequences
-                # max_output_tokens=..., # Can be useful, but JSON structure might vary
-                temperature=0.2 # Lower temperature for more deterministic code analysis
-            ),
-             safety_settings=[ # Adjust as needed, might need to be less strict for code
-                {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-                {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-                {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-                # Be cautious with dangerous content, code analysis might trigger it
-                {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-            ]
-        )
-        st.write("✅ Response received from AI.")
-
-        # Debug: Print raw response
-        # print(f"--- RAW API RESPONSE ---\n{response.text}\n------------------------")
-
-        # Attempt to parse the JSON response - more robust extraction
+    # --- MOCK MODE LOGIC ---
+    if st.session_state.mock_api_call:
+        st.info(" MOCK MODE: Simulating API call...")
+        time.sleep(2) # Simulate network/processing delay
+
+        # --- CHOOSE YOUR MOCK RESPONSE ---
+        # Option 1: Simulate successful response with some data
+        mock_json_response = json.dumps({
+            "documentation_suggestions": [{"file": "mock/core.py", "line": 15, "suggestion": "def process_data(data):\n    \"\"\"Processes the input data using mock logic.\"\"\""}],
+            "potential_bugs": [{"file":"mock/utils.py", "line": 22, "description":"Potential division by zero if denominator is not checked.", "severity":"Medium"}],
+            "style_issues": [{"file": "mock/core.py", "line": 5, "description": "Variable 'varName' does not follow snake_case convention."}],
+            "module_summaries": [{"file": "mock/core.py", "summary": "This file contains the core mock processing logic."}, {"file":"mock/utils.py", "summary": "Utility functions for mocking."}],
+            "refactoring_suggestions": [{"file":"mock/utils.py", "line": 30, "area":"calculate_metrics function", "suggestion": "Function is too long (> 50 lines), consider breaking it down."}]
+        })
+        st.success("Mock response generated successfully.")
+        return json.loads(mock_json_response), None # Return insights, no error
+
+        # Option 2: Simulate API error
+        # st.error("Simulating API error.")
+        # return None, "MOCK ERROR: Simulated API Quota Exceeded."
+
+        # Option 3: Simulate invalid JSON response
+        # st.warning("Simulating invalid JSON response from AI.")
+        # return {"raw_response": "{malformed json'"}, "AI response was not valid JSON, showing raw text."
+        #
+        # Option 4: Simulate empty results
+        # mock_empty_json = json.dumps({
+        #     "documentation_suggestions": [], "potential_bugs": [], "style_issues": [],
+        #     "module_summaries": [], "refactoring_suggestions": []
+        # })
+        # st.success("Mock response generated (empty results).")
+        # return json.loads(mock_empty_json), None
+    # --- END MOCK MODE LOGIC ---
+
+
+    # --- REAL API CALL LOGIC ---
+    else:
+        if not initialize_gemini_model(): # Ensure model is ready
+             return None, "Gemini Model Initialization Failed."
+        if model is None: # Should not happen if initialize check passed, but safeguard
+             return None, "Gemini model not available."
+
         try:
-            # Find the start and end of the JSON block
-            json_start = response.text.find('{')
-            json_end = response.text.rfind('}') + 1
-            if json_start != -1 and json_end != -1:
-                 json_response_text = response.text[json_start:json_end]
-                 insights = json.loads(json_response_text)
-                 return insights, None
-            else:
-                 # Fallback if no {} found - maybe simple text response?
-                 st.warning("⚠️ Could not find JSON structure in response. Displaying raw text.")
-                 return {"raw_response": response.text}, "AI response was not valid JSON, showing raw text."
-
-        except json.JSONDecodeError as json_err:
-            st.error(f"🚨 Error parsing JSON response from AI: {json_err}")
-            st.error("Raw AI Response:")
-            st.code(response.text, language='text')
-            return None, f"AI response was not valid JSON: {json_err}"
-        except Exception as e:
-             st.error(f"🚨 Unexpected issue processing AI response: {e}")
-             try: st.code(f"Response object: {response}", language='text')
-             except: pass
-             return None, f"Unexpected response structure: {e}"
+            st.write(f"📡 Sending request to {GEMINI_MODEL_NAME}...")
+            response = model.generate_content(
+                prompt,
+                generation_config=genai.types.GenerationConfig(temperature=0.2),
+                safety_settings=[
+                    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+                    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+                    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+                    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+                ]
+            )
+            st.write("✅ Response received from AI.")
+
+            # Debug: Print raw response text
+            # print(f"--- RAW API RESPONSE ---\n{response.text}\n------------------------")
+
+            try:
+                # Try to extract JSON robustly
+                json_response_text = response.text.strip()
+                # Handle potential markdown code block fences
+                if json_response_text.startswith("```json"):
+                    json_response_text = json_response_text[7:]
+                if json_response_text.startswith("```"): # Handle case where ```json wasn't used
+                     json_response_text = json_response_text[3:]
+                if json_response_text.endswith("```"):
+                    json_response_text = json_response_text[:-3]
+
+                # Find the first '{' and the last '}'
+                json_start = json_response_text.find('{')
+                json_end = json_response_text.rfind('}') + 1
+
+                if json_start != -1 and json_end != -1 and json_end > json_start:
+                    final_json_text = json_response_text[json_start:json_end]
+                    insights = json.loads(final_json_text)
+                    return insights, None
+                else:
+                    st.warning("⚠️ Could not find valid JSON object boundaries ({...}) in response. Displaying raw text.")
+                    return {"raw_response": response.text}, "AI response did not contain clear JSON object, showing raw text."
+
+            except json.JSONDecodeError as json_err:
+                st.error(f"🚨 Error parsing JSON response from AI: {json_err}")
+                st.error("Raw AI Response:")
+                st.code(response.text, language='text')
+                return None, f"AI response was not valid JSON: {json_err}"
+            except AttributeError:
+                 # Handle cases where response structure might be different (e.g. blocked)
+                 st.error(f"🚨 Unexpected API response structure.")
+                 st.code(f"Response object: {response}", language='text') # Log the problematic response
+                 # Try to get blocked reason if available
+                 try:
+                     block_reason = response.prompt_feedback.block_reason
+                     if block_reason:
+                          return None, f"Content blocked by API. Reason: {block_reason}"
+                 except Exception:
+                      pass # Ignore if feedback structure isn't as expected
+                 return None, "Unexpected response structure from API."
+            except Exception as e:
+                st.error(f"🚨 Unexpected issue processing AI response: {e}")
+                try: st.code(f"Response object: {response}", language='text')
+                except: pass
+                return None, f"Unexpected response structure: {e}"
 
-    except Exception as e:
-        st.error(f"🚨 An error occurred during API call: {e}")
-        # Add specific error checks if possible (e.g., quota, safety blocks)
-        error_msg = f"API call failed: {e}"
-        if "block_reason: SAFETY" in str(e):
-             error_msg = "Content blocked due to safety settings. Code or AI response may have triggered filters."
-        elif "429" in str(e): # Crude check for quota error
-             error_msg = "API Quota Exceeded or Rate Limit hit. Check your Google AI Studio dashboard."
-        return None, error_msg
+        except Exception as e:
+            st.error(f"🚨 An error occurred during API call: {e}")
+            error_msg = f"API call failed: {e}"
+            # Improved error identification
+            if hasattr(e, 'message'): # For google.api_core.exceptions
+                 if "429" in e.message:
+                     error_msg = "API Quota Exceeded or Rate Limit hit. Check your Google Cloud/AI Studio dashboard."
+                 elif "API key not valid" in e.message:
+                     error_msg = "Invalid Gemini API Key. Please check `.streamlit/secrets.toml`."
+                 elif "blocked" in e.message.lower(): # General check for safety blocks
+                     error_msg = "Content blocked due to safety settings. Review input code or adjust safety settings if appropriate."
+            elif "block_reason: SAFETY" in str(e): # Fallback check
+                 error_msg = "Content blocked due to safety settings. Review input code or adjust safety settings if appropriate."
+
+            return None, error_msg
 
 
 def display_results(results_json, requested_analyses):
@@ -225,72 +299,73 @@ def display_results(results_json, requested_analyses):
 
     if not isinstance(results_json, dict):
         st.error("Invalid results format received.")
-        st.json(results_json) # Show what was received
+        st.json(results_json)
         return
 
-    # Handle raw response fallback
     if "raw_response" in results_json:
         st.subheader("Raw AI Response (JSON Parsing Failed)")
         st.code(results_json["raw_response"], language='text')
         return
 
-    # Display each requested section
-    if "generate_docs" in requested_analyses:
-        st.subheader(AVAILABLE_ANALYSES["generate_docs"])
-        suggestions = results_json.get("documentation_suggestions", [])
-        if suggestions:
-            for item in suggestions:
-                st.markdown(f"- **File:** `{item.get('file', 'N/A')}` (Line: {item.get('line', 'N/A')})")
-                st.code(item.get('suggestion', ''), language='text') # Show suggestion as code/text
-        else:
-            st.markdown("_No documentation suggestions provided._")
-        st.divider()
-
-    if "find_bugs" in requested_analyses:
-        st.subheader(AVAILABLE_ANALYSES["find_bugs"])
-        bugs = results_json.get("potential_bugs", [])
-        if bugs:
-            for item in bugs:
-                 st.markdown(f"- **File:** `{item.get('file', 'N/A')}` (Line: {item.get('line', 'N/A')}) - **Severity:** {item.get('severity', 'Unknown')}")
-                 st.markdown(f"  Description: {item.get('description', 'N/A')}")
-        else:
-            st.markdown("_No potential bugs identified._")
-        st.divider()
-
-    if "check_style" in requested_analyses:
-        st.subheader(AVAILABLE_ANALYSES["check_style"])
-        issues = results_json.get("style_issues", [])
-        if issues:
-             for item in issues:
-                st.markdown(f"- **File:** `{item.get('file', 'N/A')}` (Line: {item.get('line', 'N/A')})")
-                st.markdown(f"  Issue: {item.get('description', 'N/A')}")
-        else:
-            st.markdown("_No style issues identified._")
-        st.divider()
-
-    if "summarize_modules" in requested_analyses:
-        st.subheader(AVAILABLE_ANALYSES["summarize_modules"])
-        summaries = results_json.get("module_summaries", [])
-        if summaries:
-             for item in summaries:
-                 st.markdown(f"**File:** `{item.get('file', 'N/A')}`")
-                 st.markdown(f"> {item.get('summary', 'N/A')}") # Blockquote for summary
-        else:
-            st.markdown("_No module summaries provided._")
-        st.divider()
-
-    if "suggest_refactoring" in requested_analyses:
-        st.subheader(AVAILABLE_ANALYSES["suggest_refactoring"])
-        suggestions = results_json.get("refactoring_suggestions", [])
-        if suggestions:
-             for item in suggestions:
-                 st.markdown(f"- **File:** `{item.get('file', 'N/A')}` (Line: {item.get('line', 'N/A')}) - **Area:** {item.get('area', 'N/A')}")
-                 st.markdown(f"  Suggestion: {item.get('suggestion', 'N/A')}")
+    # Define display functions for clarity
+    def display_list_items(items, fields):
+        if items:
+            for item in items:
+                details = []
+                for field_key, field_label in fields.items():
+                    value = item.get(field_key, 'N/A')
+                    if value != 'N/A': # Only show if value exists
+                       details.append(f"**{field_label}:** {value}")
+                st.markdown("- " + " - ".join(details))
+                # Handle specific multi-line outputs like suggestions/summaries
+                if 'suggestion' in item:
+                     st.code(item['suggestion'], language='text')
+                elif 'description' in item:
+                     st.markdown(f"  > {item['description']}") # Indent description
+                elif 'summary' in item:
+                     st.markdown(f"  > {item['summary']}") # Indent summary
         else:
-            st.markdown("_No refactoring suggestions provided._")
+            st.markdown("_No items found for this category._")
         st.divider()
 
-    # Option to download the raw JSON results
+    # Map keys to display configurations
+    display_config = {
+        "generate_docs": {
+            "key": "documentation_suggestions", "title": AVAILABLE_ANALYSES["generate_docs"],
+            "fields": {"file": "File", "line": "Line"} # Suggestion shown by st.code
+        },
+        "find_bugs": {
+            "key": "potential_bugs", "title": AVAILABLE_ANALYSES["find_bugs"],
+            "fields": {"file": "File", "line": "Line", "severity": "Severity"} # Description shown separately
+        },
+        "check_style": {
+            "key": "style_issues", "title": AVAILABLE_ANALYSES["check_style"],
+            "fields": {"file": "File", "line": "Line"} # Description shown separately
+        },
+        "summarize_modules": {
+            "key": "module_summaries", "title": AVAILABLE_ANALYSES["summarize_modules"],
+            "fields": {"file": "File"} # Summary shown separately
+        },
+        "suggest_refactoring": {
+            "key": "refactoring_suggestions", "title": AVAILABLE_ANALYSES["suggest_refactoring"],
+            "fields": {"file": "File", "line": "Line", "area": "Area"} # Suggestion shown separately
+        }
+    }
+
+    # Iterate and display selected sections
+    any_results = False
+    for analysis_key in requested_analyses:
+        if analysis_key in display_config:
+            config = display_config[analysis_key]
+            st.subheader(config["title"])
+            items = results_json.get(config["key"], [])
+            display_list_items(items, config["fields"])
+            if items: any_results = True
+
+    if not any_results:
+         st.info("No specific findings were identified in the analysis based on your selections.")
+
+    # Download button
     st.download_button(
          label="Download Full Report (JSON)",
          data=json.dumps(results_json, indent=4),
@@ -298,85 +373,112 @@ def display_results(results_json, requested_analyses):
          mime="application/json"
      )
 
-
 # --- Streamlit App Main Interface ---
 st.set_page_config(page_title="Codebase Audit Assistant", layout="wide")
 
 st.title("🤖 Codebase Audit & Documentation Assistant")
 st.markdown(f"Upload your codebase (`.zip`) for analysis using **{GEMINI_MODEL_NAME}**.")
-st.warning("⚠️ **Privacy Notice:** Your code content will be sent to the Google Gemini API for analysis. Do not upload highly sensitive or proprietary code if you are not comfortable with this.")
-
-# Sidebar for options
-st.sidebar.header("🛠️ Analysis Options")
-selected_analyses = []
-for key, name in AVAILABLE_ANALYSES.items():
-     if st.sidebar.checkbox(name, value=True, key=f"cb_{key}"):
-         selected_analyses.append(key)
-
-st.sidebar.header("📄 How To Use")
-st.sidebar.info(
-    "1. Ensure `GEMINI_API_KEY` is in `.streamlit/secrets.toml`.\n"
-    "2. Select desired analyses in the sidebar.\n"
-    "3. Create a **ZIP archive** of your codebase.\n"
-    "4. Upload the `.zip` file below.\n"
-    "5. Click 'Analyze Codebase'.\n"
-    "6. Review the report generated."
-)
-st.sidebar.info(f"**Note:** Only files with common code extensions ({', '.join(CODE_EXTENSIONS)}) within the ZIP will be processed. Analysis might be limited by token estimates (~{MAX_PROMPT_TOKENS_ESTIMATE} tokens).")
+
+# Sidebar controls
+with st.sidebar:
+    st.header("⚙️ Analysis Controls")
+    # Mock Mode Toggle
+    st.session_state.mock_api_call = st.toggle("🧪 Enable Mock API Mode (for Testing)", value=st.session_state.mock_api_call,
+                                               help="If enabled, uses fake data instead of calling the real Gemini API. Saves cost during testing.")
+    if st.session_state.mock_api_call:
+        st.info("Mock API Mode ACTIVE")
+    else:
+        st.info("Using REAL Gemini API")
+
+    st.divider()
+    st.header("🔎 Select Analyses")
+    selected_analyses = []
+    for key, name in AVAILABLE_ANALYSES.items():
+        if st.checkbox(name, value=True, key=f"cb_{key}"):
+            selected_analyses.append(key)
+
+    st.divider()
+    st.header("📄 How To Use")
+    st.info(
+        "1. Set API Key in `.streamlit/secrets.toml` (if not using Mock Mode).\n"
+        "2. Toggle Mock Mode if needed.\n"
+        "3. Select desired analyses.\n"
+        "4. Create a **ZIP archive** of your codebase.\n"
+        "5. Upload the `.zip` file.\n"
+        "6. Click 'Analyze Codebase'.\n"
+        "7. Review the report."
+    )
+    st.info(f"**Note:** Only files with common code extensions ({', '.join(CODE_EXTENSIONS)}) are processed. Analysis might be limited (~{MAX_PROMPT_TOKENS_ESTIMATE:,} est. tokens).")
+
+    st.divider()
+    st.warning("⚠️ **Privacy:** Code content is sent to the Google Gemini API if Mock Mode is OFF. Do not upload sensitive code if uncomfortable.")
 
 
 # Main content area
-uploaded_file = st.file_uploader("📁 Upload Codebase ZIP File", type=['zip'])
+uploaded_file = st.file_uploader("📁 Upload Codebase ZIP File", type=['zip'], key="file_uploader")
+
+analysis_triggered = False
+results_cache = None # To store results briefly
 
 if uploaded_file:
-    st.success(f"✅ File '{uploaded_file.name}' uploaded successfully.")
+    st.success(f"✅ File '{uploaded_file.name}' uploaded.")
 
-    # Process the zip file immediately to give feedback
     with st.spinner("Inspecting ZIP file..."):
         code_files, total_chars, file_count, ignored_files = process_zip_file(uploaded_file)
 
     if code_files is not None:
-        st.info(f"Found **{file_count}** relevant code files ({total_chars:,} characters). Estimated tokens: ~{estimate_token_count(total_chars):,}")
+        st.info(f"Found **{file_count}** relevant code files ({total_chars:,} characters). Est. tokens: ~{estimate_token_count(total_chars):,}")
         if ignored_files:
              with st.expander(f"View {len(ignored_files)} Skipped/Ignored Files"):
-                 st.json(ignored_files)
-
-        # Analysis Button
-        analyze_button = st.button("Analyze Codebase", type="primary", disabled=(not selected_analyses or file_count == 0))
-
-        if not selected_analyses and analyze_button:
-             st.warning("Please select at least one analysis type from the sidebar.")
-        elif file_count == 0 and analyze_button:
-             st.warning("No relevant code files found in the ZIP archive to analyze.")
-
-
-        if analyze_button and selected_analyses and file_count > 0:
-            st.divider()
-            with st.spinner(f"🚀 Preparing prompt and contacting {GEMINI_MODEL_NAME}... This may take several minutes for large codebases."):
-                # 1. Construct the prompt
-                analysis_prompt, included_files_in_prompt = construct_analysis_prompt(code_files, selected_analyses)
-
-                if analysis_prompt and included_files_in_prompt:
-                     st.write(f"Analyzing {len(included_files_in_prompt)} files...")
-                     # 2. Call the API
-                     results_json, error_message = call_gemini_api(analysis_prompt)
-
-                     # 3. Display Results
-                     if error_message:
-                         st.error(f"Analysis Failed: {error_message}")
-                     elif results_json:
-                         display_results(results_json, selected_analyses)
-                     else:
-                          st.error("Analysis did not return results or an unknown error occurred.")
-                elif not included_files_in_prompt:
-                     st.error("Could not proceed: No files were included in the prompt (likely due to token limits or processing errors).")
-
-
+                 # Use st.code for better formatting of list
+                 st.code("\n".join(ignored_files), language='text')
+
+        analyze_button_disabled = (not selected_analyses or file_count == 0)
+        analyze_button_label = "Analyze Codebase" if not analyze_button_disabled else "Select Analyses or Upload Valid Code"
+        if st.button(analyze_button_label, type="primary", disabled=analyze_button_disabled):
+            analysis_triggered = True
+
+            if not selected_analyses:
+                 st.warning("Please select at least one analysis type from the sidebar.")
+            elif file_count == 0:
+                 st.warning("No relevant code files found in the ZIP archive to analyze.")
+            else:
+                st.divider()
+                with st.spinner(f"🚀 Preparing prompt & contacting AI ({'Mock Mode' if st.session_state.mock_api_call else GEMINI_MODEL_NAME})... This may take time."):
+                    # 1. Construct Prompt
+                    analysis_prompt, included_files_in_prompt = construct_analysis_prompt(code_files, selected_analyses)
+
+                    if analysis_prompt and included_files_in_prompt:
+                        st.write(f"Analyzing {len(included_files_in_prompt)} files...")
+                        # 2. Call API (Real or Mock)
+                        results_json, error_message = call_gemini_api(analysis_prompt)
+                        results_cache = (results_json, error_message) # Store results
+                    elif not included_files_in_prompt:
+                        results_cache = (None, "Could not proceed: No files included in prompt (check token limits/errors).")
+                    else:
+                         results_cache = (None, "Failed to generate analysis prompt.")
+
+    else: # Error during zip processing
+        pass # Error message already shown
+
+# Display results outside the button click block if analysis was triggered
+if analysis_triggered and results_cache:
+    results_json, error_message = results_cache
+    st.divider()
+    if error_message:
+        st.error(f"Analysis Failed: {error_message}")
+        # Display partial results if available (e.g., raw response on JSON error)
+        if results_json and isinstance(results_json, dict) and "raw_response" in results_json:
+             st.subheader("Raw AI Response")
+             st.code(results_json["raw_response"], language='text')
+
+    elif results_json:
+        display_results(results_json, selected_analyses)
     else:
-        # Error message already shown by process_zip_file
-        pass
+        st.error("Analysis did not return results or an unknown error occurred.")
+
 
-else:
+elif not uploaded_file:
     st.info("Upload a ZIP file containing your source code to begin.")
 
 st.divider()