Update app.py

app.py

@@ -1,516 +1,383 @@
-# advanced_archsketch_app.py
-import os
 import streamlit as st
-from streamlit_drawable_canvas import st_canvas
-from PIL import Image, ImageDraw, ImageFont, UnidentifiedImageError
-import requests # For potential real API calls later
-import openai # Used notionally
-from io import BytesIO
+import google.generativeai as genai
+import zipfile
+import io
 import json
-import uuid
-import time
-import random
-import base64 # For potential image encoding if needed
-
-# ─── 1. Configuration & Secrets ─────────────────────────────────────────────
+import os
+from pathlib import Path
+
+# --- 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
+
+# Define the types of analysis available
+AVAILABLE_ANALYSES = {
+    "generate_docs": "Generate Missing Docstrings/Comments",
+    "find_bugs": "Identify Potential Bugs & Anti-patterns",
+    "check_style": "Check Style Guide Compliance (General)",
+    "summarize_modules": "Summarize Complex Modules/Files",
+    "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'}
+
+# --- Gemini API Setup ---
 try:
-    openai.api_key = st.secrets["OPENAI_API_KEY"]
-except Exception:
-    st.error("OpenAI API Key not found. Please set it in Streamlit secrets.")
-    # openai.api_key = "YOUR_FALLBACK_KEY_FOR_LOCAL_TESTING" # Or load from env
-
-st.set_page_config(page_title="ArchSketch AI [Advanced]", layout="wide", page_icon="🏗️")
-
-# --- Simulated Backend API Endpoints ---
-# Replace with your actual endpoints if building a real backend
-API_SUBMIT_URL = "http://your-backend.com/api/v1/submit_arch_job"
-API_STATUS_URL = "http://your-backend.com/api/v1/job_status/{job_id}"
-API_RESULT_URL = "http://your-backend.com/api/v1/job_result/{job_id}" # Might return data directly or a URL
-
-# ─── 2. State Initialization & Authentication ───────────────────────────────
-
-def initialize_state():
-    """Initializes all necessary session state variables."""
-    defaults = {
-        'logged_in': False,
-        'username': None,
-        'current_job_id': None,
-        'job_status': 'IDLE', # IDLE, SUBMITTED, PENDING, PROCESSING, COMPLETED, FAILED
-        'job_progress': {}, # Progress dict per job_id
-        'job_errors': {},   # Error dict per job_id
-        'job_results': {},  # Stores result data/references per job_id {job_id: {'type': 'image'/'svg'/'json', 'data': path_or_data, 'params':{...}, 'prompt': '...'}}
-        'selected_history_job_id': None,
-        'annotations': {}, # {job_id: [annotation_objects]}
-        # Input specific state
-        'input_prompt': "",
-        'input_staging_image_bytes': None,
-        'input_staging_image_preview': None,
-        'input_filename': None, # Store filename of uploaded staging image
-    }
-    for key, value in defaults.items():
-        if key not in st.session_state:
-            st.session_state[key] = value
-
-initialize_state()
-
-def show_login_form():
-    """Displays the login form."""
-    st.warning("Login Required")
-    with st.form("login_form"):
-        username = st.text_input("Username", key="login_user")
-        password = st.text_input("Password", type="password", key="login_pass")
-        submitted = st.form_submit_button("Login")
-        if submitted:
-            # --- !!! INSECURE - DEMO ONLY !!! ---
-            if username == "arch_user" and password == "pass123":
-                st.session_state.logged_in = True
-                st.session_state.username = username
-                st.success("Login successful!")
-                time.sleep(1)
-                st.rerun()
-            else:
-                st.error("Invalid credentials.")
-
-# --- Authentication Gate ---
-if not st.session_state.logged_in:
-    show_login_form()
+    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()
 
-# ─── 3. Simulated Backend Interaction Functions ───────────────────────────────
-
-def submit_job_to_backend(payload: dict) -> tuple[str | None, str | None]:
-    """Simulates submitting job, returns (job_id, error)."""
-    st.info("Submitting job to backend simulation...")
-    print(f"SIMULATING API SUBMIT to {API_SUBMIT_URL}")
-    # In reality: response = requests.post(API_SUBMIT_URL, json=payload, headers=auth_headers)
-    time.sleep(1.5) # Simulate network + queue time
-    if random.random() < 0.95:
-        job_id = f"archjob_{uuid.uuid4().hex[:12]}"
-        print(f"API Submit SUCCESS: Job ID = {job_id}")
-        st.session_state.job_progress[job_id] = 0
-        st.session_state.job_errors[job_id] = None
-        # Store essential info with job immediately
-        st.session_state.job_results[job_id] = {
-            'type': None, 'data': None, # Will be filled on completion
-            'params': payload.get('parameters', {}), # Store settings used
-            'prompt': payload.get('prompt', '')
-        }
-        return job_id, None
-    else:
-        error_msg = "Simulated API Error: Failed to submit (server busy/invalid payload)."
-        print(f"API Submit FAILED: {error_msg}")
-        return None, error_msg
+# --- 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
 
-def check_job_status_backend(job_id: str) -> tuple[str, dict | None]:
-    """Simulates checking job status, returns (status, result_info | None)."""
-    status_url = API_STATUS_URL.format(job_id=job_id)
-    print(f"SIMULATING API STATUS CHECK: {status_url}")
-    # In reality: response = requests.get(status_url, headers=auth_headers)
-    time.sleep(0.7) # Simulate network latency
-
-    if job_id not in st.session_state.job_progress:
-        st.session_state.job_progress[job_id] = 0
-
-    current_progress = st.session_state.job_progress[job_id]
-    status = "UNKNOWN"
-    result_info = None
-
-    # Simulate progress and potential states
-    if current_progress < 0.1:
-        status = "PENDING"
-        st.session_state.job_progress[job_id] += random.uniform(0.05, 0.15)
-    elif current_progress < 0.9:
-        status = "PROCESSING"
-        st.session_state.job_progress[job_id] += random.uniform(0.1, 0.3)
-        # Simulate potential failure during processing
-        if random.random() < 0.03: # 3% chance of failure mid-run
-             status = "FAILED"
-             st.session_state.job_errors[job_id] = "Simulated AI failure during processing."
-             print(f"API Status SIMULATION: Job {job_id} FAILED processing.")
-    elif current_progress >= 0.9: # Consider it done
-        status = "COMPLETED"
-        print(f"API Status SIMULATION: Job {job_id} COMPLETED.")
-        # Determine simulated result type based on original request stored in job_results
-        job_mode = st.session_state.job_results.get(job_id, {}).get('params', {}).get('mode', 'Unknown')
-
-        if job_mode == "Floor Plan":
-            # Simulate returning path to an SVG or structured JSON data
-            placeholder_path = "assets/placeholder_floorplan.svg" # Need this file
-            if not os.path.exists(placeholder_path): placeholder_path = "assets/placeholder_floorplan.json" # Fallback - need JSON too
-            result_info = {'type': 'svg' if '.svg' in placeholder_path else 'json', 'data_path': placeholder_path}
-        else: # Virtual Staging
-            placeholder_path = "assets/placeholder_image.png" # Need this file
-            result_info = {'type': 'image', 'data_path': placeholder_path}
-
-    print(f"API Status SIMULATION: Job {job_id} Status={status}, Progress={st.session_state.job_progress.get(job_id, 0):.2f}")
-    return status, result_info
-
-def fetch_result_data(result_info: dict):
-    """Simulates fetching/loading result data based on info from status check."""
-    result_type = result_info['type']
-    data_path = result_info['data_path'] # In real app, might be URL
-    print(f"SIMULATING Fetching {result_type} result from: {data_path}")
-    # In reality: if URL, use requests.get(data_path).content
-
-    if not os.path.exists(data_path):
-        print(f"ERROR: Result placeholder not found at {data_path}")
-        raise FileNotFoundError(f"Result file missing: {data_path}")
+def process_zip_file(uploaded_file):
+    """Extracts code files and their content from the uploaded zip file."""
+    code_files = {}
+    total_chars = 0
+    file_count = 0
+    ignored_files = []
 
     try:
-        if result_type == 'image':
-            img = Image.open(data_path).convert("RGB")
-            return img
-        elif result_type == 'svg':
-            with open(data_path, 'r', encoding='utf-8') as f:
-                svg_content = f.read()
-            return svg_content # Return raw SVG string
-        elif result_type == 'json':
-            with open(data_path, 'r', encoding='utf-8') as f:
-                json_data = json.load(f)
-            return json_data # Return parsed JSON
-        else:
-            raise ValueError(f"Unsupported result type: {result_type}")
+        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:
+                    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:
+                                try:
+                                    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
+
+                            code_files[member.filename] = content
+                            total_chars += len(content)
+                            file_count += 1
+                    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})")
+
+    except zipfile.BadZipFile:
+        st.error("🚨 Invalid or corrupted ZIP file.")
+        return None, 0, 0, []
     except Exception as e:
-        print(f"ERROR loading result from {data_path}: {e}")
-        raise
-
-# ─── 4. Sidebar UI ───────────────────────────────────────────────────────────
-with st.sidebar:
-    st.header(f"🏗️ ArchSketch AI")
-    st.caption(f"User: {st.session_state.username}")
-    if st.button("Logout", key="logout_btn"):
-        # Clear sensitive parts of state, re-initialize others
-        keys_to_clear = list(st.session_state.keys())
-        for key in keys_to_clear:
-            del st.session_state[key]
-        initialize_state()
-        st.rerun()
-    st.markdown("---")
-
-    st.header("⚙️ Project Configuration")
-
-    # Disable controls while job is active
-    ui_disabled = st.session_state.job_status in ["SUBMITTED", "PENDING", "PROCESSING"]
-
-    mode = st.radio("Mode", ["Floor Plan", "Virtual Staging"], key="mode_radio", disabled=ui_disabled)
-
-    # --- Conditional Input for Staging ---
-    if mode == "Virtual Staging":
-        staging_image_file = st.file_uploader(
-            "Upload Empty Room Image:",
-            type=["png", "jpg", "jpeg", "webp"],
-            key="staging_uploader",
-            disabled=ui_disabled,
-            help="Required for Virtual Staging mode."
+        st.error(f"🚨 Error processing ZIP file: {e}")
+        return None, 0, 0, []
+
+    return code_files, total_chars, file_count, ignored_files
+
+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:
+            concatenated_code += segment
+            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
+
+    if not included_files:
+        st.error("🚨 No code files could be included within the estimated token limit.")
+        return None, []
+
+    prompt_content += concatenated_code
+
+    # Define the requested JSON structure based on selections
+    json_structure_description = "{\n"
+    if "generate_docs" in requested_analyses:
+        json_structure_description += '    "documentation_suggestions": [{"file": "path/to/file", "line": number, "suggestion": "Suggested docstring/comment"}],\n'
+    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'
+    if "check_style" in requested_analyses:
+        json_structure_description += '    "style_issues": [{"file": "path/to/file", "line": number, "description": "Description of style deviation"}],\n'
+    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'
+    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)
+
+
+    prompt_footer = f"""
+**Analysis Task:**
+Perform the analyses corresponding to the keys present in the JSON structure below, based *only* on the provided code files ({', '.join(included_files)}).
+
+**Output Format:**
+Respond ONLY with a single, valid JSON object adhering strictly to the following structure. If no issues/suggestions are found for a category, provide an empty list `[]`. Do not include explanations outside the JSON structure.
+
+{json_structure_description}
+
+**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
+    return full_prompt, included_files
+
+
+def call_gemini_api(prompt):
+    """Calls the Gemini API and attempts to parse the JSON response."""
+    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"},
+            ]
         )
-        if staging_image_file:
-            if staging_image_file.name != st.session_state.input_filename: # Detect new upload
-                 st.info("Processing staging image...")
-                 try:
-                     img_bytes = staging_image_file.getvalue()
-                     image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-                     image.thumbnail((1024, 1024), Image.Resampling.LANCZOS) # Resize preview
-                     st.session_state.input_staging_image_bytes = img_bytes # Store bytes for API
-                     st.session_state.input_staging_image_preview = image
-                     st.session_state.input_filename = staging_image_file.name
-                     st.success("Staging image loaded.")
-                     # Don't rerun here, let user configure other options
-                 except UnidentifiedImageError:
-                     st.error("Invalid image file.")
-                     st.session_state.input_staging_image_bytes = None
-                     st.session_state.input_staging_image_preview = None
-                     st.session_state.input_filename = None
-                 except Exception as e:
-                      st.error(f"Error loading image: {e}")
-                      st.session_state.input_staging_image_bytes = None
-                      st.session_state.input_staging_image_preview = None
-                      st.session_state.input_filename = None
-
-        elif st.session_state.input_filename: # User cleared the uploader
-             st.session_state.input_staging_image_bytes = None
-             st.session_state.input_staging_image_preview = None
-             st.session_state.input_filename = None
-
-
-    st.markdown("---")
-    st.header("✨ AI Parameters")
-    # Note: Different models might be chosen by the backend based on mode/style
-    model_hint = st.selectbox("Model Preference (Hint for Backend)", ["Auto", "GPT‑4o (Text/Layout)", "Stable Diffusion (Image Gen)", "ControlNet (Editing)"], key="model_select", disabled=ui_disabled)
-    style = st.selectbox("Style Preset", ["Modern", "Minimalist", "Rustic", "Industrial", "Coastal", "Custom"], key="style_select", disabled=ui_disabled)
-    resolution = st.select_slider("Target Resolution (Approx.)", options=[512, 768, 1024], value=768, key="res_slider", disabled=ui_disabled)
-
-    with st.expander("Optional Metadata"):
-        project_id = st.text_input("Project ID", key="proj_id_input", disabled=ui_disabled)
-        location = st.text_input("Location / Address", key="loc_input", disabled=ui_disabled)
-        client_notes = st.text_area("Client Notes", key="notes_area", disabled=ui_disabled)
-
-# ─── 5. Main Area UI ─────────────────────────────────────────────────────────
-
-st.title("Advanced AI Architectural Visualizer")
-
-# --- Prompt Input Area ---
-st.subheader("📝 Describe Your Request")
-prompt_text = st.text_area(
-    "Enter detailed prompt:",
-    placeholder=(
-        "Floor Plan Example: 'Generate a detailed 2D floor plan SVG for a 4-bedroom modern farmhouse, approx 2500 sq ft, main floor master suite, large open concept kitchen/living area, separate office, mudroom entrance.'\n"
-        "Staging Example: 'Virtually stage the uploaded living room image in a minimalist Scandinavian style. Include a light grey sectional sofa, a geometric rug, light wood coffee table, and several potted plants. Ensure bright, natural lighting.'"
-    ),
-    height=150,
-    key="prompt_input",
-    disabled=ui_disabled # Disable if job running
-)
-st.session_state.input_prompt = prompt_text # Keep state updated
-
-# --- Submit Button ---
-can_submit = bool(st.session_state.input_prompt.strip())
-if mode == "Virtual Staging":
-    can_submit = can_submit and (st.session_state.input_staging_image_bytes is not None)
-
-submit_button = st.button(
-    "🚀 Submit Visualization Job",
-    key="submit_btn",
-    use_container_width=True,
-    disabled=ui_disabled or not can_submit,
-    help="Requires a prompt. Staging mode also requires an uploaded image."
+        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
+        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}"
+
+    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
+
+
+def display_results(results_json, requested_analyses):
+    """Renders the analysis results in Streamlit."""
+    st.header("📊 Analysis Report")
+
+    if not isinstance(results_json, dict):
+        st.error("Invalid results format received.")
+        st.json(results_json) # Show what was received
+        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')}")
+        else:
+            st.markdown("_No refactoring suggestions provided._")
+        st.divider()
+
+    # Option to download the raw JSON results
+    st.download_button(
+         label="Download Full Report (JSON)",
+         data=json.dumps(results_json, indent=4),
+         file_name="code_audit_report.json",
+         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).")
 
-if not can_submit and not ui_disabled:
-     if mode == "Virtual Staging" and not st.session_state.input_staging_image_bytes:
-         st.warning("Please upload an image for Virtual Staging mode.")
-     elif not st.session_state.input_prompt.strip():
-          st.warning("Please enter a prompt describing your request.")
-
-
-# --- Job Submission Logic ---
-if submit_button:
-    st.session_state.job_status = "SUBMITTED"
-    st.session_state.current_job_id = None # Clear old ID before new submission attempt
-    st.session_state.ai_result_image = None # Clear old result display
-
-    # Prepare Payload
-    api_payload = {
-        "prompt": st.session_state.input_prompt,
-        "parameters": {
-            "mode": mode,
-            "model_preference": model_hint,
-            "style": style,
-            "resolution": resolution,
-            "project_id": project_id,
-            "location": location,
-            "client_notes": client_notes,
-        },
-        "user_id": st.session_state.username,
-    }
-
-    # Add image data for staging mode (handle carefully in production!)
-    if mode == "Virtual Staging" and st.session_state.input_staging_image_bytes:
-        # Option 1: Send as base64 (simpler for demo, BAD for large files)
-        api_payload["base_image_b64"] = base64.b64encode(st.session_state.input_staging_image_bytes).decode('utf-8')
-        api_payload["base_image_filename"] = st.session_state.input_filename
-        # Option 2 (Production): Upload to S3/GCS first, send URL/key
-        # api_payload["base_image_url"] = "s3://bucket/path/to/uploaded_image.jpg"
-
-    job_id, error = submit_job_to_backend(api_payload)
-
-    if job_id:
-        st.session_state.current_job_id = job_id
-        st.session_state.job_status = "PENDING" # Move to pending after successful submit
-        st.session_state.selected_history_job_id = job_id # Auto-select the new job
-        # Store params with result structure immediately
-        if job_id in st.session_state.job_results:
-            st.session_state.job_results[job_id]['params'] = api_payload['parameters']
-            st.session_state.job_results[job_id]['prompt'] = api_payload['prompt']
-
-        st.success(f"Job submitted! ID: {job_id}. Status will update below.")
-        st.rerun() # Start the polling loop
-    else:
-        st.error(f"Job submission failed: {error}")
-        st.session_state.job_status = "FAILED"
 
+# Main content area
+uploaded_file = st.file_uploader("📁 Upload Codebase ZIP File", type=['zip'])
 
-# --- Status & Result Display Area ---
-st.markdown("---")
-st.subheader("📊 Job Status & Result")
+if uploaded_file:
+    st.success(f"✅ File '{uploaded_file.name}' uploaded successfully.")
+
+    # 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):,}")
+        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).")
 
-current_job_id = st.session_state.current_job_id
-status = st.session_state.job_status
 
-if not current_job_id:
-    st.info("Submit a job using the controls above.")
-else:
-    # Display status updates
-    if status == "SUBMITTED":
-        st.warning(f"Job Status: Submitted... Waiting for confirmation (ID: {current_job_id})")
-        time.sleep(2) # Short delay before first poll
-        st.rerun()
-    elif status == "PENDING":
-        st.info(f"Job Status: Pending in queue... (ID: {current_job_id})")
-        time.sleep(5) # Poll interval
-        st.rerun()
-    elif status == "PROCESSING":
-        progress = st.session_state.job_progress.get(current_job_id, 0)
-        st.progress(min(progress, 1.0), text=f"Job Status: Processing... ({int(min(progress,1.0)*100)}%) (ID: {current_job_id})")
-        time.sleep(3) # Poll interval during processing
-        st.rerun()
-    elif status == "COMPLETED":
-        st.success(f"Job Status: Completed! (ID: {current_job_id})")
-        # Result display handled below in results/history section
-    elif status == "FAILED":
-        error_msg = st.session_state.job_errors.get(current_job_id, "Unknown error")
-        st.error(f"Job Status: Failed! (ID: {current_job_id}) - Error: {error_msg}")
-    elif status == "IDLE":
-         st.info("Submit a job to see status.")
-    else: # Should not happen
-        st.error(f"Unknown Job Status: {status}")
-
-    # --- Status Update Logic (if job is active) ---
-    if status in ["SUBMITTED", "PENDING", "PROCESSING"]:
-        new_status, result_info = check_job_status_backend(current_job_id)
-        st.session_state.job_status = new_status
-
-        if new_status == "COMPLETED" and result_info:
-            try:
-                result_data = fetch_result_data(result_info)
-                # Store result data associated with job_id
-                st.session_state.job_results[current_job_id]['type'] = result_info['type']
-                st.session_state.job_results[current_job_id]['data'] = result_data
-                st.session_state.selected_history_job_id = current_job_id # Ensure completed job is selected
-                st.rerun() # Rerun to display result
-            except Exception as e:
-                st.error(f"Failed to load result data: {e}")
-                st.session_state.job_status = "FAILED"
-                st.session_state.job_errors[current_job_id] = f"Failed to load result: {e}"
-                st.rerun()
-        elif new_status == "FAILED":
-             if not st.session_state.job_errors.get(current_job_id):
-                  st.session_state.job_errors[current_job_id] = "Job failed during processing (unknown reason)."
-             st.rerun() # Rerun to show failed status
-
-
-# --- Result Display / History / Annotation Area ---
-st.markdown("---")
-col_results, col_history = st.columns([3, 1]) # Main area for result, smaller sidebar for history
-
-with col_history:
-    st.subheader("📚 History")
-    if not st.session_state.job_results:
-        st.caption("No jobs run yet in this session.")
-    else:
-        # Display history items (most recent first)
-        sorted_job_ids = sorted(st.session_state.job_results.keys(), reverse=True)
-        for job_id in sorted_job_ids:
-            job_info = st.session_state.job_results[job_id]
-            prompt_short = job_info.get('prompt', 'No Prompt')[:40] + "..." if len(job_info.get('prompt', '')) > 40 else job_info.get('prompt', 'No Prompt')
-            mode_display = job_info.get('params',{}).get('mode', '?')
-            item_label = f"[{mode_display}] {prompt_short}"
-
-            # Use button to select history item
-            if st.button(item_label, key=f"history_{job_id}", use_container_width=True,
-                         help=f"View result for Job ID: {job_id}\nPrompt: {job_info.get('prompt', '')}"):
-                st.session_state.selected_history_job_id = job_id
-                st.rerun() # Rerun to update the main display
-
-        if st.session_state.job_results:
-             st.download_button(
-                 "⬇️ Export History (JSON)",
-                 data=json.dumps(st.session_state.job_results, indent=2, default=str), # Default=str for non-serializable
-                 file_name="archsketch_history.json",
-                 mime="application/json"
-            )
-
-
-with col_results:
-    selected_job_id = st.session_state.selected_history_job_id
-    if not selected_job_id or selected_job_id not in st.session_state.job_results:
-         st.info("Select a job from the history panel to view details and annotate.")
     else:
-        result_info = st.session_state.job_results[selected_job_id]
-        result_type = result_info.get('type')
-        result_data = result_info.get('data')
-        result_params = result_info.get('params', {})
-        result_prompt = result_info.get('prompt', 'N/A')
-
-        st.subheader(f"🔍 Viewing Result: {selected_job_id}")
-        st.caption(f"**Mode:** {result_params.get('mode', 'N/A')} | **Style:** {result_params.get('style', 'N/A')}")
-        st.markdown(f"**Prompt:** *{result_prompt}*")
-
-        display_image = None # Image to use for canvas background
-
-        if result_type == 'image' and isinstance(result_data, Image.Image):
-            st.image(result_data, caption="Generated Visualization", use_column_width=True)
-            display_image = result_data
-            # Add image download button
-            buf = BytesIO(); result_data.save(buf, format="PNG")
-            st.download_button("⬇️ Download Image (PNG)", buf.getvalue(), f"{selected_job_id}_result.png", "image/png")
-
-        elif result_type == 'svg' and isinstance(result_data, str):
-             st.image(result_data, caption="Generated Floor Plan (SVG)", use_column_width=True)
-             # SVG Download
-             st.download_button("⬇️ Download SVG", result_data, f"{selected_job_id}_floorplan.svg", "image/svg+xml")
-             # Cannot easily use SVG as canvas background directly - maybe render SVG to PNG first?
-             st.warning("Annotation on SVG is not directly supported in this demo. Showing base image if available.")
-             # If staging mode produced SVG somehow (unlikely), use the input image for annotation context
-             if result_params.get('mode') == 'Virtual Staging' and st.session_state.input_staging_image_preview:
-                  display_image = st.session_state.input_staging_image_preview
-
-        elif result_type == 'json' and isinstance(result_data, dict):
-            st.json(result_data, expanded=False)
-            st.caption("Generated Structured Data (JSON)")
-            # JSON Download
-            st.download_button("⬇️ Download JSON", json.dumps(result_data, indent=2), f"{selected_job_id}_data.json", "application/json")
-            st.warning("Annotation not applicable for JSON results. Showing base image if available.")
-            if result_params.get('mode') == 'Virtual Staging' and st.session_state.input_staging_image_preview:
-                  display_image = st.session_state.input_staging_image_preview
-        elif result_data is None:
-             st.warning("Result data is not available for this job (may still be processing or failed).")
-        else:
-            st.error("Result type or data is invalid.")
-
-
-        # --- Annotation Canvas ---
-        if display_image:
-            st.markdown("---")
-            st.subheader("✏️ Annotate / Edit")
-            # Load existing annotations for this job_id if they exist
-            initial_drawing = {"objects": st.session_state.annotations.get(selected_job_id, [])}
-
-            canvas = st_canvas(
-                fill_color="rgba(255, 0, 0, 0.2)", # Red annotation
-                stroke_width=3,
-                stroke_color="#FF0000",
-                background_image=display_image,
-                update_streamlit=[" Mosul", "mouseup"], # Update on drawing release
-                height=500, # Adjust height as needed
-                width=700, # Adjust width as needed
-                drawing_mode=st.selectbox("Drawing tool:", ("freedraw", "line", "rect", "circle", "transform"), key=f"draw_mode_{selected_job_id}"),
-                key=f"canvas_{selected_job_id}" # Key tied to job ID
-                # Removed initial_drawing for simplicity now, add back if needed carefully
-            )
-
-            # Save annotations when canvas updates
-            if canvas.json_data is not None and canvas.json_data["objects"]:
-                 st.session_state.annotations[selected_job_id] = canvas.json_data["objects"]
-
-            # Display current annotations (optional) & Export
-            current_annotations = st.session_state.annotations.get(selected_job_id)
-            if current_annotations:
-                with st.expander("View/Export Current Annotations (JSON)"):
-                    st.json(current_annotations)
-                    st.download_button(
-                         "⬇️ Export Annotations",
-                         data=json.dumps({selected_job_id: current_annotations}, indent=2),
-                         file_name=f"{selected_job_id}_annotations.json",
-                         mime="application/json"
-                    )
-        else:
-             st.caption("Annotation requires a viewable image result.")
+        # Error message already shown by process_zip_file
+        pass
 
+else:
+    st.info("Upload a ZIP file containing your source code to begin.")
 
-# ─── Footer & Disclaimer ─────────────────────────────────────────────────────
-st.markdown("---")
-st.warning("""
-    **Disclaimer:** This is an **advanced conceptual blueprint**. User authentication is **not secure**.
-    Backend API calls, asynchronous job handling, status polling, AI model execution (image generation, floor plan logic, staging),
-    and result data fetching are **simulated**. Building the real backend requires substantial AI and infrastructure expertise.
-""")
+st.divider()
+st.markdown("_Assistant powered by Google Gemini._")