Update app.py

app.py

@@ -1,180 +1,224 @@
 import streamlit as st
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image, UnidentifiedImageError
 import io
 import time
 import os
-import random # For simulating variations
+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",
+    page_title="AI Real Estate Visualization Suite [PRO]",
     layout="wide",
-    page_icon="💎",
+    page_icon="🚀",
     initial_sidebar_state="expanded"
 )
 
 # ------------------------------------------------------------------------------
-# Advanced Placeholder for AI Model Call
+# Simulated Backend API Interaction
 # ------------------------------------------------------------------------------
-# This function now accepts many more parameters to simulate advanced control.
-# In reality, this would likely involve multiple specialized AI models or a very
-# complex multi-modal model API.
-
-def run_advanced_ai_model(
-    input_data, # Could be PIL Image, or bytes for other formats
-    input_type: str, # 'image', 'floorplan', '3dmodel'
-    output_mode: str, # 'staging', 'renovation', 'material_swap', 'layout_variation'
-    room_type: str,
-    style: str,
-    furniture_prefs: str = "",
-    materials: dict = {}, # e.g., {'wall_color': '#FFFFFF', 'floor_type': 'Oak Wood'}
-    lighting_time: float = 0.5, # 0.0 (dawn) to 1.0 (dusk)
-    camera_angle: str = "Eye-Level",
-    remove_objects: bool = False,
-    renovation_instructions: str = "" # e.g., "Remove wall between kitchen and living room"
-    ) -> Image.Image:
+# 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]:
     """
-    Advanced placeholder simulating a sophisticated AI visualization model.
+    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(f"Simulating Advanced AI: Mode='{output_mode}', Style='{style}'...")
-    print(f"--- ADVANCED AI SIMULATION ---")
-    print(f"Input Type: {input_type}")
-    print(f"Output Mode: {output_mode}")
-    print(f"Room Type: {room_type}")
-    print(f"Style: {style}")
-    print(f"Furniture Prefs: {furniture_prefs if furniture_prefs else 'N/A'}")
-    print(f"Materials: {materials}")
-    print(f"Lighting Time (0-1): {lighting_time:.2f}")
-    print(f"Camera Angle: {camera_angle}")
-    print(f"Remove Existing Objects: {remove_objects}")
-    print(f"Renovation Instructions: {renovation_instructions if renovation_instructions else 'N/A'}")
-    print(f"-----------------------------")
-
-    # Simulate processing time based on complexity
-    processing_time = 3 + len(furniture_prefs) * 0.01 + len(materials) * 0.5
-    if output_mode == 'renovation': processing_time += 2
-    time.sleep(min(processing_time, 8)) # Cap simulation time
-
-    # --- Simulation Logic ---
-    # Determine the placeholder image based on mode
-    base_placeholder = "assets/staged_result_placeholder.png"
-    if output_mode == 'renovation':
-        base_placeholder = "assets/renovation_placeholder.png" # Need another dummy image
-    elif output_mode == 'material_swap':
-        base_placeholder = "assets/material_swap_placeholder.png" # Need another dummy image
-    elif output_mode == 'layout_variation':
-         base_placeholder = f"assets/layout_{random.randint(1,2)}_placeholder.png" # Need layout_1/2 dummies
-
-    if isinstance(input_data, Image.Image):
-        base_image = input_data # Use original if input was image
+    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:
-        base_image = None # Cannot use non-image input directly for overlay
+        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
 
-    # Try loading the specific placeholder
-    staged_image = None
-    if os.path.exists(base_placeholder):
-        try:
-            staged_image = Image.open(base_placeholder).convert("RGB")
-            print(f"SIMULATING AI: Loaded placeholder: {base_placeholder}")
-        except Exception as e:
-            print(f"Error loading placeholder '{base_placeholder}': {e}")
-            staged_image = None
-
-    # Fallback if placeholder failed or input wasn't an image for overlay
-    if staged_image is None:
-        if base_image:
-             staged_image = base_image.copy()
-             print("SIMULATING AI: Placeholder failed, using copy of original image.")
-        else:
-             # Create a generic image if no input image and placeholder failed
-             staged_image = Image.new('RGB', (800, 600), color = (200, 200, 200))
-             print("SIMULATING AI: Placeholder failed and no base image, using blank canvas.")
-
-    # Add text overlay summarizing the simulation
-    draw = ImageDraw.Draw(staged_image)
-    try:
-        font_path = "arial.ttf" # Check system fonts or provide path
-        font_size = 20
-        font = ImageFont.truetype(font_path, font_size)
-    except IOError:
-        font = ImageFont.load_default()
-
-    text_lines = [
-        f"AI Sim Result: {output_mode.replace('_', ' ').title()}",
-        f"Style: {style}, Room: {room_type}",
-        f"Lighting: {lighting_time:.2f}, Camera: {camera_angle}",
-    ]
-    if furniture_prefs: text_lines.append(f"Prefs: {furniture_prefs[:30]}...")
-    if materials.get('wall_color'): text_lines.append(f"Wall: {materials['wall_color']}")
-    if materials.get('floor_type'): text_lines.append(f"Floor: {materials['floor_type']}")
-    if remove_objects: text_lines.append(f"[Removed Objects]")
-    if renovation_instructions: text_lines.append(f"Reno: {renovation_instructions[:30]}...")
-
-    y_pos = 10
-    for line in text_lines:
-        try:
-             # Simple background box
-             text_bbox = draw.textbbox((10, y_pos), line, font=font)
-             # Add padding to bbox
-             padded_bbox = (text_bbox[0]-2, text_bbox[1]-2, text_bbox[2]+2, text_bbox[3]+2)
-             draw.rectangle(padded_bbox, fill="rgba(0, 0, 0, 0.6)")
-             draw.text((10, y_pos), line, fill="white", font=font)
-             y_pos += font_size + 4 # Move down for next line
-        except Exception as draw_err:
-            print(f"Error drawing text overlay: {draw_err}") # Log error but continue
+# --- 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'.
 
-    print("SIMULATING AI: Processing complete.")
-    st.success(f"Advanced AI Simulation for '{output_mode}' Complete!")
-    return staged_image
+    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
+
+    st.session_state.job_progress[job_id] += random.uniform(0.1, 0.3) # Increment progress
+
+    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"
+
+    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
 
 # ------------------------------------------------------------------------------
-# Initialize Session State
+# Authentication Simulation
 # ------------------------------------------------------------------------------
-if 'input_data' not in st.session_state:
-    st.session_state.input_data = None # Can hold PIL Image or bytes
-if 'input_type' not in st.session_state:
-    st.session_state.input_type = None # 'image', 'floorplan', '3dmodel'
-if 'ai_result_image' not in st.session_state:
-    st.session_state.ai_result_image = None
-if 'uploaded_filename' not in st.session_state:
-    st.session_state.uploaded_filename = None
-if 'last_run_params' not in st.session_state: # Store params used for the run
-    st.session_state.last_run_params = {}
+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.")
 
 # ------------------------------------------------------------------------------
-# Helper Functions
+# Initialize Session State (More Robust)
 # ------------------------------------------------------------------------------
-def display_input_placeholder(input_type, filename):
-    """Shows appropriate info for non-image inputs."""
-    if input_type == 'floorplan':
-        st.info(f"Floor Plan File Loaded: **{filename}**\n\n*(Rendering from floor plans is simulated)*")
-        # In a real app, you might use a library like `ezdxf` to show a basic plot
-    elif input_type == '3dmodel':
-        st.info(f"3D Model File Loaded: **{filename}**\n\n*(Rendering from 3D models is simulated)*")
-        # In a real app, you might use `pyvista` or `trimesh` for basic viewing
-    else:
-        st.warning("Unsupported input type for display.")
+def initialize_state():
+    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': {}
+    }
+    for key, value in defaults.items():
+        if key not in st.session_state:
+            st.session_state[key] = value
+
+initialize_state()
 
 # ------------------------------------------------------------------------------
-# Main Application UI
+# Main Application Logic
 # ------------------------------------------------------------------------------
-st.title("💎 Advanced AI Real Estate Visualization Suite")
-st.caption("Simulating cutting-edge AI for staging, renovation previews, material swaps, and more.")
+
+# --- Authentication Gate ---
+if not st.session_state.logged_in:
+    show_login_form()
+    st.stop() # Stop execution if not logged in
+
+# --- 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("---")
 
 # --- Sidebar ---
 with st.sidebar:
     st.header("⚙️ Input & Configuration")
+    st.caption(f"User: {st.session_state.username}")
+    if st.button("Logout"):
+        for key in list(st.session_state.keys()): # Clear state on logout
+            del st.session_state[key]
+        initialize_state() # Re-init default state
+        st.rerun()
+
+    st.markdown("---")
+
+    # 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"], # Accept more types
+        type=["png", "jpg", "jpeg", "webp", "dxf", "dwg", "obj", "fbx"], # Add more as needed
         key="file_uploader",
         accept_multiple_files=False,
-        help="Upload Room Image (JPG, PNG), Floor Plan (DXF - simulated), or 3D Model (OBJ - simulated)."
+        help="Upload Room Image, Floor Plan (DXF - simulated), or 3D Model (OBJ - simulated).",
+        disabled=input_disabled
     )
 
     # --- Input Processing ---
@@ -184,203 +228,214 @@ with st.sidebar:
             file_ext = os.path.splitext(input_file.name)[1].lower()
             file_bytes = input_file.getvalue()
             processed = False
+            input_image_preview = None
 
-            # Image Input
-            if file_ext in ['.png', '.jpg', '.jpeg', '.webp']:
-                try:
+            try:
+                if file_ext in ['.png', '.jpg', '.jpeg', '.webp']:
                     image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
-                    max_size = (1024, 1024) # Resize for consistency
+                    max_size = (1024, 1024)
                     image.thumbnail(max_size, Image.Resampling.LANCZOS)
-                    st.session_state.input_data = image
-                    st.session_state.input_type = 'image'
+                    input_image_preview = image # Store PIL for preview
+                    input_type = 'image'
                     processed = True
-                except Exception as e: st.error(f"Error loading image: {e}")
-            # Floor Plan Input (Simulated)
-            elif file_ext in ['.dxf', '.dwg']:
-                 st.session_state.input_data = file_bytes # Store bytes
-                 st.session_state.input_type = 'floorplan'
-                 processed = True
-                 st.warning("Floor plan processing is simulated.")
-            # 3D Model Input (Simulated)
-            elif file_ext in ['.obj', '.fbx']:
-                 st.session_state.input_data = file_bytes # Store bytes
-                 st.session_state.input_type = '3dmodel'
-                 processed = True
-                 st.warning("3D model processing is simulated.")
-            else:
-                st.error(f"Unsupported file type: {file_ext}")
-
-            # Update state if processed successfully
-            if processed:
-                st.session_state.ai_result_image = None # Reset result
-                st.session_state.uploaded_filename = input_file.name
-                st.session_state.last_run_params = {} # Clear old params
-                st.success(f"File '{input_file.name}' ({st.session_state.input_type}) loaded.")
-            else:
-                 st.session_state.input_data = None
+                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
 
-    elif st.session_state.uploaded_filename: # File was removed by user
-         st.session_state.input_data = None
-         st.session_state.input_type = None
-         st.session_state.ai_result_image = None
-         st.session_state.uploaded_filename = None
-         st.session_state.last_run_params = {}
-
 
-    # --- Configuration Options (Only if input is loaded) ---
-    if st.session_state.input_data is not 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',
-            help="Choose the type of visualization you want."
-        ).lower().replace(' ', '_') # Convert to key format like 'virtual_staging'
+            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"
-        )
-        style = st.selectbox(
-            "Primary Style:",
-            ["Modern", "Contemporary", "Minimalist", "Scandinavian", "Industrial", "Traditional", "Coastal", "Farmhouse", "Bohemian"],
-            key="style_select"
-        )
-
-        # --- Advanced Controls Expander ---
-        with st.expander("✨ Advanced Controls"):
-            furniture_prefs = st.text_input(
-                "Furniture Preferences (Optional)",
-                placeholder="e.g., 'Large velvet green sofa', 'minimalist oak desk'",
-                key="furniture_input",
-                help="Describe specific furniture items or characteristics."
-            )
-
-            st.markdown("**Material Customization:**")
-            wall_color = st.color_picker("Wall Color (Approx.)", value="#FFFFFF", key="wall_color_picker")
-            floor_type = st.selectbox(
-                "Floor Type",
-                ["(Auto)", "Oak Wood", "Dark Wood", "Light Wood", "Carpet (Neutral)", "Concrete", "Tile (Light)", "Tile (Dark)"],
-                key="floor_select"
-            )
-            materials_dict = {'wall_color': wall_color, 'floor_type': floor_type}
-
-            st.markdown("**Lighting & Camera:**")
-            lighting_time = st.slider(
-                "Time of Day (Simulated)",
-                min_value=0.0, max_value=1.0, value=0.5, step=0.1, format="%.1f",
-                key="lighting_slider",
-                help="0.0 = Dawn/Sunrise, 0.5 = Midday, 1.0 = Dusk/Sunset"
-            )
-            camera_angle = st.selectbox(
-                "Camera Angle", ["Eye-Level", "High Angle", "Low Angle", "Wide Angle (Simulated)"],
-                key="camera_select"
-            )
-
-            st.markdown("**AI Assist Features:**")
-            remove_objects = st.checkbox(
-                "Attempt to Remove Existing Objects/Furniture", value=False, key="remove_obj_check",
-                help="If staging an image that isn't empty, AI will try to remove existing items first (Simulated)."
-            )
-
+        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': # Show only in renovation mode
-                renovation_instructions = st.text_input(
-                    "Renovation Instructions (Simulated)",
-                    placeholder="e.g., 'Remove center wall', 'add window on left wall'",
-                    key="renovation_input"
-                )
+            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")
-        if st.button("🚀 Generate Advanced Visualization", key="generate_button", use_container_width=True):
-            with st.spinner("Engaging Advanced AI Simulation... ✨"):
-                try:
-                    run_params = {
-                        "input_data": st.session_state.input_data,
-                        "input_type": st.session_state.input_type,
-                        "output_mode": output_mode,
-                        "room_type": room_type,
-                        "style": style,
-                        "furniture_prefs": furniture_prefs,
-                        "materials": materials_dict,
-                        "lighting_time": lighting_time,
-                        "camera_angle": camera_angle,
-                        "remove_objects": remove_objects,
-                        "renovation_instructions": renovation_instructions,
-                    }
-                    st.session_state.last_run_params = run_params # Store params
-                    st.session_state.ai_result_image = run_advanced_ai_model(**run_params)
-
-                except Exception as e:
-                    st.error(f"An error occurred during AI simulation: {e}")
-                    st.session_state.ai_result_image = None
+
+        # 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
 
     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_data is not None:
-        if st.session_state.input_type == 'image':
-            st.image(st.session_state.input_data, caption=f"Original: {st.session_state.uploaded_filename}", use_column_width=True)
+    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:
-            # Display info for non-image types
-            display_input_placeholder(st.session_state.input_type, st.session_state.uploaded_filename)
+             st.warning("Input loaded, but preview unavailable.")
     else:
-        st.markdown("<div style='height: 400px; border: 2px dashed #ccc; display: flex; align-items: center; justify-content: center; text-align: center; color: #aaa; font-style: italic;'>Upload Input File</div>", unsafe_allow_html=True)
+        st.markdown("<div style='height: 400px; border: 2px dashed #ccc; ...'>Upload Input File</div>", unsafe_allow_html=True)
+
 
 with col2:
     st.subheader("AI Visualization Result")
-    if st.session_state.ai_result_image is not None:
-        # Display the resulting 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} Simulation)", use_column_width=True)
-
-        # Display parameters used for this run
-        with st.expander("View Parameters Used for This Result", expanded=False):
-             st.json(st.session_state.last_run_params, expanded=True) # Cannot directly display image data in json
-
 
-        # Download Button
-        try:
-            buf = io.BytesIO()
-            img_to_save = st.session_state.ai_result_image
-            save_format = 'PNG' # Use PNG to preserve potential overlay quality
-            img_to_save.save(buf, format=save_format)
-            byte_im = buf.getvalue()
-
-            # Generate filename based on input and mode
-            base_name = os.path.splitext(st.session_state.uploaded_filename or 'ai_result')[0]
-            mode_name = st.session_state.last_run_params.get('output_mode', 'result')
-            download_filename = f"{base_name}_{mode_name}.{save_format.lower()}"
-
-            st.download_button(
-                label=f"Download Result ({save_format})",
-                data=byte_im,
-                file_name=download_filename,
-                mime=f"image/{save_format.lower()}"
-            )
-        except Exception as e:
-            st.error(f"Could not prepare image for download: {e}")
-
-    else:
-        st.markdown("<div style='height: 400px; border: 2px dashed #ccc; display: flex; align-items: center; justify-content: center; text-align: center; color: #aaa; font-style: italic;'>Result will appear here</div>", unsafe_allow_html=True)
+    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)
 
 
 st.markdown("---")
 st.warning("""
-    **Disclaimer:** This is an advanced **demonstration** application.
-    All AI functionalities (staging, renovation, material swaps, floor plan/3D model processing, object removal etc.) are **simulated** using placeholder logic and images.
-    The purpose is to showcase a potential user interface and feature set for a sophisticated AI visualization tool.
+    **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.
 """)