app.py

import base64
import io
import os
import pandas as pd
from docx import Document
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, ALL, no_update
from dash.exceptions import PreventUpdate
import google.generativeai as genai
from docx.shared import Pt
from docx.enum.style import WD_STYLE_TYPE
from PyPDF2 import PdfReader
import logging
import uuid
import xlsxwriter # Needed for Excel export engine

# --- Logging Configuration ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# --- Initialize Dash app ---
# dash==3.0.3
# dash-bootstrap-components==2.0.2
app = dash.Dash(__name__,
                external_stylesheets=[dbc.themes.BOOTSTRAP],
                suppress_callback_exceptions=True,
                meta_tags=[{"name": "viewport", "content": "width=device-width, initial-scale=1"}])
server = app.server

# --- Configure Gemini AI ---
# IMPORTANT: Set the GEMINI_API_KEY environment variable.
try:
    # Prefer direct CUDA GPU configuration in app.py - Note: Not applicable for cloud APIs like Gemini.
    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)
        # Using 'gemini-1.5-pro-latest' or similar advanced model is recommended.
        model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
        logging.info("Gemini AI configured successfully using 'gemini-2.5-pro-preview-03-25'.")
except Exception as e:
    logging.error(f"Error configuring Gemini AI: {e}", exc_info=True)
    model = None

# --- Global Variables ---
# Consider dcc.Store for more robust multi-user state management.
uploaded_files = {} # {filename: content_text}

# Stores the *results* of generation/review steps
shredded_document = None
pink_review_document = None
red_review_document = None
gold_review_document = None
loe_document = None
virtual_board_document = None

# Stores the *generated* proposal drafts
pink_document = None
red_document = None
gold_document = None

# 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 ---
document_types = {
    "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."
}

# --- Layout Definition ---
app.layout = dbc.Container(fluid=True, className="dbc", children=[
    # Title Row
    dbc.Row(
        dbc.Col(html.H1("Proposal AI Assistant", className="text-center my-4", style={'color': '#1C304A'}), width=12)
    ),

    # Progress Indicator Row
    dbc.Row(
        dbc.Col(
            dcc.Loading(
                id="loading-indicator",
                type="dots",
                children=[html.Div(id="loading-output", style={'height': '10px'})],
                overlay_style={"visibility":"hidden", "opacity": 0},
                style={'visibility':'hidden', 'height': '30px'},
                fullscreen=False,
                className="justify-content-center"
            ),
            width=12,
            className="text-center mb-3"
        )
    ),

    # Main Content Row
    dbc.Row([
        # Left Column (Nav/Upload)
        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')]),
                        style={
                            'width': '100%', 'height': '60px', 'lineHeight': '60px',
                            'borderWidth': '1px', 'borderStyle': 'dashed', 'borderRadius': '5px',
                            'textAlign': 'center', 'margin': '10px 0', 'backgroundColor': '#ffffff'
                        },
                        multiple=True
                    ),
                    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"),
                    dbc.Card(
                        dbc.CardBody([
                            *[dbc.Button(
                                doc_type,
                                id={'type': 'action-button', 'index': doc_type},
                                color="primary",
                                className="mb-2 w-100 d-block",
                                style={'textAlign': 'left', 'whiteSpace': 'normal', 'height': 'auto', 'wordWrap': 'break-word'}
                              ) for doc_type in document_types.keys()]
                         ])
                    )
                ])
             , color="light"),
            width=12, lg=4,
            className="mb-3 mb-lg-0",
            style={'padding': '15px'}
        ),

        # Right Column (Status/Preview/Controls/Chat)
        dbc.Col(
            dbc.Card(
                dbc.CardBody([
                    dbc.Alert(id='status-bar', children="Upload source documents and select an action.", color="info"),
                    dbc.Card(id='review-controls-card', children=[dbc.CardBody(id='review-controls')], className="mb-3", style={'display': 'none'}),
                    dbc.Card(
                        dbc.CardBody([
                            html.H5("Document Preview / Output", className="card-title"),
                             dcc.Loading(
                                 id="loading-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 me-2", style={'display': 'none'}),
                    dcc.Download(id="download-document"),
                    html.Hr(),
                    dbc.Card(
                         dbc.CardBody([
                            html.H5("Refine Output (Chat)", className="card-title"),
                            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'}),
                                    dbc.ButtonGroup([
                                        dbc.Button("Send Chat", id="btn-send-chat", color="secondary"),
                                        dbc.Button("Clear Chat", id="btn-clear-chat", color="tertiary")
                                    ], className="mb-3"),
                                    html.Div(id="chat-output", style={'whiteSpace': 'pre-wrap', 'marginTop': '10px', 'border': '1px solid #eee', 'padding': '10px', 'borderRadius': '5px', 'minHeight': '50px'})
                                ]
                            )
                         ]), className="mb-3"
                     )
                ])
            ),
            width=12, lg=8,
            style={'backgroundColor': '#ffffff', 'padding': '15px'}
        )
    ])
], style={'maxWidth': '100%', 'padding': '0 15px'})


# --- Helper Functions ---

def process_document(contents, filename):
    """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(io.BytesIO(decoded))
            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(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)
            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:
            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:
        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."""
    if not uploaded_files:
        return ""
    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

    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(doc.strip() for doc in input_docs if doc):
        logging.warning(f"generate_ai_document called for {doc_type} with no valid input documents.")
        return f"Error: Missing required input document(s) for {doc_type} generation."

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

    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). Use '|' as the primary delimiter for tables.
    *   **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}