app.py

# 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 json
import uuid
import time
import random
import base64 # For potential image encoding if needed

# ─── 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

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()
    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

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(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."
        )
        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."
)

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"


# --- 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


# --- 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 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.
""")