Update app.py

app.py

@@ -3,339 +3,1135 @@ import io
 import os
 import pandas as pd
 from docx import Document
-from io import BytesIO
-import dash
-import dash_bootstrap_components as dbc
-from dash import html, dcc, Input, Output, State, callback_context
+from io import BytesIO, StringIO
+import dash # Version 3.0.3
+import dash_bootstrap_components as dbc # Version 2.0.2
+from dash import html, dcc, Input, Output, State, callback_context, PreventUpdate, ALL
 import google.generativeai as genai
 from docx.shared import Pt
 from docx.enum.style import WD_STYLE_TYPE
 from PyPDF2 import PdfReader
-from io import StringIO
+import logging
+import uuid # For unique IDs if needed with pattern matching
 
-# Initialize Dash app
-app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
+# --- Logging Configuration ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Configure Gemini AI
-genai.configure(api_key=os.environ["GEMINI_API_KEY"])
-model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
+# --- Initialize Dash app ---
+# Using Bootstrap for layout and styling. Added meta tags for responsiveness.
+# dash==3.0.3
+# dash-bootstrap-components==2.0.2
+app = dash.Dash(__name__,
+                external_stylesheets=[dbc.themes.BOOTSTRAP],
+                suppress_callback_exceptions=True, # Needed because controls are dynamically added
+                meta_tags=[{"name": "viewport", "content": "width=device-width, initial-scale=1"}])
+server = app.server # Expose server for Gunicorn
 
-# Global variables
-uploaded_files = {}
-current_document = None
-document_type = None
-shredded_document = None
-pink_review_document = None
+# --- Configure Gemini AI ---
+# IMPORTANT: Set the GEMINI_API_KEY environment variable before running the app.
+try:
+    # Prefer direct CUDA GPU configuration in app.py - Note: This is not applicable for cloud-based APIs like Gemini. Configuration happens via API key.
+    api_key = os.environ.get("GEMINI_API_KEY")
+    if not api_key:
+        logging.warning("GEMINI_API_KEY environment variable not found. AI features will be disabled.")
+        model = None
+    else:
+        genai.configure(api_key=api_key)
+        # Specify a model compatible with function calling or more advanced generation if needed.
+        # Using 'gemini-pro' as a generally available and capable model.
+        model = genai.GenerativeModel('gemini-pro')
+        logging.info("Gemini AI configured successfully using 'gemini-pro'.")
+except Exception as e:
+    logging.error(f"Error configuring Gemini AI: {e}", exc_info=True)
+    model = None
+
+# --- Global Variables ---
+# Using dictionaries to store session-specific data might be better for multi-user,
+# but for simplicity with current constraints, we use global vars.
+# Consider using dcc.Store for better state management in complex scenarios.
+uploaded_files = {} # {filename: content_text}
+
+# Stores the *results* of generation/review steps
+shredded_document = None # Text content of the shredded PWS/requirements
+pink_review_document = None # Text content of the generated Pink Review
+red_review_document = None # Text content of the generated Red Review
+gold_review_document = None # Text content of the generated Gold Review
+loe_document = None # Text content of the generated LOE
+virtual_board_document = None # Text content of the generated Virtual Board
+
+# Stores the *generated* proposal drafts
+pink_document = None # Text content of the generated Pink Team document
+red_document = None # Text content of the generated Red Team document
+gold_document = None # Text content of the generated Gold Team document
 
-# Document types and their descriptions
+# Store uploaded content specifically for review inputs
+uploaded_pink_content = None
+uploaded_red_content = None
+uploaded_gold_content = None
+
+# Store the currently displayed document and its type for download/chat
+current_display_document = None
+current_display_type = None
+
+# --- Document Types ---
+# Descriptions adjusted slightly for clarity
 document_types = {
-    "Shred": "Generate a requirements spreadsheet of the Project Work Statement (PWS) identified by action words like shall, will, perform etc. by pws section, requirement.  Do not write as if you're responding to the proposal.  Its a spreadsheet to distill the requirements, not microhealth's approach",
-    "Pink": "Create a Pink Team document based on the PWS outline.  Your goal is to be compliant and compelling.",
-    "Pink Review": "Evaluate compliance of the Pink Team document against the requirements and output a spreadsheet of non compliant findings by pws number, the goal of that pws section, what made it non compliant and your recommendations for recovery",
-    "Red": "Produce a Red Team document based on the Pink Review by pws sections. Your goal is to be compliant and compelling by recovering all the findings in Pink Review",
-    "Red Review": "Evaluate compliance of the Red Team document against the requirements and output a spreadsheet of non compliant findings by pws number, the goal of that pws section, what made it non compliant and your recommendations for recovery",
-    "Gold": "Create a Pink Team document based on the PWS response by pws sections. Your goal is to be compliant and compelling by recovering all the findings in Red Review",
-    "Gold Review": "Perform a final compliance review against the requirements and output a spreadsheet of non compliant findings by pws number, the goal of that pws section, what made it non compliant and your recommendations for recovery",
-    "Virtual Board": "Based on the requirements and in particular the evaulation criteria, you will evaluate the proposal as if you were a contracting office and provide section by section evaluation as unsatisfactory, satisfactory, good, very good, excellent and why in a spreadsheet",
-    "LOE": "Generate a Level of Effort (LOE) breakdown as a spreadsheet"
+    "Shred": "Generate a requirements spreadsheet from the PWS/Source Docs, identifying action words (shall, will, perform, etc.) by section.",
+    "Pink": "Create a compliant and compelling Pink Team proposal draft based on the Shredded requirements.",
+    "Pink Review": "Evaluate a Pink Team draft against Shredded requirements. Output findings (compliance, gaps, recommendations) in a spreadsheet.",
+    "Red": "Create a Red Team proposal draft, addressing feedback from the Pink Review and enhancing compliance/compellingness.",
+    "Red Review": "Evaluate a Red Team draft against Shredded requirements and Pink Review findings. Output findings in a spreadsheet.",
+    "Gold": "Create a Gold Team proposal draft, addressing feedback from the Red Review for final compliance and polish.",
+    "Gold Review": "Perform a final compliance review of the Gold Team draft against Shredded requirements and Red Review findings. Output findings.",
+    "Virtual Board": "Simulate a source selection board evaluation of the final proposal against PWS/Shred requirements and evaluation criteria (Sec L/M). Output evaluation.",
+    "LOE": "Generate a Level of Effort (LOE) estimate spreadsheet based on the Shredded requirements."
 }
 
-app.layout = dbc.Container([
-    dbc.Row([
-        dbc.Col([
-            html.H4("Proposal Documents", className="mt-3 mb-4"),
-            dcc.Upload(
-                id='upload-document',
-                children=html.Div([
-                    'Drag and Drop or ',
-                    html.A('Select Files')
-                ]),
-                style={
-                    'width': '100%',
-                    'height': '60px',
-                    'lineHeight': '60px',
-                    'borderWidth': '1px',
-                    'borderStyle': 'dashed',
-                    'borderRadius': '5px',
-                    'textAlign': 'center',
-                    'margin': '10px 0'
-                },
-                multiple=True
-            ),
-            html.Div(id='file-list'),
-            html.Hr(),
-            html.Div([
-                dbc.Button(
-                    doc_type,
-                    id=f'btn-{doc_type.lower().replace("_", "-")}',
-                    color="link",
-                    className="mb-2 w-100 text-left custom-button",
-                    style={'overflow': 'hidden', 'text-overflow': 'ellipsis', 'white-space': 'nowrap'}
-                ) for doc_type in document_types.keys()
-            ])
-        ], width=3),
-        dbc.Col([
-            html.Div(id='status-bar', className="alert alert-info", style={'marginBottom': '20px'}),
+# --- Layout Definition ---
+# Using Dash Bootstrap Components for layout and Cards for logical separation.
+# Single form layout functions for modern design.
+app.layout = dbc.Container(fluid=True, className="dbc", children=[
+    # Title Row - Full Width
+    dbc.Row(
+        dbc.Col(html.H1("Proposal AI Assistant", className="text-center my-4", style={'color': '#1C304A'}), width=12)
+    ),
+
+    # Progress Indicator Row (Initially Hidden) - Full Width below title, above columns
+    dbc.Row(
+        dbc.Col(
+            # Blinking triple dot for progress
             dcc.Loading(
                 id="loading-indicator",
-                type="dot",
-                children=[html.Div(id="loading-output")]
+                type="dots", # Changed type to dots as requested
+                children=[html.Div(id="loading-output", style={'height': '10px'})], # Placeholder content
+                overlay_style={"visibility":"hidden", "opacity": 0}, # Make overlay invisible
+                style={'visibility':'hidden', 'height': '30px'}, # Hide initially via style, give some height
+                fullscreen=False, # Keep it contained
+                className="justify-content-center"
             ),
-            html.Div(id='document-preview', className="border p-3 mb-3"),
-            dbc.Button("Download Document", id="btn-download", color="success", className="mt-3"),
-            dcc.Download(id="download-document"),
-            html.Hr(),
-            html.Div(id='pink-review-upload', style={'display': 'none'}, children=[
-                dcc.Upload(
-                    id='upload-pink-review',
-                    children=html.Div(['Drag and Drop or ', html.A('Select Pink Review File')]),
-                    style={
-                        'width': '100%',
-                        'height': '60px',
-                        'lineHeight': '60px',
-                        'borderWidth': '1px',
-                        'borderStyle': 'dashed',
-                        'borderRadius': '5px',
-                        'textAlign': 'center',
-                        'margin': '10px 0'
-                    },
-                    multiple=False
-                ),
-                html.Div(id='pink-review-file-name')
-            ]),
-            dcc.Loading(
-                id="chat-loading",
-                type="dot",
-                children=[
-                    dbc.Input(id="chat-input", type="text", placeholder="Chat with AI to update document...", className="mb-2"),
-                    dbc.Button("Send", id="btn-send-chat", color="primary", className="mb-3"),
-                    html.Div(id="chat-output")
-                ]
-            )
-        ], width=9)
+            width=12,
+            className="text-center mb-3" # Center the dots
+        )
+    ),
+
+    # Main Content Row (Two Columns)
+    dbc.Row([
+        # Left Column (Navigation / Upload) - 30% width, light gray background
+        dbc.Col(
+            dbc.Card(
+                dbc.CardBody([
+                    html.H4("1. Upload Source Documents", className="card-title"),
+                    dcc.Upload(
+                        id='upload-document',
+                        children=html.Div(['Drag and Drop or ', html.A('Select PWS/Source Files (.docx, .pdf)')]),
+                        style={
+                            'width': '100%', 'height': '60px', 'lineHeight': '60px',
+                            'borderWidth': '1px', 'borderStyle': 'dashed', 'borderRadius': '5px',
+                            'textAlign': 'center', 'margin': '10px 0', 'backgroundColor': '#ffffff' # White background for contrast
+                        },
+                        multiple=True # Allow multiple source files
+                    ),
+                    # Use Card for file list for better visual grouping
+                    dbc.Card(
+                        dbc.CardBody(
+                           html.Div(id='file-list', style={'maxHeight': '150px', 'overflowY': 'auto', 'fontSize': '0.9em'})
+                        ), className="mb-3" , style={'backgroundColor': '#ffffff'}
+                    ),
+                    html.Hr(),
+                    html.H4("2. Select Action", className="card-title mt-3"),
+                    # Buttons for actions - Use Card for button group
+                    dbc.Card(
+                        dbc.CardBody([
+                            # Use primary button style defined in CSS request (implicitly via dbc class)
+                            *[dbc.Button(
+                                doc_type,
+                                id={'type': 'action-button', 'index': doc_type}, # Use pattern-matching ID
+                                color="primary", # Primary style
+                                className="mb-2 w-100 d-block", # d-block for full width buttons
+                                style={'textAlign': 'left', 'whiteSpace': 'normal', 'height': 'auto', 'wordWrap': 'break-word'} # Allow wrap
+                              ) for doc_type in document_types.keys()]
+                         ])
+                    )
+                ])
+             , color="light"), # Use Bootstrap 'light' color for card background -> light gray
+            width=12, lg=4, # Full width on small screens, 30% (4/12) on large
+            className="mb-3 mb-lg-0", # Margin bottom on small screens
+            style={'padding': '15px'}
+        ),
+
+        # Right Column (Status / Preview / Controls / Chat) - 70% width, white background
+        dbc.Col(
+            dbc.Card(
+                dbc.CardBody([
+                    # Status Bar
+                    dbc.Alert(id='status-bar', children="Upload source documents and select an action.", color="info"),
+
+                    # Dynamic Controls for Reviews - Use Card for visual separation
+                    dbc.Card(id='review-controls-card', children=[dbc.CardBody(id='review-controls')], className="mb-3", style={'display': 'none'}), # Hidden initially
+
+                    # Document Preview Area - Use Card
+                    dbc.Card(
+                        dbc.CardBody([
+                            html.H5("Document Preview / Output", className="card-title"),
+                             # Wrap preview in Loading
+                             dcc.Loading(
+                                 id="loading-preview", # Separate loading for preview
+                                 type="circle",
+                                 children=[html.Div(id='document-preview', style={'whiteSpace': 'pre-wrap', 'maxHeight': '400px', 'overflowY': 'auto', 'border': '1px solid #ccc', 'padding': '10px', 'borderRadius': '5px', 'background': '#f8f9fa'})]
+                            )
+                        ]), className="mb-3"
+                    ),
+                    dbc.Button("Download Output", id="btn-download", color="success", className="mt-3", style={'display': 'none'}), # Hidden initially
+                    dcc.Download(id="download-document"),
+
+                    html.Hr(),
+
+                    # Chat Section - Use Card
+                    dbc.Card(
+                         dbc.CardBody([
+                            html.H5("Refine Output (Chat)", className="card-title"),
+                            # Wrap chat in loading
+                            dcc.Loading(
+                                id="chat-loading",
+                                type="circle",
+                                children=[
+                                    dbc.Textarea(id="chat-input", placeholder="Enter instructions to refine the document shown above...", className="mb-2", style={'whiteSpace': 'normal', 'wordWrap': 'break-word'}), # Ensure word wrap
+                                    dbc.Button("Send Chat", id="btn-send-chat", color="secondary", className="mb-3"), # Use secondary style
+                                    html.Div(id="chat-output", style={'whiteSpace': 'pre-wrap', 'marginTop': '10px', 'border': '1px solid #eee', 'padding': '10px', 'borderRadius': '5px', 'minHeight': '50px'}) # Add border/padding
+                                ]
+                            )
+                         ]), className="mb-3"
+                     )
+                ])
+            ),
+            width=12, lg=8, # Full width on small screens, 70% (8/12) on large
+            style={'backgroundColor': '#ffffff', 'padding': '15px'} # White background
+        )
     ])
-], fluid=True)
+], style={'maxWidth': '100%', 'padding': '0 15px'}) # Max width and padding for container
+
+
+# --- Helper Functions ---
 
 def process_document(contents, filename):
-    content_type, content_string = contents.split(',')
-    decoded = base64.b64decode(content_string)
+    """Processes uploaded file content (PDF or DOCX) and returns text, or None and error message."""
+    if contents is None:
+        logging.warning(f"process_document called with None contents for {filename}")
+        return None, f"Error: No content provided for {filename}."
+
     try:
+        content_type, content_string = contents.split(',')
+        decoded = base64.b64decode(content_string)
+        logging.info(f"Processing file: {filename}")
+        text = None
+        error_message = None
+
         if filename.lower().endswith('.docx'):
-            doc = Document(BytesIO(decoded))
-            text = "\n".join([para.text for para in doc.paragraphs])
-            return text
+            doc = Document(io.BytesIO(decoded))
+            # Extract text, ensuring paragraphs are separated and empty ones are skipped
+            text = "\n".join([para.text for para in doc.paragraphs if para.text.strip()])
+            logging.info(f"Successfully processed DOCX: {filename}")
         elif filename.lower().endswith('.pdf'):
-            pdf = PdfReader(BytesIO(decoded))
-            text = ""
-            for page in pdf.pages:
-                text += page.extract_text()
-            return text
+            pdf = PdfReader(io.BytesIO(decoded))
+            extracted_pages = []
+            for i, page in enumerate(pdf.pages):
+                try:
+                    page_text = page.extract_text()
+                    if page_text:
+                        extracted_pages.append(page_text)
+                except Exception as page_e:
+                    logging.warning(f"Could not extract text from page {i+1} of {filename}: {page_e}")
+            text = "\n\n".join(extracted_pages) # Separate pages clearly
+            if not text:
+                 logging.warning(f"No text extracted from PDF: {filename}. It might be image-based or corrupted.")
+                 error_message = f"Error: No text could be extracted from PDF {filename}. It might be image-based or require OCR."
+            else:
+                logging.info(f"Successfully processed PDF: {filename}")
         else:
-            return f"Unsupported file format: {filename}. Please upload a PDF or DOCX file."
+            logging.warning(f"Unsupported file format: {filename}")
+            error_message = f"Unsupported file format: {filename}. Please upload PDF or DOCX."
+
+        return text, error_message
     except Exception as e:
-        return f"Error processing document: {str(e)}"
+        logging.error(f"Error processing document {filename}: {e}", exc_info=True)
+        return None, f"Error processing file {filename}: {str(e)}"
+
+def get_combined_uploaded_text():
+    """Combines text content of all successfully uploaded files, separated clearly."""
+    if not uploaded_files:
+        return ""
+    # Join contents with a separator indicating file breaks
+    return "\n\n--- FILE BREAK ---\n\n".join(uploaded_files.values())
+
+def generate_ai_document(doc_type, input_docs, context_docs=None):
+    """Generates document using Gemini AI. Updates current_display."""
+    global current_display_document, current_display_type # Allow modification
+
+    if not model:
+        logging.error("Gemini AI model not initialized.")
+        return "Error: AI Model not configured. Please check API Key."
+    if not input_docs or not any(input_docs): # Check if list exists and has content
+        logging.warning(f"generate_ai_document called for {doc_type} with no input documents.")
+        return f"Error: Missing required input document(s) for {doc_type} generation."
+
+    # Combine input documents into a single string
+    combined_input = "\n\n---\n\n".join(filter(None, input_docs))
+    combined_context = "\n\n---\n\n".join(filter(None, context_docs)) if context_docs else ""
+
+    # Enhanced prompt structure based on user feedback and best practices
+    prompt = f"""**Objective:** Generate the '{doc_type}' document.
+
+**Your Role:** Act as an expert proposal writer/analyst.
+
+**Core Instructions:**
+1.  **Adhere Strictly to the Task:** Generate *only* the content for the '{doc_type}'. Do not add introductions, summaries, or conversational filler unless it's part of the requested document format itself.
+2.  **Follow Format Guidelines:**
+    *   **Spreadsheet Types (Shred, Reviews, LOE, Board):** Structure output clearly. Use Markdown tables or a delimited format (like CSV) suitable for parsing. Define clear columns (e.g., `PWS_Section | Requirement | Finding | Recommendation` for reviews; `Section | Task | Estimated_Hours | Resource_Type` for LOE).
+    *   **Proposal Sections (Pink, Red, Gold):** Write professional, compelling prose. Use active voice ("MicroHealth will..."). Directly address requirements from context (Shredded PWS). Detail the 'how' (technical approach, methodology, workflow, tools). Incorporate innovation and benefits (efficiency, quality, outcomes). Substantiate claims (e.g., cite Gartner, Forrester if applicable). Clearly state roles/responsibilities (labor categories). Ensure compliance with Section L/M (Evaluation Criteria) from context. Avoid vague terms ('might', 'could', 'potentially'); be assertive and confident. Use paragraphs primarily; limit bullet points to lists where essential.
+3.  **Utilize Provided Documents:**
+    *   **Context Document(s):** Use these as the primary reference or baseline (e.g., Shredded Requirements are the basis for compliance).
+    *   **Primary Input Document(s):** This is the main subject of the task (e.g., the PWS to be Shredded, the Pink draft to be Reviewed, the Review findings to incorporate into the next draft).
+
+**Provided Documents:**
+
+**Context Document(s) (e.g., Shredded Requirements, PWS Section L/M):**
+```text
+{combined_context if combined_context else "N/A"}
+```
+
+**Primary Input Document(s) (e.g., PWS text, Pink Draft text, Review Findings text):**
+```text
+{combined_input}
+```
+
+**Detailed Instructions for '{doc_type}':**
+{document_types.get(doc_type, "Generate the requested document based on the inputs and context.")}
+
+**Begin '{doc_type}' Output:**
+"""
+
+    logging.info(f"Generating AI document for: {doc_type}")
+    # logging.debug(f"Prompt for {doc_type}: {prompt[:500]}...") # Uncomment for debugging prompt starts
+
+    try:
+        # Increased timeout might be needed for complex generations
+        response = model.generate_content(prompt) # Consider adding request_options={'timeout': 300} if needed
+
+        # Handle potential safety blocks or empty responses
+        if not response.parts:
+             logging.warning(f"Gemini AI returned no parts for {doc_type}. Potential safety block or empty response.")
+             generated_text = f"Error: AI returned no content for {doc_type}. This might be due to safety filters or an issue with the prompt/input."
+        else:
+            generated_text = response.text
+
+        logging.info(f"Successfully generated document for: {doc_type}")
+        # Update global state for download/chat *only if successful*
+        if not generated_text.startswith("Error:"):
+             current_display_document = generated_text
+             current_display_type = doc_type
+        else:
+             # Ensure error message is displayed if AI returns an error internally
+             current_display_document = generated_text
+             current_display_type = doc_type # Still set type so user knows what failed
+
+        return generated_text
+    except Exception as e:
+        logging.error(f"Error during Gemini AI call for {doc_type}: {e}", exc_info=True)
+        # Update display with error message
+        current_display_document = f"Error generating document via AI for {doc_type}: {str(e)}"
+        current_display_type = doc_type
+        return current_display_document
+
+
+# --- Callbacks ---
 
+# 1. Handle File Uploads (Source Documents)
 @app.callback(
     Output('file-list', 'children'),
-    Output('status-bar', 'children'),
+    Output('status-bar', 'children', allow_duplicate=True),
     Input('upload-document', 'contents'),
     State('upload-document', 'filename'),
-    State('file-list', 'children')
+    State('file-list', 'children'),
+    prevent_initial_call=True
 )
-def update_output(list_of_contents, list_of_names, existing_files):
-    global uploaded_files, shredded_document
-    if list_of_contents is not None:
-        new_files = []
-        for i, (content, name) in enumerate(zip(list_of_contents, list_of_names)):
-            file_content = process_document(content, name)
-            uploaded_files[name] = file_content
-            new_files.append(html.Div([
-                html.Button('×', id={'type': 'remove-file', 'index': name}, style={'marginRight': '5px', 'fontSize': '10px'}),
-                html.Span(name)
-            ]))
-        if existing_files is None:
-            existing_files = []
-        shredded_document = None  # Reset shredded document when new files are uploaded
-        return existing_files + new_files, "Document uploaded. Please click 'Shred' to proceed."
-    return existing_files, "Please upload a document and click 'Shred' to begin."
+def handle_file_upload(list_of_contents, list_of_names, existing_files_display):
+    global uploaded_files
+    # Reset downstream data when new source files are uploaded, as context changes
+    global shredded_document, pink_document, pink_review_document, red_document, red_review_document, gold_document, gold_review_document, loe_document, virtual_board_document, current_display_document, current_display_type, uploaded_pink_content, uploaded_red_content, uploaded_gold_content
+
+    status_message = "Please upload source documents (.pdf, .docx) and select an action."
+    if list_of_contents is None:
+        raise PreventUpdate
+
+    new_files_display = []
+    processed_count = 0
+    error_count = 0
+    reset_needed = False
+
+    if existing_files_display is None:
+        existing_files_display = []
+
+    # Get current filenames from the display to avoid duplicates
+    current_filenames = set()
+    if existing_files_display:
+         # Handle potential list vs single item
+        file_list_items = existing_files_display if isinstance(existing_files_display, list) else [existing_files_display]
+        for item in file_list_items:
+             # Check structure carefully based on Div/Button/Span
+             if isinstance(item, html.Div) and len(item.children) > 1 and isinstance(item.children[1], html.Span):
+                 current_filenames.add(item.children[1].children)
+
+
+    for i, (content, name) in enumerate(zip(list_of_contents, list_of_names)):
+        if name in current_filenames:
+            logging.warning(f"Skipping duplicate upload attempt for source file: {name}")
+            continue # Avoid processing duplicates
+
+        file_content_text, error = process_document(content, name)
+
+        if error:
+            logging.error(f"Failed to process source file {name}: {error}")
+            error_count += 1
+            status_message = f"Error processing {name}. {error}" # Show last error
+            continue # Skip adding failed files
+
+        if file_content_text is not None: # Allow empty files if processing is successful
+             uploaded_files[name] = file_content_text
+             # Use dbc.Button for remove, styled small
+             new_files_display.append(html.Div([
+                dbc.Button('X', id={'type': 'remove-file', 'index': name}, size="sm", color="danger", className="me-2 py-0 px-1", n_clicks=0),
+                html.Span(name, title=name) # Add tooltip with full name
+             ], className="d-flex align-items-center mb-1"))
+             processed_count += 1
+             current_filenames.add(name) # Add to tracking set
+             reset_needed = True # Mark that downstream docs should be cleared
+
+    if reset_needed:
+        logging.info("New source files uploaded, resetting downstream generated documents.")
+        shredded_document = None
+        pink_document = None
+        pink_review_document = None
+        red_document = None
+        red_review_document = None
+        gold_document = None
+        gold_review_document = None
+        loe_document = None
+        virtual_board_document = None
+        current_display_document = None # Clear preview
+        current_display_type = None
+        uploaded_pink_content = None # Also clear review uploads if source changes
+        uploaded_red_content = None
+        uploaded_gold_content = None
 
+
+    if processed_count > 0:
+        status_message = f"Successfully uploaded {processed_count} source file(s). Ready for 'Shred' or other actions."
+    elif error_count > 0 and processed_count == 0:
+         status_message = "Failed to process uploaded file(s). Check logs. Ensure they are valid PDF/DOCX with extractable text."
+    elif not new_files_display: # Means only duplicates were uploaded or upload was empty
+        status_message = "No new valid source files were added."
+
+
+    # Combine existing and new display items
+    final_display_list = (existing_files_display if isinstance(existing_files_display, list) else [existing_files_display] if existing_files_display else []) + new_files_display
+
+    return final_display_list, status_message
+
+# 2. Handle File Removal (Source Documents)
 @app.callback(
     Output('file-list', 'children', allow_duplicate=True),
     Output('status-bar', 'children', allow_duplicate=True),
-    Input({'type': 'remove-file', 'index': dash.ALL}, 'n_clicks'),
+    # Use pattern-matching ID for the Input
+    Input({'type': 'remove-file', 'index': ALL}, 'n_clicks'),
     State('file-list', 'children'),
     prevent_initial_call=True
 )
-def remove_file(n_clicks, existing_files):
-    global uploaded_files, shredded_document
-    ctx = dash.callback_context
-    if not ctx.triggered:
-        raise dash.exceptions.PreventUpdate
-    removed_file = ctx.triggered[0]['prop_id'].split(',')[0].split(':')[-1].strip('}')
-    uploaded_files.pop(removed_file, None)
-    shredded_document = None  # Reset shredded document when a file is removed
-    return [file for file in existing_files if file['props']['children'][1]['props']['children'] != removed_file], "Document removed. Please upload a document and click 'Shred' to begin."
-
-def generate_document(document_type, file_contents):
-    prompt = f"""Generate a {document_type} based on the following project artifacts:
-{' '.join(file_contents)}
-Instructions:
-1. Create the {document_type} as a detailed document.
-2. Use proper formatting and structure.
-3. Include all necessary sections and details.
-4. Start the output immediately with the document content.
-5. IMPORTANT: If the document type is Pink, Red, Gold and not review type, loe or board
-then your goal is to be compliant and compelling based on the
-requrements, write in paragraph in active voice as
-MicroHealth, limit bullets, answer the
-requirement with what MicroHealth will do 
-to satisfy the requirement, the technical 
-approach with innovation for efficiency, 
-productivity, quality and measurable 
-outcomes, the industry standard that 
-methodology is based on if applicable, 
-detail the workflow or steps to accomplish 
-the requirement with labor categories that 
-will do those tasks in that workflow, 
-reference reputable research like gartner, 
-forrester, IDC, Deloitte, Accenture etc 
-with measures of success and substantiation 
-of MicroHealth's approach. Never use soft words
-like maybe, could be, should, possible be definitive in your language and confident.
-6. you must also take into account section L&M of the document which is the evaluation criteria
-to be sure we address them.
-Now, generate the {document_type}:
-"""
+def handle_file_remove(n_clicks, current_file_list_display):
+    global uploaded_files
+    # Reset downstream data when a source file is removed
+    global shredded_document, pink_document, pink_review_document, red_document, red_review_document, gold_document, gold_review_document, loe_document, virtual_board_document, current_display_document, current_display_type, uploaded_pink_content, uploaded_red_content, uploaded_gold_content
+
+    triggered_id_dict = callback_context.triggered_id
+    # Check if the callback was triggered by a pattern-matching ID and n_clicks increased
+    if not triggered_id_dict or not isinstance(triggered_id_dict, dict) or 'index' not in triggered_id_dict:
+        raise PreventUpdate
+
+    # Check if any click count is > 0 (or just check the specific one that triggered)
+    if not any(nc > 0 for nc in n_clicks if nc is not None): # Ensure click happened
+         raise PreventUpdate
+
+
+    file_to_remove = triggered_id_dict['index']
+    logging.info(f"Attempting to remove source file: {file_to_remove}")
+
+    if file_to_remove in uploaded_files:
+        del uploaded_files[file_to_remove]
+        logging.info(f"Removed {file_to_remove} from uploaded_files dictionary.")
+
+        # Reset downstream docs since context changed
+        logging.info("Source file removed, resetting downstream generated documents.")
+        shredded_document = None
+        pink_document = None
+        pink_review_document = None
+        red_document = None
+        red_review_document = None
+        gold_document = None
+        gold_review_document = None
+        loe_document = None
+        virtual_board_document = None
+        current_display_document = None # Clear preview
+        current_display_type = None
+        uploaded_pink_content = None # Also clear review uploads
+        uploaded_red_content = None
+        uploaded_gold_content = None
 
-    response = model.generate_content(prompt)
-    return response.text
+    # Filter the display list
+    updated_file_list_display = []
+    if current_file_list_display:
+        # Handle potential list vs single item
+        file_list_items = current_file_list_display if isinstance(current_file_list_display, list) else [current_file_list_display]
+        updated_file_list_display = [
+            item for item in file_list_items
+            # Check structure carefully based on Div/Button/Span
+            if not (isinstance(item, html.Div) and
+                    item.children and isinstance(item.children[0], dbc.Button) and
+                    isinstance(item.children[0].id, dict) and
+                    item.children[0].id.get('index') == file_to_remove)
+        ]
 
+
+    status_message = f"Removed source file: {file_to_remove}. "
+    if not uploaded_files:
+        status_message += "No source files remaining. Please upload documents."
+    else:
+        status_message += "Ready for 'Shred' or other actions."
+
+    return updated_file_list_display, status_message
+
+
+# 3. Handle Action Button Clicks (Show Controls or Trigger Generation)
 @app.callback(
-    Output('document-preview', 'children'),
-    Output('loading-output', 'children'),
+    Output('review-controls-card', 'style'), # Show/hide the whole card
+    Output('review-controls', 'children'),   # Content of the card body
     Output('status-bar', 'children', allow_duplicate=True),
-    Output('pink-review-upload', 'style'),
-    [Input(f'btn-{doc_type.lower().replace("_", "-")}', 'n_clicks') for doc_type in document_types.keys()],
-    State('pink-review-file-name', 'children'),
+    Output('document-preview', 'children', allow_duplicate=True),
+    Output('loading-indicator', 'style'),    # Show/hide main loading indicator (dots)
+    Output('btn-download', 'style', allow_duplicate=True), # Show/hide download button
+    # Use pattern-matching ID for the Input
+    Input({'type': 'action-button', 'index': ALL}, 'n_clicks'),
     prevent_initial_call=True
 )
-def generate_document_preview(*args):
-    global current_document, document_type, shredded_document, pink_review_document
-    ctx = dash.callback_context
-    if not ctx.triggered:
-        raise dash.exceptions.PreventUpdate
-    button_id = ctx.triggered[0]['prop_id'].split('.')[0]
-    document_type = button_id.replace('btn-', '').replace('-', '_').title()
-    pink_review_file = args[-1]
-
-    if not uploaded_files and document_type != "Shred":
-        return html.Div("Please upload and shred a document first."), "", "Please upload and shred a document first.", {'display': 'none'}
-
-    if document_type == "Shred":
-        if not uploaded_files:
-            return html.Div("Please upload a document before shredding."), "", "Please upload a document before shredding.", {'display': 'none'}
-        file_contents = list(uploaded_files.values())
-        try:
-            shredded_document = generate_document(document_type, file_contents)
-            return dcc.Markdown(shredded_document), f"{document_type} generated", "Document shredded. You can now proceed with other operations.", {'display': 'none'}
-        except Exception as e:
-            print(f"Error generating document: {str(e)}")
-            return html.Div(f"Error generating document: {str(e)}"), "Error", "An error occurred while shredding the document.", {'display': 'none'}
+def handle_action_button(n_clicks):
+    global shredded_document, pink_document, red_document, gold_document, pink_review_document, red_review_document, gold_review_document, loe_document, virtual_board_document
+    # Reset potentially uploaded review files when a *new* main action is selected from the left nav
+    global uploaded_pink_content, uploaded_red_content, uploaded_gold_content
 
-    if shredded_document is None:
-        return html.Div("Please shred a document first."), "", "Please shred a document first.", {'display': 'none'}
+    triggered_id_dict = callback_context.triggered_id
+    if not triggered_id_dict or not isinstance(triggered_id_dict, dict) or 'index' not in triggered_id_dict:
+        raise PreventUpdate
 
-    if document_type == "Pink Review":
-        return html.Div("Please upload a Pink Team document or use the generated one."), "", "Please upload a Pink Team document or use the generated one.", {'display': 'block'}
+    # Check if any click count is > 0
+    if not any(nc > 0 for nc in n_clicks if nc is not None):
+         raise PreventUpdate
 
-    if document_type in ["Red", "Red Review"] and pink_review_document is None:
-        return html.Div("Please complete Pink Review first."), "", "Please complete Pink Review first.", {'display': 'none'}
+    action_type = triggered_id_dict['index']
+    logging.info(f"Action button clicked: {action_type}")
 
-    try:
-        if document_type == "Pink Review" and pink_review_file:
-            current_document = generate_document(document_type, [pink_review_file, shredded_document])
-        elif document_type in ["Red", "Red Review"]:
-            current_document = generate_document(document_type, [pink_review_document, shredded_document])
+    # Default states
+    review_controls_style = {'display': 'none'} # Hide review controls by default
+    review_controls_children = []
+    status_message = f"Selected action: {action_type}"
+    doc_preview_children = dash.no_update # Avoid clearing preview unless needed
+    loading_style = {'visibility':'hidden'} # Hide loading dots by default
+    download_style = {'display': 'none'} # Hide download button by default
+
+
+    # Reset previously uploaded review files content when a *new* action is selected.
+    # This prevents using an old uploaded file for a new review type accidentally.
+    uploaded_pink_content = None
+    uploaded_red_content = None
+    uploaded_gold_content = None
+    logging.debug("Cleared any previously uploaded review document content.")
+
+    # --- Actions Requiring Review Controls (Pink/Red/Gold Review) ---
+    if action_type in ["Pink Review", "Red Review", "Gold Review"]:
+        review_controls_style = {'display': 'block'} # Show the review controls card
+        base_doc_type = action_type.split(" ")[0] # Pink, Red, or Gold
+        prereq_doc = None
+        prereq_doc_name = ""
+        generated_doc_to_review = None
+        generated_doc_name = f"Generated {base_doc_type} Document"
+        upload_file_prompt = f"Select {base_doc_type} File"
+
+        # Check common prerequisite: Shredded document
+        if not shredded_document:
+            status_message = "Error: Please 'Shred' the source documents first."
+            doc_preview_children = html.Div(status_message, className="text-danger")
+            return review_controls_style, [], status_message, doc_preview_children, loading_style, download_style
+
+        # Check specific prerequisites and get the document to review
+        if action_type == "Pink Review":
+            prereq_doc = shredded_document # Base requirement
+            prereq_doc_name = "Shredded Document"
+            generated_doc_to_review = pink_document # Generated Pink to review
+        elif action_type == "Red Review":
+            prereq_doc = pink_review_document # Need Pink review results
+            prereq_doc_name = "Pink Review Document"
+            generated_doc_to_review = red_document # Generated Red to review
+        elif action_type == "Gold Review":
+            prereq_doc = red_review_document # Need Red review results
+            prereq_doc_name = "Red Review Document"
+            generated_doc_to_review = gold_document # Generated Gold to review
+
+        # Check if the specific prerequisite (like Pink Review for Red Review) exists
+        if prereq_doc is None and action_type != "Pink Review": # Shred is checked above
+             status_message = f"Error: Please complete '{prereq_doc_name.replace(' Document','')}' first."
+             doc_preview_children = html.Div(status_message, className="text-danger")
+             return review_controls_style, [], status_message, doc_preview_children, loading_style, download_style
+
+
+        # Configure Radio Items based on whether the generated version exists
+        radio_options = []
+        default_value = 'upload' # Default to upload as requested
+        if generated_doc_to_review:
+            radio_options.append({'label': f'Use {generated_doc_name}', 'value': 'generated'})
+            radio_options.append({'label': f'Upload {base_doc_type} Document', 'value': 'upload'})
+            # Keep default 'upload' unless generated is the *only* option (which shouldn't happen here)
         else:
-            current_document = generate_document(document_type, [shredded_document])
+            # If generated doesn't exist, only allow upload
+            radio_options.append({'label': f'Upload {base_doc_type} Document', 'value': 'upload'})
+            status_message = f"Warning: No '{base_doc_type}' document was generated in this session. You must upload one to proceed with {action_type}."
 
-        if document_type == "Pink Review":
-            pink_review_document = current_document
 
-        return dcc.Markdown(current_document), f"{document_type} generated", f"{document_type} document generated successfully.", {'display': 'none'}
-    except Exception as e:
-        print(f"Error generating document: {str(e)}")
-        return html.Div(f"Error generating document: {str(e)}"), "Error", "An error occurred while generating the document.", {'display': 'none'}
+        # Build the controls
+        review_controls_children = [
+            html.H5(f"Configure Input for {action_type}"),
+            dbc.Label(f"Select {base_doc_type} document source:"),
+            dbc.RadioItems(
+                id='review-source-radio', # Single ID for the radio group
+                options=radio_options,
+                value=default_value, # Default to 'upload'
+                inline=True,
+                className='mb-2'
+            ),
+            # Single Upload component, dynamically shown/hidden by radio button callback
+            dcc.Upload(
+                id='upload-review-doc', # Single ID for the upload component
+                children=html.Div(['Drag and Drop or ', html.A(upload_file_prompt)]),
+                style={'display': 'block' if default_value == 'upload' else 'none', # Show/hide based on default value
+                       'width': '100%', 'height': '60px', 'lineHeight': '60px', 'borderWidth': '1px',
+                       'borderStyle': 'dashed', 'borderRadius': '5px', 'textAlign': 'center', 'margin': '10px 0',
+                       'backgroundColor': '#f8f9fa'},
+                multiple=False
+            ),
+            html.Div(id='review-upload-status', className='mb-2 text-muted small'), # For upload confirmation/error
+            # Generate Review button with pattern-matching ID
+            dbc.Button(f"Generate {action_type}", id={'type': 'generate-review-button', 'index': action_type}, color="primary")
+        ]
+        # Clear preview when showing controls, provide instruction
+        doc_preview_children = html.Div(f"Configure input source for {base_doc_type} document and click 'Generate {action_type}'.", style={'padding':'10px'})
+        status_message = f"Ready to configure input for {action_type}."
+
+
+    # --- Actions Triggering Direct Generation (Shred, Pink, Red, Gold, LOE, Virtual Board) ---
+    else:
+        review_controls_style = {'display': 'none'} # Hide review controls
+        review_controls_children = []
+        loading_style = {'visibility':'visible'} # Show loading dots
+        doc_preview_children = "" # Clear preview while loading/generating
+        status_message = f"Generating {action_type}..."
+
+        # Determine inputs based on action type
+        input_docs = []
+        context_docs = []
+        generation_possible = True
 
+        if action_type == "Shred":
+            source_docs_text = get_combined_uploaded_text()
+            if not source_docs_text:
+                status_message = "Error: Please upload source document(s) first."
+                generation_possible = False
+            else:
+                input_docs = [source_docs_text]
+        elif action_type == "Pink":
+            if not shredded_document:
+                status_message = "Error: Please 'Shred' the source documents first."
+                generation_possible = False
+            else:
+                input_docs = [get_combined_uploaded_text()] # Pink is based on source docs
+                context_docs = [shredded_document] # With context of shredded requirements
+        elif action_type == "Red":
+            if not shredded_document or not pink_review_document:
+                status_message = "Error: Please complete 'Shred' and 'Pink Review' first."
+                generation_possible = False
+            else:
+                # Red uses Pink Review findings as primary input to address them
+                input_docs = [pink_review_document]
+                # Context includes Shredded requirements and maybe original Pink? Let's stick to Shred+Review for now.
+                context_docs = [shredded_document]
+        elif action_type == "Gold":
+            if not shredded_document or not red_review_document:
+                status_message = "Error: Please complete 'Shred' and 'Red Review' first."
+                generation_possible = False
+            else:
+                 # Gold uses Red Review findings as primary input
+                input_docs = [red_review_document]
+                context_docs = [shredded_document]
+        elif action_type in ["LOE", "Virtual Board"]:
+             if not shredded_document:
+                status_message = f"Error: Please 'Shred' the source documents first before generating {action_type}."
+                generation_possible = False
+             else:
+                # These likely only need the shredded requirements as input
+                input_docs = [shredded_document]
+        else:
+            status_message = f"Action '{action_type}' is not recognized for direct generation."
+            generation_possible = False
+
+        # Perform generation if possible
+        if generation_possible:
+            result_doc = generate_ai_document(action_type, input_docs, context_docs)
+
+            # Store result in the correct global variable
+            if result_doc and not result_doc.startswith("Error:"):
+                if action_type == "Shred": shredded_document = result_doc
+                elif action_type == "Pink": pink_document = result_doc
+                elif action_type == "Red": red_document = result_doc
+                elif action_type == "Gold": gold_document = result_doc
+                elif action_type == "LOE": loe_document = result_doc
+                elif action_type == "Virtual Board": virtual_board_document = result_doc
+                # Reviews are handled separately
+
+                doc_preview_children = dcc.Markdown(result_doc, style={'wordWrap': 'break-word'})
+                status_message = f"{action_type} generated successfully."
+                download_style = {'display': 'inline-block'} # Show download button on success
+            else:
+                # If generation failed, result_doc contains the error message from generate_ai_document
+                doc_preview_children = html.Div(result_doc, className="text-danger") # Display error in preview
+                status_message = f"Error generating {action_type}. See preview for details."
+                download_style = {'display': 'none'} # Hide download button on error
+
+        else:
+            # Generation not possible due to prerequisites
+            doc_preview_children = html.Div(status_message, className="text-danger")
+
+        loading_style = {'visibility':'hidden'} # Hide loading dots when finished/failed
+
+
+    return review_controls_style, review_controls_children, status_message, doc_preview_children, loading_style, download_style
+
+
+# 4. Toggle Review Upload Component Visibility based on Radio Button
 @app.callback(
-    Output('pink-review-file-name', 'children'),
-    Input('upload-pink-review', 'contents'),
-    State('upload-pink-review', 'filename')
+    Output('upload-review-doc', 'style'),
+    Input('review-source-radio', 'value'),
+    State('upload-review-doc', 'style'), # Get current style to prevent unnecessary updates
+    prevent_initial_call=True
 )
-def update_pink_review_filename(contents, filename):
-    if contents is not None:
-        return filename
-    return ""
+def toggle_review_upload_visibility(radio_value, current_style):
+    # Preserves existing style attributes while toggling 'display'
+    new_style = current_style.copy() if current_style else {}
+    should_display = (radio_value == 'upload')
 
+    if should_display:
+        new_style['display'] = 'block'
+    else:
+        new_style['display'] = 'none'
+
+    # Prevent update if display style is already correct
+    if ('display' in current_style and current_style['display'] == new_style['display']):
+        raise PreventUpdate
+    else:
+        logging.debug(f"Toggling review upload visibility. Radio: {radio_value}, New Style: {new_style}")
+        return new_style
+
+
+# 5. Handle Upload of Document for Review Input
+@app.callback(
+    Output('review-upload-status', 'children'),
+    Output('status-bar', 'children', allow_duplicate=True),
+    Input('upload-review-doc', 'contents'),
+    State('upload-review-doc', 'filename'),
+    # Get the current review type from the button ID that generated the controls
+    State({'type': 'generate-review-button', 'index': ALL}, 'id'),
+    prevent_initial_call=True
+)
+def handle_review_upload(contents, filename, button_ids):
+    global uploaded_pink_content, uploaded_red_content, uploaded_gold_content
+
+    if contents is None or filename is None or not button_ids:
+        # No file uploaded or controls not fully rendered yet
+        raise PreventUpdate
+
+    # Determine which review type this upload is for based on the visible button ID
+    # Assumes only one set of review controls is visible, thus only one button ID in the list
+    if not isinstance(button_ids, list) or len(button_ids) == 0:
+         logging.warning("handle_review_upload: Could not determine review type from button ID state.")
+         raise PreventUpdate
+
+    review_type = button_ids[0]['index'] # e.g., "Pink Review"
+    base_type = review_type.split(" ")[0] # e.g., "Pink"
+
+    logging.info(f"Handling upload of file '{filename}' for {review_type} input.")
+
+    file_content_text, error = process_document(contents, filename)
+
+    upload_status_display = ""
+    status_bar_message = ""
+
+    # Clear previous uploads for this type before storing new one
+    if base_type == "Pink": uploaded_pink_content = None
+    elif base_type == "Red": uploaded_red_content = None
+    elif base_type == "Gold": uploaded_gold_content = None
+
+    if error:
+        status_bar_message = f"Error processing uploaded {base_type} file: {error}"
+        upload_status_display = html.Div(f"Failed to load {filename}: {error}", className="text-danger small")
+    else:
+        status_bar_message = f"Uploaded '{filename}' successfully for {review_type} input."
+        upload_status_display = html.Div(f"Using uploaded file: {filename}", className="text-success small")
+        # Store the content in the correct variable
+        if base_type == "Pink": uploaded_pink_content = file_content_text
+        elif base_type == "Red": uploaded_red_content = file_content_text
+        elif base_type == "Gold": uploaded_gold_content = file_content_text
+        logging.info(f"Stored uploaded content for {base_type} review input.")
+
+    return upload_status_display, status_bar_message
+
+
+# 6. Generate Review Document on Button Click
 @app.callback(
-    Output('chat-output', 'children'),
     Output('document-preview', 'children', allow_duplicate=True),
+    Output('status-bar', 'children', allow_duplicate=True),
+    Output('loading-indicator', 'style', allow_duplicate=True), # Show/hide main loading dots
+    Output('btn-download', 'style', allow_duplicate=True),
+    # Use pattern-matching ID for the Input trigger
+    Input({'type': 'generate-review-button', 'index': ALL}, 'n_clicks'),
+    State('review-source-radio', 'value'), # State of the radio button choice
+    # Get the button ID again to know which review type triggered it
+    State({'type': 'generate-review-button', 'index': ALL}, 'id'),
+    prevent_initial_call=True
+)
+def generate_review_document(n_clicks, source_option, button_ids):
+    global shredded_document, pink_document, red_document, gold_document
+    global pink_review_document, red_review_document, gold_review_document
+    global uploaded_pink_content, uploaded_red_content, uploaded_gold_content
+
+    triggered_id_dict = callback_context.triggered_id
+    if not triggered_id_dict or not isinstance(triggered_id_dict, dict) or 'index' not in triggered_id_dict:
+        raise PreventUpdate
+
+     # Check if any click count is > 0
+    if not any(nc > 0 for nc in n_clicks if nc is not None):
+         raise PreventUpdate
+
+    review_type = triggered_id_dict['index'] # e.g., "Pink Review"
+    base_type = review_type.split(" ")[0] # e.g., "Pink"
+    logging.info(f"Generate button clicked for: {review_type}, Source option chosen: {source_option}")
+
+    doc_preview_children = "" # Clear preview
+    status_message = f"Generating {review_type}..."
+    loading_style = {'visibility':'visible'} # Show loading dots
+    download_style = {'display': 'none'} # Hide download initially
+
+
+    # --- Prerequisite Check ---
+    if not shredded_document:
+        status_message = "Error: 'Shredded' document is missing. Please perform 'Shred' first."
+        loading_style = {'visibility':'hidden'}
+        doc_preview_children = html.Div(status_message, className="text-danger")
+        return doc_preview_children, status_message, loading_style, download_style
+
+    # --- Determine Input Document based on Radio Choice ---
+    input_document_content = None
+    input_doc_source_name = "" # For logging/status messages
+
+    if source_option == 'generated':
+        input_doc_source_name = f"Generated {base_type} Document"
+        if base_type == "Pink": input_document_content = pink_document
+        elif base_type == "Red": input_document_content = red_document
+        elif base_type == "Gold": input_document_content = gold_document
+
+        if not input_document_content:
+            status_message = f"Error: Cannot use 'generated' option. The {input_doc_source_name} was not found (was it generated successfully?)."
+            loading_style = {'visibility':'hidden'}
+            doc_preview_children = html.Div(status_message, className="text-danger")
+            return doc_preview_children, status_message, loading_style, download_style
+
+    elif source_option == 'upload':
+        input_doc_source_name = f"Uploaded {base_type} Document"
+        if base_type == "Pink": input_document_content = uploaded_pink_content
+        elif base_type == "Red": input_document_content = uploaded_red_content
+        elif base_type == "Gold": input_document_content = uploaded_gold_content
+
+        if not input_document_content:
+            # Check if the upload component has children (file name displayed)
+            status_message = f"Error: Cannot use 'upload' option. No {base_type} document was successfully uploaded and processed for this review step."
+            loading_style = {'visibility':'hidden'}
+            doc_preview_children = html.Div(status_message, className="text-danger")
+            return doc_preview_children, status_message, loading_style, download_style
+    else:
+        status_message = f"Error: Invalid source option '{source_option}' selected."
+        loading_style = {'visibility':'hidden'}
+        doc_preview_children = html.Div(status_message, className="text-danger")
+        return doc_preview_children, status_message, loading_style, download_style
+
+
+    # --- Generate Review Document ---
+    logging.info(f"Generating {review_type} using '{input_doc_source_name}' as input and Shredded document as context.")
+
+    # Reviews need the document being reviewed (Pink/Red/Gold) as primary input
+    # and the Shredded PWS as context/requirements basis.
+    review_result = generate_ai_document(review_type, [input_document_content], context_docs=[shredded_document])
+
+    if review_result and not review_result.startswith("Error:"):
+        doc_preview_children = dcc.Markdown(review_result, style={'wordWrap': 'break-word'})
+        status_message = f"{review_type} generated successfully using {input_doc_source_name}."
+        # Store the result in the correct global variable
+        if review_type == "Pink Review": pink_review_document = review_result
+        elif review_type == "Red Review": red_review_document = review_result
+        elif review_type == "Gold Review": gold_review_document = review_result
+        download_style = {'display': 'inline-block'} # Show download button
+    else:
+        # review_result contains the error message
+        doc_preview_children = html.Div(f"Error generating {review_type}: {review_result}", className="text-danger")
+        status_message = f"Failed to generate {review_type}. See preview for details."
+        download_style = {'display': 'none'}
+
+    loading_style = {'visibility':'hidden'} # Hide loading dots
+    return doc_preview_children, status_message, loading_style, download_style
+
+
+# 7. Handle Chat Interaction to Refine Displayed Document
+@app.callback(
+    Output('chat-output', 'children'),                   # Display chat confirmation/error
+    Output('document-preview', 'children', allow_duplicate=True), # Update the preview
+    Output('status-bar', 'children', allow_duplicate=True),     # Update main status
     Input('btn-send-chat', 'n_clicks'),
-    State('chat-input', 'value'),
+    State('chat-input', 'value'),                       # Get the chat instruction
     prevent_initial_call=True
 )
-def update_document_via_chat(n_clicks, chat_input):
-    global current_document, document_type
-    if not chat_input or current_document is None:
-        raise dash.exceptions.PreventUpdate
-    
-    prompt = f"""Update the following {document_type} based on this instruction: {chat_input}
-Current document:
-{current_document}
-Instructions:
-1. Provide the updated document content.
-2. Maintain proper formatting and structure.
-3. Incorporate the requested changes seamlessly.
-Now, provide the updated {document_type}:
+def handle_chat(n_clicks, chat_input):
+    global current_display_document, current_display_type
+    # Also need to update the specific underlying document variable (e.g., pink_document)
+    # so the chat changes persist if that document is used later.
+    global shredded_document, pink_document, red_document, gold_document, pink_review_document, red_review_document, gold_review_document, loe_document, virtual_board_document
+
+    if not n_clicks or not chat_input or not chat_input.strip():
+        # No click or empty input
+        raise PreventUpdate
+    if not current_display_document or not current_display_type:
+         # No document currently loaded in the preview to refine
+         return html.Div("Error: No document is currently displayed to refine.", className="text-warning"), dash.no_update, "Cannot refine: No document loaded in preview."
+
+    logging.info(f"Chat refinement requested for displayed document type: {current_display_type}. Instruction: '{chat_input[:100]}...'")
+
+    # Construct prompt for refinement
+    prompt = f"""**Objective:** Refine the following '{current_display_type}' document based *only* on the user's instruction below.
+
+**Your Role:** Act as an editor making precise changes.
+
+**Core Instructions:**
+1.  **Apply Instruction:** Modify the 'Original Document' solely based on the 'User Instruction'.
+2.  **Maintain Context:** Preserve the overall structure, tone, and format of the original document unless the instruction explicitly directs otherwise.
+3.  **Output Only Updated Document:** Provide *only* the complete, updated '{current_display_type}' document. Do not add any conversational text, preamble, or explanation of changes.
+
+**Original Document:**
+```text
+{current_display_document}
+```
+
+**User Instruction:**
+```text
+{chat_input}
+```
+
+**Begin Updated '{current_display_type}' Output:**
 """
-    
-    response = model.generate_content(prompt)
-    current_document = response.text
-    
-    return f"Document updated based on: {chat_input}", dcc.Markdown(current_document)
 
+    try:
+        # Show loading indicator for chat refinement? (Optional, maybe use inner loading)
+        status_message = f"Refining {current_display_type} based on chat instruction..."
+        # Note: Consider adding a loading state specifically for the chat output area.
+
+        response = model.generate_content(prompt)
+
+        # Handle potential safety blocks or empty responses
+        if not response.parts:
+             logging.warning(f"Gemini AI returned no parts for chat refinement of {current_display_type}.")
+             updated_document = f"Error: AI returned no content during refinement. This might be due to safety filters or an issue with the instruction."
+             chat_response_display = html.Div(updated_document, className="text-danger")
+             status_message = f"Error refining {current_display_type} via chat."
+             return chat_response_display, dash.no_update, status_message # Don't update preview if refinement failed
+        else:
+            updated_document = response.text
+            logging.info(f"Successfully refined {current_display_type} via chat.")
+
+            # --- CRITICAL: Update the correct underlying global variable ---
+            # This ensures the refined document is used in subsequent steps if needed.
+            original_doc_updated = False
+            if current_display_type == "Shred": shredded_document = updated_document; original_doc_updated = True
+            elif current_display_type == "Pink": pink_document = updated_document; original_doc_updated = True
+            elif current_display_type == "Pink Review": pink_review_document = updated_document; original_doc_updated = True
+            elif current_display_type == "Red": red_document = updated_document; original_doc_updated = True
+            elif current_display_type == "Red Review": red_review_document = updated_document; original_doc_updated = True
+            elif current_display_type == "Gold": gold_document = updated_document; original_doc_updated = True
+            elif current_display_type == "Gold Review": gold_review_document = updated_document; original_doc_updated = True
+            elif current_display_type == "LOE": loe_document = updated_document; original_doc_updated = True
+            elif current_display_type == "Virtual Board": virtual_board_document = updated_document; original_doc_updated = True
+
+            if original_doc_updated:
+                 logging.info(f"Updated the underlying global variable for {current_display_type} with chat refinement.")
+            else:
+                 logging.warning(f"Could not map displayed type '{current_display_type}' to a specific global variable for persistent update after chat.")
+
+
+            # Update the preview display immediately
+            current_display_document = updated_document # Keep preview consistent
+
+            # Display confirmation in chat output area
+            chat_response_display = html.Div([
+                html.Strong("Refinement applied successfully."),
+                #html.P(f'Instruction: "{chat_input}"'), # Optional: Echo instruction
+                html.Hr(),
+                html.Em("Preview above has been updated. The changes will be used if this document is input for subsequent steps.")
+            ])
+            status_message = f"{current_display_type} updated via chat instruction."
+            # Update the document preview itself
+            doc_preview_update = dcc.Markdown(updated_document, style={'wordWrap': 'break-word'})
+
+            return chat_response_display, doc_preview_update, status_message
+
+    except Exception as e:
+        logging.error(f"Error during chat refinement call for {current_display_type}: {e}", exc_info=True)
+        chat_response_display = html.Div(f"Error refining document via chat: {str(e)}", className="text-danger")
+        status_message = f"Error refining {current_display_type} via chat."
+        # Do not update the main document preview if chat refinement fails
+        return chat_response_display, dash.no_update, status_message
+
+
+# 8. Handle Download Button Click
 @app.callback(
     Output("download-document", "data"),
     Input("btn-download", "n_clicks"),
     prevent_initial_call=True
 )
-def download_document(n_clicks):
-    global current_document, document_type
-    if current_document is None:
-        raise dash.exceptions.PreventUpdate
-    
-    if document_type == "LOE":
-        # Create a pandas DataFrame for LOE
-        df = pd.read_csv(StringIO(current_document))
-        
-        # Save the DataFrame to an Excel file
-        output = BytesIO()
-        with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
-            df.to_excel(writer, sheet_name='LOE', index=False)
-        
-        return dcc.send_bytes(output.getvalue(), f"{document_type}.xlsx")
-    else:
-        # Create an in-memory Word document
-        doc = Document()
-        doc.add_paragraph(current_document)
-        
-        # Save the document to a BytesIO object
-        output = BytesIO()
-        doc.save(output)
-        
-        return dcc.send_bytes(output.getvalue(), f"{document_type}.docx")
+def download_generated_document(n_clicks):
+    """Prepares the currently displayed document for download."""
+    global current_display_document, current_display_type
+
+    if not n_clicks or current_display_document is None or current_display_type is None:
+        # No clicks or nothing to download
+        raise PreventUpdate
+
+    logging.info(f"Download requested for displayed document: {current_display_type}")
+
+    # Sanitize filename
+    safe_filename_base = "".join(c if c.isalnum() else "_" for c in current_display_type)
+
+    # Determine if output should be spreadsheet (Excel) or document (Word)
+    is_spreadsheet_type = current_display_type in ["Shred", "Pink Review", "Red Review", "Gold Review", "Virtual Board", "LOE"]
+
+    if is_spreadsheet_type:
+        filename = f"{safe_filename_base}.xlsx"
+        logging.info(f"Attempting to format {current_display_type} as Excel.")
+        try:
+            # Use StringIO to treat the string data as a file for pandas
+            data_io = StringIO(current_display_document)
+            df = None
+
+            # Attempt 1: Try parsing as CSV (or detect delimiter)
+            try:
+                 # Read a sample to sniff delimiter
+                 sniffer_sample = data_io.read(2048) # Read more data for better sniffing
+                 data_io.seek(0) # Reset pointer after reading sample
+                 dialect = pd.io.parsers.readers.csv.Sniffer().sniff(sniffer_sample, delimiters=',|\t') # Sniff common delimiters
+                 df = pd.read_csv(data_io, sep=dialect.delimiter)
+                 logging.info(f"Successfully parsed {current_display_type} using detected delimiter '{dialect.delimiter}'.")
+            except Exception as e_csv:
+                 logging.warning(f"Could not parse {current_display_type} as standard CSV/TSV ({e_csv}). Trying Markdown Table parsing.")
+                 data_io.seek(0) # Reset pointer
+
+                 # Attempt 2: Try parsing as a Markdown table (simple version)
+                 lines = [line.strip() for line in data_io.readlines() if line.strip()]
+                 header = []
+                 data = []
+                 header_found = False
+                 separator_found = False
+
+                 for line in lines:
+                      if line.startswith('|') and line.endswith('|'):
+                          parts = [p.strip() for p in line.strip('|').split('|')]
+                          if not header_found:
+                              header = parts
+                              header_found = True
+                          elif '---' in line: # Detect separator line
+                              separator_found = True
+                              # Optional: Check if separator alignment matches header count
+                              if len(line.strip('|').split('|')) != len(header):
+                                   logging.warning("Markdown table header/separator mismatch detected.")
+                                   # Decide whether to proceed or fail parsing
+                          elif header_found and separator_found: # Only add data lines after header and separator
+                               if len(parts) == len(header):
+                                  data.append(parts)
+                               else:
+                                   logging.warning(f"Markdown table row data mismatch (expected {len(header)}, got {len(parts)}): {line}")
+
+
+                 if header and data:
+                     df = pd.DataFrame(data, columns=header)
+                     logging.info(f"Successfully parsed {current_display_type} as Markdown Table.")
+                 else:
+                     logging.warning(f"Could not parse {current_display_type} as Markdown Table after CSV attempt failed.")
+                     # Fallback: If no DataFrame could be created, send as text
+                     return dict(content=current_display_document, filename=f"{safe_filename_base}.txt")
+
+            # If DataFrame was created, save to Excel
+            output = BytesIO()
+            # Use xlsxwriter engine for better compatibility/features if needed
+            with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
+                df.to_excel(writer, sheet_name=current_display_type[:31], index=False) # Sheet name limit is 31 chars
+            logging.info(f"Sending {filename} (Excel format)")
+            return dcc.send_bytes(output.getvalue(), filename)
+
+        except Exception as e_excel:
+            logging.error(f"Error creating Excel file for {current_display_type}: {e_excel}. Sending as text.", exc_info=True)
+            # Fallback to sending as a text file if any DataFrame/Excel processing fails
+            return dict(content=current_display_document, filename=f"{safe_filename_base}.txt")
+
+    else: # Assume DOCX for Pink, Red, Gold
+        filename = f"{safe_filename_base}.docx"
+        logging.info(f"Formatting {current_display_type} as DOCX.")
+        try:
+            doc = Document()
+            # Add paragraph by paragraph to potentially preserve some structure like line breaks
+            for paragraph_text in current_display_document.split('\n'):
+                # Add paragraph only if it contains non-whitespace characters
+                if paragraph_text.strip():
+                    doc.add_paragraph(paragraph_text)
+                else:
+                     # Optionally add an empty paragraph to represent blank lines, or skip
+                     doc.add_paragraph() # Add blank lines if they were intentional
+
+            # Save the document to an in-memory BytesIO object
+            output = BytesIO()
+            doc.save(output)
+            logging.info(f"Sending {filename} (DOCX format)")
+            return dcc.send_bytes(output.getvalue(), filename) # Use dcc.send_bytes for BytesIO
+        except Exception as e_docx:
+            logging.error(f"Error creating DOCX file for {current_display_type}: {e_docx}. Sending as text.", exc_info=True)
+             # Fallback to sending as a text file
+            return dict(content=current_display_document, filename=f"{safe_filename_base}.txt")
+
 
+# --- Main Execution ---
+# Always use this structure for running the app
 if __name__ == '__main__':
     print("Starting the Dash application...")
+    # Set debug=False for production/deployment environments like Hugging Face Spaces
+    # Set host='0.0.0.0' to make the app accessible on the network (required for Docker/Spaces)
+    # Default port 8050, using 7860 as often used for ML demos/Spaces
+    # Use server=app.server for Gunicorn compatibility (multi-threading via workers)
+    # Multi-threading for multiple simultaneous user support is handled by the deployment server (e.g., Gunicorn with workers), not directly in app.run for production.
     app.run(debug=False, host='0.0.0.0', port=7860)
     print("Dash application has finished running.")