Update app.py

app.py

@@ -1,181 +1,51 @@
+# advanced_archsketch_app.py
+import os
 import streamlit as st
-from PIL import Image, UnidentifiedImageError
-import io
+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 json
+import uuid
 import time
-import os
 import random
-import json # For API simulation
-import requests # To simulate API calls (though we won't make real ones here)
-import uuid # For job IDs
-
-# ------------------------------------------------------------------------------
-# Page Configuration
-# ------------------------------------------------------------------------------
-st.set_page_config(
-    page_title="AI Real Estate Visualization Suite [PRO]",
-    layout="wide",
-    page_icon="🚀",
-    initial_sidebar_state="expanded"
-)
-
-# ------------------------------------------------------------------------------
-# Simulated Backend API Interaction
-# ------------------------------------------------------------------------------
-# Replace with your actual backend API endpoint
-BACKEND_API_URL = "http://your-production-backend.com/api/v1/visualize"
-# Replace with your actual status check endpoint
-STATUS_API_URL = "http://your-production-backend.com/api/v1/status/{job_id}"
-
-# --- Simulate API Call ---
-def submit_visualization_job(payload: dict) -> tuple[str | None, str | None]:
-    """
-    Simulates submitting a job to the backend API.
-    In reality, this would use 'requests.post'.
-
-    Args:
-        payload (dict): Data to send (params, input reference, etc.)
-
-    Returns:
-        tuple[str | None, str | None]: (job_id, error_message)
-    """
-    st.info("Submitting job to backend simulation...")
-    print(f"SIMULATING API SUBMIT to {BACKEND_API_URL}")
-    print("Payload (summary):")
-    print(f"  Mode: {payload.get('output_mode')}")
-    print(f"  Style: {payload.get('style')}")
-    print(f"  Input Type: {payload.get('input_type')}")
-    # Omit potentially large data like image bytes from console log in real scenario
-    # print(json.dumps(payload, indent=2)) # Don't print potentially large data
-
-    # Simulate network latency & backend processing start
-    time.sleep(1.5)
-
-    # Simulate success/failure
-    if random.random() < 0.95: # 95% success rate
-        job_id = f"job_{uuid.uuid4()}"
-        print(f"API Submit Simulation SUCCESS: Job ID = {job_id}")
-        # In a real app, store job_id associated with user/session
-        return job_id, None
-    else:
-        error_msg = "Simulated API Error: Failed to submit job (e.g., server busy, invalid params)."
-        print(f"API Submit Simulation FAILED: {error_msg}")
-        return None, error_msg
-
-# --- Simulate Status Check ---
-def check_job_status(job_id: str) -> tuple[str, str | None, str | None]:
-    """
-    Simulates checking the status of a job via the backend API.
-    In reality, this would use 'requests.get'.
-
-    Args:
-        job_id (str): The ID of the job to check.
-
-    Returns:
-        tuple[str, str | None, str | None]: (status, result_url, error_message)
-        status: 'PENDING', 'PROCESSING', 'COMPLETED', 'FAILED'
-    """
-    status_url = STATUS_API_URL.format(job_id=job_id)
-    print(f"SIMULATING API STATUS CHECK for {job_id} at {status_url}")
-    time.sleep(0.5) # Simulate network latency
-
-    # Simulate different states based on time or random chance
-    # This is highly simplified; a real backend manages state.
-    if job_id not in st.session_state.job_progress:
-        st.session_state.job_progress[job_id] = 0 # Start progress counter
+import base64 # For potential image encoding if needed
 
-    st.session_state.job_progress[job_id] += random.uniform(0.1, 0.3) # Increment progress
+# ─── 1. Configuration & Secrets ─────────────────────────────────────────────
+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
 
-    if st.session_state.job_progress[job_id] < 0.2:
-        status = "PENDING"
-    elif st.session_state.job_progress[job_id] < 0.8:
-        status = "PROCESSING"
-    elif st.session_state.job_progress[job_id] < 1.0:
-         # Simulate occasional failure during processing
-        if random.random() < 0.05: # 5% chance of failure during processing
-             print(f"API Status Simulation: Job {job_id} FAILED during processing.")
-             st.session_state.job_progress[job_id] = 99 # Mark as failed state
-             return "FAILED", None, "Simulated AI Processing Error."
-        else:
-            status = "PROCESSING" # Still processing
-    else:
-        status = "COMPLETED"
+st.set_page_config(page_title="ArchSketch AI [Advanced]", layout="wide", page_icon="🏗️")
 
-    if status == "COMPLETED":
-        # Simulate getting a result URL (e.g., to a generated image in cloud storage)
-        # Use a *real placeholder path* accessible by the Streamlit app for the demo
-        result_placeholder = "assets/staged_result_placeholder.png" # Make sure this exists!
-        print(f"API Status Simulation: Job {job_id} COMPLETED. Result URL (simulated): {result_placeholder}")
-        return "COMPLETED", result_placeholder, None
-    elif status == "FAILED":
-         print(f"API Status Simulation: Job {job_id} FAILED.")
-         # Error message might have been set earlier
-         error = st.session_state.job_errors.get(job_id, "Unknown processing error.")
-         return "FAILED", None, error
-    else:
-        print(f"API Status Simulation: Job {job_id} is {status}.")
-        return status, None, None
-
-# --- Simulate Fetching Result Image ---
-def fetch_result_image(image_path_or_url: str) -> Image.Image | None:
-    """
-    Simulates fetching the result image from a URL or path.
-    In reality, this would use requests.get for a URL.
-    For this demo, we just load from the local placeholder path.
-    """
-    print(f"SIMULATING Fetching result image from: {image_path_or_url}")
-    if os.path.exists(image_path_or_url):
-        try:
-            img = Image.open(image_path_or_url).convert("RGB")
-            return img
-        except UnidentifiedImageError:
-            print(f"ERROR: Placeholder at '{image_path_or_url}' is not a valid image.")
-            return None
-        except Exception as e:
-            print(f"ERROR: Failed to load placeholder '{image_path_or_url}': {e}")
-            return None
-    else:
-        print(f"ERROR: Placeholder image not found at '{image_path_or_url}'.")
-        return None
+# --- 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
 
-# ------------------------------------------------------------------------------
-# Authentication Simulation
-# ------------------------------------------------------------------------------
-def show_login_form():
-    st.warning("Please log in to use the Visualization Suite.")
-    with st.form("login_form"):
-        username = st.text_input("Username")
-        password = st.text_input("Password", type="password")
-        submitted = st.form_submit_button("Login")
-        if submitted:
-            # --- !!! WARNING: NEVER use hardcoded passwords in production !!! ---
-            # --- This is purely for demonstration. Use secure auth libraries ---
-            # --- like streamlit-authenticator or integrate with OAuth/etc. ---
-            if username == "admin" and password == "password123":
-                st.session_state.logged_in = True
-                st.session_state.username = username
-                st.success("Login successful!")
-                time.sleep(1) # Give user time to see success message
-                st.rerun() # Rerun to show the main app
-            else:
-                st.error("Invalid username or password.")
+# ─── 2. State Initialization & Authentication ───────────────────────────────
 
-# ------------------------------------------------------------------------------
-# Initialize Session State (More Robust)
-# ------------------------------------------------------------------------------
 def initialize_state():
+    """Initializes all necessary session state variables."""
     defaults = {
         'logged_in': False,
         'username': None,
-        'input_data_bytes': None, # Store raw bytes
-        'input_image_preview': None, # Store PIL for display if image
-        'input_type': None,
-        'uploaded_filename': None,
         'current_job_id': None,
-        'job_status': None, # PENDING, PROCESSING, COMPLETED, FAILED
-        'job_progress': {}, # Dict to track progress simulation per job_id
-        'job_errors': {}, # Dict to store errors per job_id
-        'ai_result_image': None,
-        'last_run_params': {}
+        '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:
@@ -183,259 +53,464 @@ def initialize_state():
 
 initialize_state()
 
-# ------------------------------------------------------------------------------
-# Main Application Logic
-# ------------------------------------------------------------------------------
+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()
-    st.stop() # Stop execution if not logged in
+    st.stop()
 
-# --- Main App UI (if logged in) ---
-st.title("🚀 AI Real Estate Visualization Suite [PRO]")
-st.caption(f"Welcome, {st.session_state.username}! Advanced AI tools at your fingertips.")
-st.markdown("---")
+# ─── 3. Simulated Backend Interaction Functions ───────────────────────────────
 
-# --- Sidebar ---
+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
+
+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}")
+
+    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}")
+    except Exception as e:
+        print(f"ERROR loading result from {data_path}: {e}")
+        raise
+
+# ─── 4. Sidebar UI ───────────────────────────────────────────────────────────
 with st.sidebar:
-    st.header("⚙️ Input & Configuration")
+    st.header(f"🏗️ ArchSketch AI")
     st.caption(f"User: {st.session_state.username}")
-    if st.button("Logout"):
-        for key in list(st.session_state.keys()): # Clear state on logout
+    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() # Re-init default state
+        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."
+        )
+        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."
+)
 
-    # Prevent changing input while a job is running
-    input_disabled = st.session_state.job_status in ["PENDING", "PROCESSING"]
-
-    input_file = st.file_uploader(
-        "1. Upload File",
-        type=["png", "jpg", "jpeg", "webp", "dxf", "dwg", "obj", "fbx"], # Add more as needed
-        key="file_uploader",
-        accept_multiple_files=False,
-        help="Upload Room Image, Floor Plan (DXF - simulated), or 3D Model (OBJ - simulated).",
-        disabled=input_disabled
-    )
-
-    # --- Input Processing ---
-    if input_file is not None:
-        if input_file.name != st.session_state.uploaded_filename:
-            st.info(f"Processing '{input_file.name}'...")
-            file_ext = os.path.splitext(input_file.name)[1].lower()
-            file_bytes = input_file.getvalue()
-            processed = False
-            input_image_preview = None
+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,
+    }
 
-            try:
-                if file_ext in ['.png', '.jpg', '.jpeg', '.webp']:
-                    image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
-                    max_size = (1024, 1024)
-                    image.thumbnail(max_size, Image.Resampling.LANCZOS)
-                    input_image_preview = image # Store PIL for preview
-                    input_type = 'image'
-                    processed = True
-                elif file_ext in ['.dxf', '.dwg']:
-                     input_type = 'floorplan'
-                     processed = True
-                     st.warning("Floor plan processing is simulated.")
-                elif file_ext in ['.obj', '.fbx']:
-                     input_type = '3dmodel'
-                     processed = True
-                     st.warning("3D model processing is simulated.")
-                else:
-                    st.error(f"Unsupported file type: {file_ext}")
-
-                if processed:
-                    st.session_state.input_data_bytes = file_bytes # Store bytes for API
-                    st.session_state.input_image_preview = input_image_preview
-                    st.session_state.input_type = input_type
-                    st.session_state.uploaded_filename = input_file.name
-                    # Reset job state
-                    st.session_state.current_job_id = None
-                    st.session_state.job_status = None
-                    st.session_state.ai_result_image = None
-                    st.session_state.last_run_params = {}
-                    st.success(f"File '{input_file.name}' ({input_type}) loaded.")
-                    st.rerun() # Rerun to reflect loaded state
-                else:
-                    st.session_state.input_data_bytes = None
-                    st.session_state.input_image_preview = None
-                    st.session_state.input_type = None
-                    st.session_state.uploaded_filename = None
-
-            except UnidentifiedImageError:
-                 st.error("Uploaded file is not a valid image or is corrupted.")
-            except Exception as e:
-                 st.error(f"Error processing file: {e}")
-                 st.session_state.input_data_bytes = None
-                 st.session_state.input_image_preview = None
-                 st.session_state.input_type = None
-                 st.session_state.uploaded_filename = None
-
-
-    # --- Configuration Options ---
-    if st.session_state.input_type is not None:
-        st.markdown("---")
-        st.subheader("2. Visualization Mode")
-        output_mode = st.selectbox(
-            "Select Mode:",
-            options=['Virtual Staging', 'Renovation Preview', 'Material Swap', 'Layout Variation'],
-            key='output_mode_select',
-            disabled=input_disabled
-        ).lower().replace(' ', '_')
-
-        st.markdown("---")
-        st.subheader("3. Scene Parameters")
-        room_type = st.selectbox("Room Type:", ["Living Room", "Bedroom", "Kitchen", "Dining Room", "Office", "Bathroom", "Other"], key="room_select", disabled=input_disabled)
-        style = st.selectbox("Primary Style:", ["Modern", "Contemporary", "Minimalist", "Scandinavian", "Industrial", "Traditional", "Coastal", "Farmhouse", "Bohemian"], key="style_select", disabled=input_disabled)
-
-        # Advanced Controls
-        with st.expander("✨ Advanced Controls", expanded=False):
-            furniture_prefs = st.text_input("Furniture Preferences", placeholder="e.g., 'Large velvet green sofa'", key="furniture_input", disabled=input_disabled)
-            wall_color = st.color_picker("Wall Color (Approx.)", value="#FFFFFF", key="wall_color_picker", disabled=input_disabled)
-            floor_type = st.selectbox("Floor Type", ["(Auto)", "Oak Wood", "Dark Wood", "Light Wood", "Carpet (Neutral)", "Concrete", "Tile (Light)", "Tile (Dark)"], key="floor_select", disabled=input_disabled)
-            lighting_time = st.slider("Time of Day", 0.0, 1.0, 0.5, 0.1, "%.1f", key="lighting_slider", help="0.0=Dawn, 0.5=Midday, 1.0=Dusk", disabled=input_disabled)
-            camera_angle = st.selectbox("Camera Angle", ["Eye-Level", "High Angle", "Low Angle", "Wide (Simulated)"], key="camera_select", disabled=input_disabled)
-            remove_objects = st.checkbox("Attempt Remove Existing Objects", value=False, key="remove_obj_check", disabled=input_disabled)
-            renovation_instructions = ""
-            if output_mode == 'renovation_preview':
-                renovation_instructions = st.text_input("Renovation Instructions", placeholder="e.g., 'Remove center wall'", key="renovation_input", disabled=input_disabled)
-
-        st.markdown("---")
-        st.subheader("4. Generate Visualization")
-
-        # Disable button if no input or job already running
-        disable_generate = st.session_state.input_type is None or st.session_state.job_status in ["PENDING", "PROCESSING"]
-        if st.button("🚀 Submit Visualization Job", key="generate_button", use_container_width=True, disabled=disable_generate):
-            # Prepare payload for the backend
-            # In production, you'd likely upload the file to S3 first and send a URL/key
-            # For image data, you might send base64 encoded string or upload separately
-            payload = {
-                "user_id": st.session_state.username, # Identify the user
-                # "input_reference": "s3://bucket/user_uploads/filename.jpg", # Example
-                "input_filename": st.session_state.uploaded_filename, # For info
-                "input_type": st.session_state.input_type,
-                "output_mode": output_mode,
-                "room_type": room_type,
-                "style": style,
-                "furniture_prefs": furniture_prefs,
-                "materials": {'wall_color': wall_color, 'floor_type': floor_type},
-                "lighting_time": lighting_time,
-                "camera_angle": camera_angle,
-                "remove_objects": remove_objects,
-                "renovation_instructions": renovation_instructions,
-            }
-            # Attach input data - simplistic approach for demo, REAL API would handle uploads better
-            # Never send large files directly in JSON payload in production!
-            # payload['input_data_b64'] = base64.b64encode(st.session_state.input_data_bytes).decode()
-
-            job_id, error = submit_visualization_job(payload)
-
-            if job_id:
-                st.session_state.current_job_id = job_id
-                st.session_state.job_status = "PENDING"
-                st.session_state.ai_result_image = None # Clear previous result
-                st.session_state.last_run_params = payload # Store params for display
-                st.session_state.job_progress = {job_id: 0} # Reset progress for new job
-                st.session_state.job_errors = {} # Clear old errors
-                st.success(f"Job submitted successfully! Job ID: {job_id}")
-                st.rerun() # Start polling
-            else:
-                st.error(f"Failed to submit job: {error}")
-                st.session_state.current_job_id = None
-                st.session_state.job_status = "FAILED" # Mark status
+    # 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.info("⬆️ Upload a file to begin.")
-
-# --- Main Display Area ---
-col1, col2 = st.columns(2)
-
-with col1:
-    st.subheader("Input")
-    if st.session_state.input_type is not None:
-        if st.session_state.input_type == 'image' and st.session_state.input_image_preview:
-            st.image(st.session_state.input_image_preview, caption=f"Input: {st.session_state.uploaded_filename}", use_column_width=True)
-        elif st.session_state.input_type != 'image':
-             # Display info for non-image types
-            st.info(f"{st.session_state.input_type.capitalize()} File:\n**{st.session_state.uploaded_filename}**")
-            st.caption("(Preview not available for non-image inputs in this demo)")
-        else:
-             st.warning("Input loaded, but preview unavailable.")
-    else:
-        st.markdown("<div style='height: 400px; border: 2px dashed #ccc; ...'>Upload Input File</div>", unsafe_allow_html=True)
+        st.error(f"Job submission failed: {error}")
+        st.session_state.job_status = "FAILED"
 
 
-with col2:
-    st.subheader("AI Visualization Result")
+# --- Status & Result Display Area ---
+st.markdown("---")
+st.subheader("📊 Job Status & Result")
+
+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
 
-    job_id = st.session_state.current_job_id
-    status = st.session_state.job_status
 
-    if job_id:
-        # --- Job Status Display ---
-        if status == "PENDING":
-            st.info(f"Job Status: Pending... (ID: {job_id})")
-            # Trigger periodic rerun for polling
-            time.sleep(5) # Poll every 5 seconds (adjust as needed)
-            st.rerun()
-        elif status == "PROCESSING":
-            progress_value = st.session_state.job_progress.get(job_id, 0)
-            st.progress(min(progress_value, 1.0), text=f"Job Status: Processing... ({int(min(progress_value, 1.0)*100)}%) (ID: {job_id})")
-            # Trigger periodic rerun for polling
-            time.sleep(3) # Poll faster while processing
-            st.rerun()
-        elif status == "COMPLETED":
-            st.success(f"Job Status: Completed! (ID: {job_id})")
-            if st.session_state.ai_result_image:
-                 run_mode_display = st.session_state.last_run_params.get('output_mode', 'N/A').replace('_', ' ').title()
-                 st.image(st.session_state.ai_result_image, caption=f"Result ({run_mode_display})", use_column_width=True)
-                 # Download Button logic (similar to before)
-                 # ... (Add download button code here) ...
-            else:
-                 st.error("Job completed, but failed to load the result image.")
-            with st.expander("View Parameters Used", expanded=False):
-                 # Display JSON but remove potentially large data fields first
-                 params_to_display = {k:v for k,v in st.session_state.last_run_params.items()} # if k != 'input_data_b64'}
-                 st.json(params_to_display, expanded=True)
-
-        elif status == "FAILED":
-            error_msg = st.session_state.job_errors.get(job_id, "An unknown error occurred.")
-            st.error(f"Job Status: Failed! (ID: {job_id})\nError: {error_msg}")
-            with st.expander("View Parameters Used", expanded=False):
-                 params_to_display = {k:v for k,v in st.session_state.last_run_params.items()} # if k != 'input_data_b64'}
-                 st.json(params_to_display, expanded=True)
-
-        # --- Logic to update status (if job is active) ---
-        if status in ["PENDING", "PROCESSING"]:
-             new_status, result_url, error = check_job_status(job_id)
-             st.session_state.job_status = new_status
-             if error:
-                 st.session_state.job_errors[job_id] = error
-             if new_status == "COMPLETED" and result_url:
-                 # Fetch the result image only when completed
-                 st.session_state.ai_result_image = fetch_result_image(result_url)
-                 if st.session_state.ai_result_image is None:
-                     st.session_state.job_status = "FAILED" # Mark as failed if image fetch fails
-                     st.session_state.job_errors[job_id] = "Failed to retrieve/load result image."
-             elif new_status == "FAILED" and not error: # If status is failed but no error stored yet
-                  st.session_state.job_errors[job_id] = "Job failed for an unknown reason."
-
-             # Rerun immediately if status changed significantly to update UI
-             if status != new_status and new_status in ["COMPLETED", "FAILED"]:
-                 st.rerun()
-
-    else: # No active job
-         st.markdown("<div style='height: 400px; border: 2px dashed #ccc; ...'>Result will appear here</div>", unsafe_allow_html=True)
+# --- 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.")
 
 
+# ─── Footer & Disclaimer ─────────────────────────────────────────────────────
 st.markdown("---")
 st.warning("""
-    **Disclaimer:** This is a **conceptual blueprint** for a production front-end.
-    User authentication is **not secure**. Backend API calls, job handling, status polling, and AI processing are **simulated**.
-    Building the actual backend infrastructure and state-of-the-art AI models requires significant resources and expertise.
+    **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.
 """)