app.py

import base64
import io
import os
import threading
import time
from typing import List, Tuple
import re

import dash
import dash_bootstrap_components as dbc
from dash import html, dcc, Input, Output, State, ctx
import google.generativeai as genai
from docx import Document
from PyPDF2 import PdfReader

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

def process_document(contents: str, filename: str) -> str:
    content_type, content_string = contents.split(',')
    decoded = base64.b64decode(content_string)
    
    if filename.endswith('.pdf'):
        pdf = PdfReader(io.BytesIO(decoded))
        text = ""
        for page in pdf.pages:
            text += page.extract_text()
    elif filename.endswith('.docx'):
        doc = Document(io.BytesIO(decoded))
        text = "\n".join([para.text for para in doc.paragraphs])
    else:
        return "Unsupported file format. Please upload a PDF or DOCX file."
    
    return text

def generate_outline(text: str, instructions: str) -> str:
    prompt = f"""
    Analyze the following Project Work Statement (PWS) and create an outline 
    focusing on sections L&M. Extract the main headers, subheaders, and specific 
    requirements in each section. Pay special attention to requirements indicated 
    by words like "shall", "will", "must", and similar imperative language.

    Additional instructions: {instructions}

    Document text:
    {text}

    Provide the outline in a structured format, clearly highlighting the specific 
    requirements and their associated sections.
    """
    response = model.generate_content(prompt)
    return response.text

def generate_pink_team_document(outline: str, instructions: str) -> str:
    prompt = f"""
    Based on the following outline of a Project Work Statement (PWS):

    {outline}

    Additional instructions: {instructions}

    Create a detailed response document as if MicroHealth is responding to this PWS. 
    Follow these guidelines:
    1. Use Wikipedia style writing with active voice.
    2. For each requirement, describe in detail how MicroHealth will innovate to address it.
    3. Explain the industry best practices that will be applied.
    4. Provide measurable outcomes for the customer.
    5. Limit the use of bullet points and write predominantly in paragraph format.
    6. Ensure a logical flow of steps taken by MicroHealth for each requirement.

    Generate a comprehensive response that showcases MicroHealth's expertise and approach.
    """
    response = model.generate_content(prompt)
    return response.text

def evaluate_compliance(document: str, requirements: str) -> str:
    prompt = f"""
    Evaluate the following document against the requirements from sections L&M of the PWS:

    Document:
    {document}

    Requirements:
    {requirements}

    Provide a compliance report by section number, highlighting:
    1. Areas that need improvement
    2. Suggestions on how MicroHealth can better respond to the requirements
    3. Best industry practices that should be applied
    4. Measurable outcomes that should be included

    Format the report clearly by section number.
    """
    response = model.generate_content(prompt)
    return response.text

def generate_red_document(document: str, compliance_report: str) -> str:
    prompt = f"""
    Based on the following document and compliance report:

    Original Document:
    {document}

    Compliance Report:
    {compliance_report}

    Generate a revised "Red Team" document that addresses all issues found in the compliance report. 
    Ensure that the new document:
    1. Responds to all requirements with sufficient detail
    2. Explains how MicroHealth will implement solutions
    3. Incorporates best industry practices
    4. Includes measurable outcomes for each requirement

    Provide the revised document in a clear, structured format.
    """
    response = model.generate_content(prompt)
    return response.text

# Layout
app.layout = dbc.Container([
    html.H1("MicroHealth PWS Analysis and Response Generator", className="my-4"),
    dbc.Tabs([
        dbc.Tab(label="Shred", tab_id="shred", children=[
            dbc.Textarea(
                id='shred-instructions',
                placeholder="Enter any additional instructions for shredding the document...",
                style={'height': '100px', 'marginBottom': '10px'}
            ),
            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'
                },
                multiple=False
            ),
            dbc.Spinner(html.Div(id='shred-output')),
            dbc.Button("Download Outline", id="download-shred", className="mt-3"),
            dcc.Download(id="download-shred-doc")
        ]),
        dbc.Tab(label="Pink", tab_id="pink", children=[
            dbc.Textarea(
                id='pink-instructions',
                placeholder="Enter any additional instructions for generating the Pink Team document...",
                style={'height': '100px', 'marginBottom': '10px'}
            ),
            dbc.Button("Generate Pink Team Document", id="generate-pink", className="mt-3"),
            dbc.Spinner(html.Div(id='pink-output')),
            dbc.Button("Download Pink Team Document", id="download-pink", className="mt-3"),
            dcc.Download(id="download-pink-doc")
        ]),
        dbc.Tab(label="P.Review", tab_id="p-review", children=[
            dcc.Upload(
                id='upload-p-review',
                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'
                },
                multiple=False
            ),
            dbc.Button("Evaluate Compliance", id="evaluate-p-review", className="mt-3"),
            dbc.Spinner(html.Div(id='p-review-output')),
            dbc.Button("Download P.Review Report", id="download-p-review", className="mt-3"),
            dcc.Download(id="download-p-review-doc")
        ]),
        dbc.Tab(label="Red", tab_id="red", children=[
            dcc.Upload(
                id='upload-red',
                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'
                },
                multiple=False
            ),
            dbc.Button("Generate Red Team Document", id="generate-red", className="mt-3"),
            dbc.Spinner(html.Div(id='red-output')),
            dbc.Button("Download Red Team Document", id="download-red", className="mt-3"),
            dcc.Download(id="download-red-doc")
        ]),
        dbc.Tab(label="R.Review", tab_id="r-review", children=[
            dcc.Upload(
                id='upload-r-review',
                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'
                },
                multiple=False
            ),
            dbc.Button("Evaluate Compliance", id="evaluate-r-review", className="mt-3"),
            dbc.Spinner(html.Div(id='r-review-output')),
            dbc.Button("Download R.Review Report", id="download-r-review", className="mt-3"),
            dcc.Download(id="download-r-review-doc")
        ]),
    ], id="tabs", active_tab="shred"),
])

@app.callback(
    Output('shred-output', 'children'),
    Input('upload-document', 'contents'),
    State('upload-document', 'filename'),
    State('shred-instructions', 'value')
)
def update_shred_output(contents, filename, instructions):
    if contents is None:
        return "Upload a document to begin."
    
    text = process_document(contents, filename)
    outline = generate_outline(text, instructions or "")
    return dcc.Markdown(outline)

@app.callback(
    Output('pink-output', 'children'),
    Input('generate-pink', 'n_clicks'),
    State('shred-output', 'children'),
    State('pink-instructions', 'value')
)
def update_pink_output(n_clicks, shred_output, instructions):
    if n_clicks is None or shred_output is None:
        return "Generate an outline in the Shred tab first."
    
    pink_doc = generate_pink_team_document(shred_output, instructions or "")
    return dcc.Markdown(pink_doc)

@app.callback(
    Output('p-review-output', 'children'),
    Input('evaluate-p-review', 'n_clicks'),
    State('upload-p-review', 'contents'),
    State('upload-p-review', 'filename'),
    State('pink-output', 'children'),
    State('shred-output', 'children')
)
def update_p_review_output(n_clicks, contents, filename, pink_doc, requirements):
    if n_clicks is None:
        return "Click 'Evaluate Compliance' to begin."
    
    if contents:
        document = process_document(contents, filename)
    elif pink_doc:
        document = pink_doc
    else:
        return "Please upload a document or generate a Pink Team document first."
    
    compliance_report = evaluate_compliance(document, requirements)
    return dcc.Markdown(compliance_report)

@app.callback(
    Output('red-output', 'children'),
    Input('generate-red', 'n_clicks'),
    State('upload-red', 'contents'),
    State('upload-red', 'filename'),
    State('p-review-output', 'children')
)
def update_red_output(n_clicks, contents, filename, p_review_output):
    if n_clicks is None:
        return "Click 'Generate Red Team Document' to begin."
    
    if contents:
        document = process_document(contents, filename)
    elif p_review_output:
        document = p_review_output
    else:
        return "Please upload a document or complete the P.Review first."
    
    red_doc = generate_red_document(document, p_review_output)
    return dcc.Markdown(red_doc)

@app.callback(
    Output('r-review-output', 'children'),
    Input('evaluate-r-review', 'n_clicks'),
    State('upload-r-review', 'contents'),
    State('upload-r-review', 'filename'),
    State('red-output', 'children'),
    State('shred-output', 'children')
)
def update_r_review_output(n_clicks, contents, filename, red_doc, requirements):
    if n_clicks is None:
        return "Click 'Evaluate Compliance' to begin."
    
    if contents:
        document = process_document(contents, filename)
    elif red_doc:
        document = red_doc
    else:
        return "Please upload a document or generate a Red Team document first."
    
    compliance_report = evaluate_compliance(document, requirements)
    return dcc.Markdown(compliance_report)

@app.callback(
    Output("download-shred-doc", "data"),
    Input("download-shred", "n_clicks"),
    State('shred-output', 'children'),
    prevent_initial_call=True,
)
def download_shred(n_clicks, shred_output):
    if shred_output is None:
        return dash.no_update
    return dict(content=shred_output, filename="shred_outline.md")

@app.callback(
    Output("download-pink-doc", "data"),
    Input("download-pink", "n_clicks"),
    State('pink-output', 'children'),
    prevent_initial_call=True,
)
def download_pink(n_clicks, pink_output):
    if pink_output is None:
        return dash.no_update
    return dict(content=pink_output, filename="pink_team_document.md")

@app.callback(
    Output("download-p-review-doc", "data"),
    Input("download-p-review", "n_clicks"),
    State('p-review-output', 'children'),
    prevent_initial_call=True,
)
def download_p_review(n_clicks, p_review_output):
    if p_review_output is None:
        return dash.no_update
    return dict(content=p_review_output, filename="p_review_report.md")

@app.callback(
    Output("download-red-doc", "data"),
    Input("download-red", "n_clicks"),
    State('red-output', 'children'),
    prevent_initial_call=True,
)
def download_red(n_clicks, red_output):
    if red_output is None:
        return dash.no_update
    return dict(content=red_output, filename="red_team_document.md")

@app.callback(
    Output("download-r-review-doc", "data"),
    Input("download-r-review", "n_clicks"),
    State('r-review-output', 'children'),
    prevent_initial_call=True,
)
def download_r_review(n_clicks, r_review_output):
    if r_review_output is None:
        return dash.no_update
    return dict(content=r_review_output, filename="r_review_report.md")

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