app.py

import base64
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
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

# Initialize Dash app
app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])

# Configure Gemini AI
genai.configure(api_key=os.environ["GEMINI_API_KEY"])
model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')

# Global variables
uploaded_files = {}
current_document = None
document_type = None
shredded_document = None
pink_review_document = None

# Document types and their descriptions
document_types = {
    "Shred": "Generate a requirements outline of the Project Work Statement (PWS) identified by action words like shall, will, perform etc.",
    "Pink": "Create a Pink Team document based on the PWS outline",
    "Pink Review": "Evaluate compliance of the Pink Team document",
    "Red": "Generate a Red Team document based on the P.Review",
    "Red Review": "Evaluate compliance of the Red Team document",
    "Gold": "Create a Pink Team document based on the PWS outline",
    "Gold Review": "Perform a final compliance review",
    "Compliance Check": "Perform a final compliance review",
    "LOE": "Generate a Level of Effort (LOE) breakdown as a spreadsheet"
}

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'}),
            dcc.Loading(
                id="loading-indicator",
                type="dot",
                children=[html.Div(id="loading-output")]
            ),
            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)
    ])
], fluid=True)

def process_document(contents, filename):
    content_type, content_string = contents.split(',')
    decoded = base64.b64decode(content_string)
    try:
        if filename.lower().endswith('.docx'):
            doc = Document(BytesIO(decoded))
            text = "\n".join([para.text for para in doc.paragraphs])
            return text
        elif filename.lower().endswith('.pdf'):
            pdf = PdfReader(BytesIO(decoded))
            text = ""
            for page in pdf.pages:
                text += page.extract_text()
            return text
        else:
            return f"Unsupported file format: {filename}. Please upload a PDF or DOCX file."
    except Exception as e:
        return f"Error processing document: {str(e)}"

@app.callback(
    Output('file-list', 'children'),
    Output('status-bar', 'children'),
    Input('upload-document', 'contents'),
    State('upload-document', 'filename'),
    State('file-list', 'children')
)
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."

@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'),
    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.
Now, generate the {document_type}:
"""

    response = model.generate_content(prompt)
    return response.text

@app.callback(
    Output('document-preview', 'children'),
    Output('loading-output', 'children'),
    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'),
    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'}

    if shredded_document is None:
        return html.Div("Please shred a document first."), "", "Please shred a document first.", {'display': 'none'}

    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'}

    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'}

    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])
        else:
            current_document = generate_document(document_type, [shredded_document])

        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'}

@app.callback(
    Output('pink-review-file-name', 'children'),
    Input('upload-pink-review', 'contents'),
    State('upload-pink-review', 'filename')
)
def update_pink_review_filename(contents, filename):
    if contents is not None:
        return filename
    return ""

@app.callback(
    Output('chat-output', 'children'),
    Output('document-preview', 'children', allow_duplicate=True),
    Input('btn-send-chat', 'n_clicks'),
    State('chat-input', 'value'),
    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}:
"""
    
    response = model.generate_content(prompt)
    current_document = response.text
    
    return f"Document updated based on: {chat_input}", dcc.Markdown(current_document)

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

if __name__ == '__main__':
    print("Starting the Dash application...")
    app.run(debug=False, host='0.0.0.0', port=7860)
    print("Dash application has finished running.")