app.py

import base64
import io
import os
import pandas as pd
from docx import Document
from io import BytesIO, StringIO
import dash
import dash_bootstrap_components as dbc
from dash import html, dcc, Input, Output, State, callback_context, MATCH, ALL
from dash.dash_table import DataTable
from docx.shared import Pt
from docx.enum.style import WD_STYLE_TYPE
from PyPDF2 import PdfReader
import logging
import threading
import re 
import markdown
from bs4 import BeautifulSoup
import google.generativeai as genai

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

GEMINI_KEY = os.environ.get("GEMINI_KEY", "")
genai.configure(api_key=GEMINI_KEY)
GEMINI_MODEL = "gemini-2.5-pro-preview-03-25"

app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP], suppress_callback_exceptions=True)

uploaded_files = {}
uploaded_doc_contents = {}

spreadsheet_types = ["Shred", "Pink Review", "Red Review", "Gold Review", "Virtual Board", "LOE"]
narrative_types = ["Pink", "Red", "Gold"]

document_types = {
    "Shred": "Ignore all other instructions and generate only 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 highly detailed Pink Team document based on the PWS outline.  Your goal is to be compliant and compelling.  Focus on describing the approach and how it will be done, the steps, workflow, people, processes and technology based on well known industry standards to accomplish the task.  Be sure to demonstrate innovation.",
    "Pink Review": "Ignore all other instructions and generate and evaluate compliance of the Pink Team document against the requirements and output only 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. you must also take into account section L&M of the document which is the evaluation criteria to be sure we address them.",
    "Red": "Produce a  highly detailed 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. Focus on describing the approach and how it will be done, the steps, workflow, people, processes and technology to accomplish the task. Be sure to refer to research that validates the approach and cite sources with measurable outcomes",
    "Red Review": "Ignore all other instructions and generate and evaluate compliance of the Red Team document against the requirements and output a only 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. you must also take into account section L&M of the document which is the evaluation criteria to be sure we address them",
    "Gold": "Create a highly detailed Gold 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.  Focus on describing the approach and how it will be done, the steps, workflow, people, processes and technology to accomplish the task. Be sure to refer to research that validates the approach and cite sources with measurable outcomes and improve on innovations of the approach",
    "Gold Review": "Ignore all other instructions and generate and perform a final compliance review against the requirements and output only 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. you must also take into account section L&M of the document which is the evaluation criteria to be sure we address them",
    "Virtual Board": "Ignore all other instructions and generate and 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 and produce only spreadsheet",
    "LOE": "Ignore all other instructions and generate and generate a Level of Effort (LOE) breakdown and produce only spreadsheet"
}

doc_dependencies = {
    "Pink": {"source": ["shred"], "require": True},
    "Pink Review": {"source": ["pink", "shred"], "require": True},
    "Red": {"source": ["pink_review"], "require": True},
    "Red Review": {"source": ["red", "shred"], "require": True},
    "Gold": {"source": ["red_review"], "require": True},
    "Gold Review": {"source": ["gold", "shred"], "require": True},
    "LOE": {"source": ["gold"], "require": True},
    "Virtual Board": {"source": ["shred"], "require": True},
}

def extract_markdown_tables(md_text):
    tables = []
    lines = md_text.split('\n')
    in_table = False
    table_lines = []
    for line in lines:
        if re.match(r'^\s*\|.*\|\s*$', line):
            in_table = True
            table_lines.append(line)
        elif in_table and (re.match(r'^\s*\|.*\|\s*$', line) or re.match(r'^\s*$', line)):
            table_lines.append(line)
        else:
            if in_table and table_lines:
                tables.append('\n'.join(table_lines))
                table_lines = []
            in_table = False
    if in_table and table_lines:
        tables.append('\n'.join(table_lines))
    return tables

def markdown_table_to_df(md_table):
    lines = [line.strip() for line in md_table.split('\n') if line.strip()]
    if len(lines) < 2:
        return None
    header = [h.strip() for h in lines[0].strip('|').split('|')]
    sep_idx = 1
    while sep_idx < len(lines) and not re.match(r'^\|\s*:?-+:?\s*(\|\s*:?-+:?\s*)+\|?$', lines[sep_idx]):
        sep_idx += 1
    data_lines = lines[sep_idx+1:] if sep_idx+1 < len(lines) else []
    rows = []
    for row in data_lines:
        if not row.strip() or not row.strip().startswith('|'):
            continue
        cells = [c.strip() for c in row.strip('|').split('|')]
        if len(cells) < len(header):
            cells += [''] * (len(header) - len(cells))
        elif len(cells) > len(header):
            cells = cells[:len(header)]
        rows.append(cells)
    df = pd.DataFrame(rows, columns=header)
    return df

def markdown_table_preview(md_text):
    tables = extract_markdown_tables(md_text)
    if not tables:
        return html.Div("No table found.")
    table_divs = []
    for i, table in enumerate(tables):
        df = markdown_table_to_df(table)
        if df is not None and not df.empty:
            table_divs.append(
                html.Div([
                    DataTable(
                        columns=[{"name": str(col), "id": str(col)} for col in df.columns],
                        data=df.to_dict('records'),
                        style_table={'overflowX': 'auto'},
                        style_cell={'whiteSpace': 'normal', 'height': 'auto', 'textAlign': 'left', 'fontFamily': 'monospace', 'fontSize': '14px', 'maxWidth': '400px', 'minWidth': '80px', 'wordBreak': 'break-word'},
                        style_header={'fontWeight': 'bold'},
                        page_size=100,
                        id={'type': 'datatable-preview', 'index': i}
                    )
                ], className="mb-4")
            )
    return html.Div(table_divs)

def markdown_narrative_preview(md_text):
    return html.Div(dcc.Markdown(md_text, dangerously_allow_html=True, style={'whiteSpace': 'pre-wrap', 'fontFamily': 'sans-serif'}))

def markdown_tables_to_xlsx(md_text):
    tables = extract_markdown_tables(md_text)
    output = BytesIO()
    with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
        for i, table in enumerate(tables):
            df = markdown_table_to_df(table)
            if df is not None:
                sheet_name = f"Table{i+1}"
                df.to_excel(writer, sheet_name=sheet_name, index=False)
    output.seek(0)
    return output

def strip_markdown(text):
    text = re.sub(r'(\*\*|__)(.*?)\1', r'\2', text)
    text = re.sub(r'(\*|_)(.*?)\1', r'\2', text)
    text = re.sub(r'`{1,3}[^`]*`{1,3}', '', text)
    text = re.sub(r'^#+ ', '', text, flags=re.MULTILINE)
    text = re.sub(r'^> ', '', text, flags=re.MULTILINE)
    text = re.sub(r'!\[.*?\]\(.*?\)', '', text)
    text = re.sub(r'\[([^\]]+)\]\([^)]+\)', r'\1', text)
    text = re.sub(r'^\s*[-*+] ', '', text, flags=re.MULTILINE)
    text = re.sub(r'^\s*\d+\.\s+', '', text, flags=re.MULTILINE)
    text = text.replace('---', '')
    text = text.replace('___', '')
    text = text.replace('***', '')
    return text.strip()

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:
                page_text = page.extract_text()
                if page_text:
                    text += page_text
            return text
        else:
            return f"Unsupported file format: {filename}. Please upload a PDF or DOCX file."
    except Exception as e:
        logging.error(f"Error processing document: {str(e)}")
        return f"Error processing document: {str(e)}"

def file_list_component():
    return html.Div(
        id='file-list-container',
        children=[
            html.Div(id='file-list')
        ]
    )

def get_left_col_content():
    chat_card = dbc.Card(
        dbc.CardBody([
            html.H5("Maiko Chat", className="mb-2"),
            dcc.Loading(
                id="chat-loading",
                type="dot",
                children=[
                    dcc.Textarea(
                        id="chat-input",
                        placeholder="Chat with AI to update document...",
                        className="mb-2",
                        style={
                            'whiteSpace':'pre-wrap',
                            'width': '100%',
                            'minHeight': '100px',
                            'maxHeight': '300px',
                            'resize': 'vertical',
                            'overflowY': 'auto'
                        },
                        rows=5,
                        wrap='soft'
                    ),
                    dcc.Store(id="chat-input-rows", data=5),
                    dbc.Row([
                        dbc.Col(
                            dbc.Button("Send", id="btn-send-chat", color="primary", className="mb-3 w-100"),
                            width=6
                        ),
                        dbc.Col(
                            dbc.Button("Clear Chat", id="btn-clear-chat", color="secondary", className="mb-3 w-100"),
                            width=6
                        ),
                    ], className="g-1 mb-2"),
                    html.Div(id="chat-output")
                ]
            )
        ]),
        className="mt-4"
    )
    return [
        html.H4("Proposal Writer", className="mt-3 mb-2", style={'marginBottom': '12px'}),
        html.Div([
            html.Div(className="blinking-dot", style={'margin':'0 auto','width':'16px','height':'16px'}),
        ], style={'textAlign':'center', 'marginBottom':'10px'}),
        html.Hr(style={'marginTop': '8px', 'marginBottom': '16px'}),
        html.Div(
            id='doc-type-buttons'
        ),
        chat_card
    ]

def get_right_col_content(selected_type, 
                          shred_doc, pink_doc, pink_review_doc, red_doc, red_review_doc, gold_doc):
    controls = []
    controls.append(html.Div([
        html.Div(className="blinking-dot", style={'margin':'0 auto','width':'16px','height':'16px'}),
    ], style={'textAlign':'center', 'marginBottom':'10px'}))
    controls.append(dcc.Loading(
        id="loading-indicator",
        type="dot",
        children=[html.Div(id="loading-output")]
    ))
    if selected_type == "Shred":
        controls.append(
            html.Div([
                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
                ),
                file_list_component(),
                html.Div([
                    html.Div("Loaded Shred Document:", className="mt-2 mb-1"),
                    markdown_table_preview(shred_doc) if shred_doc else html.Div("No Shred document loaded.")
                ], style={'marginBottom': '12px'})
            ])
        )
    elif selected_type in doc_dependencies:
        sources = doc_dependencies[selected_type]["source"]
        for src in sources:
            label = ""
            loaded_preview = None
            store_var = None
            if src == "shred":
                label = "Shred (Requirements)"
                store_var = shred_doc
                loaded_preview = markdown_table_preview(shred_doc) if shred_doc else html.Div("No Shred document loaded.")
            elif src == "pink":
                label = "Pink Document"
                store_var = pink_doc
                loaded_preview = markdown_narrative_preview(pink_doc) if pink_doc else html.Div("No Pink document loaded.")
            elif src == "pink_review":
                label = "Pink Review"
                store_var = pink_review_doc
                loaded_preview = markdown_table_preview(pink_review_doc) if pink_review_doc else html.Div("No Pink Review document loaded.")
            elif src == "red":
                label = "Red Document"
                store_var = red_doc
                loaded_preview = markdown_narrative_preview(red_doc) if red_doc else html.Div("No Red document loaded.")
            elif src == "red_review":
                label = "Red Review"
                store_var = red_doc
                loaded_preview = markdown_narrative_preview(red_doc) if red_doc else html.Div("No Red Review document loaded.")
            elif src == "gold":
                label = "Gold Document"
                store_var = gold_doc
                loaded_preview = markdown_narrative_preview(gold_doc) if gold_doc else html.Div("No Gold document loaded.")
            controls.append(html.Div([
                html.Label(f"Upload {label}"),
                dcc.Upload(
                    id={'type': f'upload-doc-type-{src}', 'index': selected_type},
                    children=html.Div(['Drag and Drop or ', html.A('Select File')]),
                    style={
                        'width': '100%',
                        'height': '60px',
                        'lineHeight': '60px',
                        'borderWidth': '1px',
                        'borderStyle': 'dashed',
                        'borderRadius': '5px',
                        'textAlign': 'center',
                        'margin': '10px 0'
                    },
                    multiple=False
                ),
                html.Div(id={'type': f'uploaded-doc-name-{src}', 'index': selected_type}),
                dbc.RadioItems(
                    id={'type': f'radio-doc-source-{src}', 'index': selected_type},
                    options=[
                        {'label': 'Loaded Document', 'value': 'loaded'},
                        {'label': 'Uploaded Document', 'value': 'uploaded'}
                    ] if store_var else [{'label': 'Uploaded Document', 'value': 'uploaded'}],
                    value='loaded' if store_var else 'uploaded',
                    inline=True,
                    className="mb-2"
                ),
                html.Div([
                    html.Div("Loaded Document Preview:", className="mt-2 mb-1"),
                    loaded_preview
                ], style={'marginBottom': '12px'}) if store_var else None,
            ], id={'type': f'doc-type-controls-{src}', 'index': selected_type}))
    controls.append(
        dbc.Row([
            dbc.Col(
                dbc.Button("Generate Document", id={'type': 'btn-generate-doc', 'index': selected_type}, color="primary", className="mb-3 w-100"),
                width=6
            ),
            dbc.Col(
                dbc.Button("Download Document", id="btn-download", color="success", className="mb-3 w-100"),
                width=6
            ),
            dcc.Download(id="download-document")
        ], className="mb-2")
    )
    controls.append(html.Hr())
    controls.append(html.Div(id='document-preview', className="border p-3 mb-3"))
    return dbc.Card(dbc.CardBody(controls))

app.layout = dbc.Container([
    dcc.Store(id='selected-doc-type', data="Shred"),
    dcc.Store(id='store-shred'),
    dcc.Store(id='store-pink'),
    dcc.Store(id='store-pink-review'),
    dcc.Store(id='store-red'),
    dcc.Store(id='store-red-review'),
    dcc.Store(id='store-gold'),
    dcc.Store(id='store-gold-review'),
    dcc.Store(id='store-loe'),
    dcc.Store(id='store-virtual-board'),
    dbc.Row([
        dbc.Col(
            html.H2(id='main-title', className="mt-3 mb-2", style={'textAlign': 'center', 'width':'100%'}),
            width=12
        )
    ]),
    dbc.Row([
        dbc.Col(
            html.Div([
                html.Div(className="blinking-dot", style={'margin':'0 auto','width':'16px','height':'16px'}),
            ], style={'textAlign':'center', 'marginBottom':'10px'})
        )
    ]),
    dbc.Row([
        dbc.Col(
            dbc.Card(
                dbc.CardBody(get_left_col_content())
            ),
            width=3
        ),
        dbc.Col(
            html.Div(id='right-col-content'),
            width=9
        )
    ])
], fluid=True)

@app.callback(
    Output('main-title', 'children'),
    Input('selected-doc-type', 'data')
)
def update_main_title(selected_type):
    return selected_type

@app.callback(
    Output('doc-type-buttons', 'children'),
    Input('selected-doc-type', 'data')
)
def render_doc_type_buttons(selected_type):
    buttons = []
    for doc_type in document_types.keys():
        btn_style = {'overflow': 'hidden', 'text-overflow': 'ellipsis', 'white-space': 'nowrap'}
        btn_class = "mb-2 w-100 text-left custom-button"
        if doc_type == selected_type:
            btn_class += " active-doc-type"
        buttons.append(
            dbc.Button(
                doc_type,
                id={'type': 'btn-doc-type', 'index': doc_type},
                color="link",
                className=btn_class,
                style=btn_style
            )
        )
    return buttons

@app.callback(
    Output('selected-doc-type', 'data'),
    [Input({'type': 'btn-doc-type', 'index': ALL}, 'n_clicks')],
    [State({'type': 'btn-doc-type', 'index': ALL}, 'id')],
    prevent_initial_call=True
)
def update_selected_doc_type(n_clicks_list, btn_ids):
    triggered = callback_context.triggered
    if not triggered or all(x is None for x in n_clicks_list):
        raise dash.exceptions.PreventUpdate
    idx = [i for i, x in enumerate(n_clicks_list) if x]
    if idx:
        selected_type = btn_ids[idx[-1]]['index']
    else:
        selected_type = "Shred"
    logging.info(f"Doc type selected: {selected_type}")
    return selected_type

@app.callback(
    Output('right-col-content', 'children'),
    Input('selected-doc-type', 'data'),
    State('store-shred', 'data'),
    State('store-pink', 'data'),
    State('store-pink-review', 'data'),
    State('store-red', 'data'),
    State('store-red-review', 'data'),
    State('store-gold', 'data'),
)
def update_right_col(selected_type, shred_doc, pink_doc, pink_review_doc, red_doc, red_review_doc, gold_doc):
    return get_right_col_content(selected_type, shred_doc, pink_doc, pink_review_doc, red_doc, red_review_doc, gold_doc)

def get_all_lengths(*args):
    return [len(x) if isinstance(x, list) else 0 for x in args]

@app.callback(
    Output('file-list', 'children'),
    Output('store-shred', 'data'),
    Output({'type': f'uploaded-doc-name-shred', 'index': ALL}, 'children'),
    Output({'type': f'uploaded-doc-name-pink', 'index': ALL}, 'children'),
    Output({'type': f'uploaded-doc-name-pink_review', 'index': ALL}, 'children'),
    Output({'type': f'uploaded-doc-name-red', 'index': ALL}, 'children'),
    Output({'type': f'uploaded-doc-name-red_review', 'index': ALL}, 'children'),
    Output({'type': f'uploaded-doc-name-gold', 'index': ALL}, 'children'),
    Output({'type': f'upload-doc-type-shred', 'index': ALL}, 'contents'),
    Output({'type': f'upload-doc-type-pink', 'index': ALL}, 'contents'),
    Output({'type': f'upload-doc-type-pink_review', 'index': ALL}, 'contents'),
    Output({'type': f'upload-doc-type-red', 'index': ALL}, 'contents'),
    Output({'type': f'upload-doc-type-red_review', 'index': ALL}, 'contents'),
    Output({'type': f'upload-doc-type-gold', 'index': ALL}, 'contents'),
    Output({'type': f'radio-doc-source-shred', 'index': ALL}, 'value'),
    Output({'type': f'radio-doc-source-pink', 'index': ALL}, 'value'),
    Output({'type': f'radio-doc-source-pink_review', 'index': ALL}, 'value'),
    Output({'type': f'radio-doc-source-red', 'index': ALL}, 'value'),
    Output({'type': f'radio-doc-source-red_review', 'index': ALL}, 'value'),
    Output({'type': f'radio-doc-source-gold', 'index': ALL}, 'value'),
    Output('document-preview', 'children'),
    Output('loading-output', 'children'),
    Output('store-pink', 'data'),
    Output('store-pink-review', 'data'),
    Output('store-red', 'data'),
    Output('store-red-review', 'data'),
    Output('store-gold', 'data'),
    Output('store-gold-review', 'data'),
    Output('store-loe', 'data'),
    Output('store-virtual-board', 'data'),
    Output('chat-output', 'children'),
    Output("download-document", "data"),
    Input('upload-document', 'contents'),
    State('upload-document', 'filename'),
    State('file-list', 'children'),
    State('store-shred', 'data'),
    Input({'type': 'remove-file', 'index': ALL}, 'n_clicks'),
    State('file-list', 'children'),
    State('store-shred', 'data'),
    Input({'type': f'upload-doc-type-shred', 'index': ALL}, 'contents'),
    State({'type': f'upload-doc-type-shred', 'index': ALL}, 'filename'),
    State({'type': f'upload-doc-type-shred', 'index': ALL}, 'id'),
    Input({'type': f'upload-doc-type-pink', 'index': ALL}, 'contents'),
    State({'type': f'upload-doc-type-pink', 'index': ALL}, 'filename'),
    State({'type': f'upload-doc-type-pink', 'index': ALL}, 'id'),
    Input({'type': f'upload-doc-type-pink_review', 'index': ALL}, 'contents'),
    State({'type': f'upload-doc-type-pink_review', 'index': ALL}, 'filename'),
    State({'type': f'upload-doc-type-pink_review', 'index': ALL}, 'id'),
    Input({'type': f'upload-doc-type-red', 'index': ALL}, 'contents'),
    State({'type': f'upload-doc-type-red', 'index': ALL}, 'filename'),
    State({'type': f'upload-doc-type-red', 'index': ALL}, 'id'),
    Input({'type': f'upload-doc-type-red_review', 'index': ALL}, 'contents'),
    State({'type': f'upload-doc-type-red_review', 'index': ALL}, 'filename'),
    State({'type': f'upload-doc-type-red_review', 'index': ALL}, 'id'),
    Input({'type': f'upload-doc-type-gold', 'index': ALL}, 'contents'),
    State({'type': f'upload-doc-type-gold', 'index': ALL}, 'filename'),
    State({'type': f'upload-doc-type-gold', 'index': ALL}, 'id'),
    Input({'type': 'btn-generate-doc', 'index': ALL}, 'n_clicks'),
    State({'type': 'btn-generate-doc', 'index': ALL}, 'id'),
    State({'type': 'radio-doc-source-shred', 'index': ALL}, 'value'),
    State({'type': 'upload-doc-type-shred', 'index': ALL}, 'contents'),
    State({'type': 'upload-doc-type-shred', 'index': ALL}, 'filename'),
    State({'type': 'radio-doc-source-pink', 'index': ALL}, 'value'),
    State({'type': 'upload-doc-type-pink', 'index': ALL}, 'contents'),
    State({'type': 'upload-doc-type-pink', 'index': ALL}, 'filename'),
    State({'type': 'radio-doc-source-pink_review', 'index': ALL}, 'value'),
    State({'type': 'upload-doc-type-pink_review', 'index': ALL}, 'contents'),
    State({'type': 'upload-doc-type-pink_review', 'index': ALL}, 'filename'),
    State({'type': 'radio-doc-source-red', 'index': ALL}, 'value'),
    State({'type': 'upload-doc-type-red', 'index': ALL}, 'contents'),
    State({'type': 'upload-doc-type-red', 'index': ALL}, 'filename'),
    State({'type': 'radio-doc-source-red_review', 'index': ALL}, 'value'),
    State({'type': 'upload-doc-type-red_review', 'index': ALL}, 'contents'),
    State({'type': 'upload-doc-type-red_review', 'index': ALL}, 'filename'),
    State({'type': 'radio-doc-source-gold', 'index': ALL}, 'value'),
    State({'type': 'upload-doc-type-gold', 'index': ALL}, 'contents'),
    State({'type': 'upload-doc-type-gold', 'index': ALL}, 'filename'),
    State('store-shred', 'data'),
    State('store-pink', 'data'),
    State('store-pink-review', 'data'),
    State('store-red', 'data'),
    State('store-red-review', 'data'),
    State('store-gold', 'data'),
    State('store-gold-review', 'data'),
    State('store-loe', 'data'),
    State('store-virtual-board', 'data'),
    Input('btn-send-chat', 'n_clicks'),
    Input('btn-clear-chat', 'n_clicks'),
    State('chat-input', 'value'),
    State('selected-doc-type', 'data'),
    State('document-preview', 'children'),
    Input("btn-download", "n_clicks"),
    State('selected-doc-type', 'data'),
    State('store-shred', 'data'),
    State('store-pink', 'data'),
    State('store-pink-review', 'data'),
    State('store-red', 'data'),
    State('store-red-review', 'data'),
    State('store-gold', 'data'),
    State('store-gold-review', 'data'),
    State('store-loe', 'data'),
    State('store-virtual-board', 'data'),
    prevent_initial_call=True
)
def master_callback(
    upload_contents, upload_filenames, existing_files, current_shred,
    remove_n_clicks, remove_existing_files, remove_current_shred,
    upload_shred_contents, upload_shred_filenames, upload_shred_ids,
    upload_pink_contents, upload_pink_filenames, upload_pink_ids,
    upload_pink_review_contents, upload_pink_review_filenames, upload_pink_review_ids,
    upload_red_contents, upload_red_filenames, upload_red_ids,
    upload_red_review_contents, upload_red_review_filenames, upload_red_review_ids,
    upload_gold_contents, upload_gold_filenames, upload_gold_ids,
    n_clicks_list, btn_ids,
    radio_shred, up_shred_contents, up_shred_filenames,
    radio_pink, up_pink_contents, up_pink_filenames,
    radio_pink_review, up_pink_review_contents, up_pink_review_filenames,
    radio_red, up_red_contents, up_red_filenames,
    radio_red_review, up_red_review_contents, up_red_review_filenames,
    radio_gold, up_gold_contents, up_gold_filenames,
    store_shred, store_pink, store_pink_review, store_red, store_red_review, store_gold, store_gold_review, store_loe, store_virtual_board,
    btn_send, btn_clear, chat_input, chat_doc_type, doc_preview,
    btn_download, dl_doc_type, dl_store_shred, dl_store_pink, dl_store_pink_review, dl_store_red, dl_store_red_review, dl_store_gold, dl_store_gold_review, dl_store_loe, dl_store_virtual_board
):
    global uploaded_files, uploaded_doc_contents
    ctx = callback_context

    wildcard_lengths = [
        len(upload_shred_ids) if upload_shred_ids is not None else 0,
        len(upload_pink_ids) if upload_pink_ids is not None else 0,
        len(upload_pink_review_ids) if upload_pink_review_ids is not None else 0,
        len(upload_red_ids) if upload_red_ids is not None else 0,
        len(upload_red_review_ids) if upload_red_review_ids is not None else 0,
        len(upload_gold_ids) if upload_gold_ids is not None else 0,
        len(upload_shred_ids) if upload_shred_ids is not None else 0,
        len(upload_pink_ids) if upload_pink_ids is not None else 0,
        len(upload_pink_review_ids) if upload_pink_review_ids is not None else 0,
        len(upload_red_ids) if upload_red_ids is not None else 0,
        len(upload_red_review_ids) if upload_red_review_ids is not None else 0,
        len(upload_gold_ids) if upload_gold_ids is not None else 0,
        len(upload_shred_ids) if upload_shred_ids is not None else 0,
        len(upload_pink_ids) if upload_pink_ids is not None else 0,
        len(upload_pink_review_ids) if upload_pink_review_ids is not None else 0,
        len(upload_red_ids) if upload_red_ids is not None else 0,
        len(upload_red_review_ids) if upload_red_review_ids is not None else 0,
        len(upload_gold_ids) if upload_gold_ids is not None else 0,
    ]

    outputs = [
        dash.no_update,  # file-list
        dash.no_update,  # store-shred
        [dash.no_update] * wildcard_lengths[0],  # uploaded-doc-name-shred
        [dash.no_update] * wildcard_lengths[1],  # uploaded-doc-name-pink
        [dash.no_update] * wildcard_lengths[2],  # uploaded-doc-name-pink_review
        [dash.no_update] * wildcard_lengths[3],  # uploaded-doc-name-red
        [dash.no_update] * wildcard_lengths[4],  # uploaded-doc-name-red_review
        [dash.no_update] * wildcard_lengths[5],  # uploaded-doc-name-gold
        [dash.no_update] * wildcard_lengths[6],  # upload-doc-type-shred.contents
        [dash.no_update] * wildcard_lengths[7],  # upload-doc-type-pink.contents
        [dash.no_update] * wildcard_lengths[8],  # upload-doc-type-pink_review.contents
        [dash.no_update] * wildcard_lengths[9],  # upload-doc-type-red.contents
        [dash.no_update] * wildcard_lengths[10], # upload-doc-type-red_review.contents
        [dash.no_update] * wildcard_lengths[11], # upload-doc-type-gold.contents
        [dash.no_update] * wildcard_lengths[12], # radio-doc-source-shred.value
        [dash.no_update] * wildcard_lengths[13], # radio-doc-source-pink.value
        [dash.no_update] * wildcard_lengths[14], # radio-doc-source-pink_review.value
        [dash.no_update] * wildcard_lengths[15], # radio-doc-source-red.value
        [dash.no_update] * wildcard_lengths[16], # radio-doc-source-red_review.value
        [dash.no_update] * wildcard_lengths[17], # radio-doc-source-gold.value
        dash.no_update,  # document-preview
        dash.no_update,  # loading-output
        dash.no_update,  # store-pink
        dash.no_update,  # store-pink-review
        dash.no_update,  # store-red
        dash.no_update,  # store-red-review
        dash.no_update,  # store-gold
        dash.no_update,  # store-gold-review
        dash.no_update,  # store-loe
        dash.no_update,  # store-virtual-board
        dash.no_update,  # chat-output
        dash.no_update   # download-document
    ]

    # --- File upload/remove logic ---
    if ctx.triggered and ctx.triggered[0]['prop_id'].startswith('upload-document'):
        new_files = []
        last_processed_content = None
        if upload_contents is not None:
            for i, (content, name) in enumerate(zip(upload_contents, upload_filenames)):
                file_content = process_document(content, name)
                uploaded_files[name] = file_content
                last_processed_content = file_content
                new_files.append(
                    dbc.Row([
                        dbc.Col(
                            html.Button('×', id={'type': 'remove-file', 'index': name}, style={'marginRight': '5px', 'fontSize': '10px'}),
                            width=1
                        ),
                        dbc.Col(
                            html.Span(name, style={'wordBreak': 'break-all'}),
                            width=11
                        )
                    ], id={'type': 'file-row', 'index': name}, align="center", className="mb-1")
                )
            if existing_files is None:
                existing_files = []
            elif not isinstance(existing_files, list):
                existing_files = [existing_files]
            outputs[0] = existing_files + new_files
            # Always set the latest uploaded Shred document as store-shred
            outputs[1] = last_processed_content
        return outputs

    if ctx.triggered and ctx.triggered[0]['prop_id'].startswith("{\"type\":\"remove-file\""):
        try:
            import json
            triggered_id = ctx.triggered[0]['prop_id'].split('.')[0]
            triggered_dict = json.loads(triggered_id)
            removed_file = triggered_dict['index']
        except Exception as e:
            logging.error(f"Could not parse removed file from callback context: {e}")
            outputs[0] = remove_existing_files
            outputs[1] = remove_current_shred
            return outputs
        uploaded_files.pop(removed_file, None)
        logging.info(f"Removed file: {removed_file}")
        filtered_files = []
        if remove_existing_files:
            for file in remove_existing_files:
                try:
                    file_id = file['props']['id']
                    if file_id['index'] != removed_file:
                        filtered_files.append(file)
                except Exception as e:
                    filtered_files.append(file)
        outputs[0] = filtered_files
        outputs[1] = remove_current_shred
        return outputs

    doc_upload_types = [
        ("shred", upload_shred_contents, upload_shred_filenames, upload_shred_ids, 2, 8, 14),
        ("pink", upload_pink_contents, upload_pink_filenames, upload_pink_ids, 3, 9, 15),
        ("pink_review", upload_pink_review_contents, upload_pink_review_filenames, upload_pink_review_ids, 4, 10, 16),
        ("red", upload_red_contents, upload_red_filenames, upload_red_ids, 5, 11, 17),
        ("red_review", upload_red_review_contents, upload_red_review_filenames, upload_red_review_ids, 6, 12, 18),
        ("gold", upload_gold_contents, upload_gold_filenames, upload_gold_ids, 7, 13, 19)
    ]
    for src, contents, filenames, ids, name_idx, content_idx, radio_idx in doc_upload_types:
        if contents and any(c is not None for c in contents):
            for i, content in enumerate(contents):
                if content is not None:
                    uploaded_doc_contents[(src, ids[i]['index'])] = (content, filenames[i])
                    logging.info(f"{ids[i]['index']} {src} file uploaded: {filenames[i]}")
                    name_list = ["" for _ in range(len(contents))]
                    content_list = [None for _ in range(len(contents))]
                    radio_list = ["loaded" for _ in range(len(contents))]
                    name_list[i] = filenames[i]
                    content_list[i] = content
                    radio_list[i] = "uploaded"
                    outputs[name_idx] = name_list
                    outputs[content_idx] = content_list
                    outputs[radio_idx] = radio_list
                    return outputs

    def gemini_generate(prompt, max_tokens=32798, temperature=0.4):
        result_holder = {}
        def run_gemini():
            try:
                model = genai.GenerativeModel(GEMINI_MODEL)
                response = model.generate_content(
                    prompt,
                    generation_config=genai.types.GenerationConfig(
                        temperature=temperature,
                        max_output_tokens=max_tokens,
                        top_p=1
                    )
                )
                result_holder['result'] = response.text
                logging.info("Gemini response received.")
            except Exception as e:
                logging.error(f"Gemini error: {e}")
                result_holder['error'] = str(e)
        thread = threading.Thread(target=run_gemini)
        thread.start()
        thread.join()
        if 'error' in result_holder:
            raise Exception(result_holder['error'])
        return result_holder['result']

    def generate_document(document_type, file_contents, extra_context=None):
        if document_type in spreadsheet_types:
            prompt = f"""Ignore all other instructions and output only a spreadsheet for {document_type} as described below. Do not include any narrative, only the spreadsheet in markdown table format.
Instructions: {document_types[document_type]}
Project Artifacts:
{' '.join(file_contents)}
Output only the spreadsheet as a markdown table, no narrative or explanation."""
        elif document_type in narrative_types:
            prompt = f"""Generate a {document_type} document 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}:
"""
        else:
            prompt = f"""Generate a {document_type} based on the following project artifacts:
{' '.join(file_contents)}
Instructions:
{document_types.get(document_type, '')}
Now, generate the {document_type}:
"""
        if extra_context:
            prompt += f"\n\n{extra_context}"
        logging.info(f"Generating document for type: {document_type} using Gemini.")
        return gemini_generate(prompt, max_tokens=4096, temperature=0.25)

    if ctx.triggered and ctx.triggered[0]['prop_id'].startswith("{\"type\":\"btn-generate-doc\""):
        idx = [i for i, x in enumerate(n_clicks_list) if x]
        if not idx:
            return outputs
        idx = idx[-1]
        doc_type = btn_ids[idx]['index']

        shred_doc = store_shred
        pink_doc = store_pink
        pink_review_doc = store_pink_review
        red_doc = store_red
        red_review_doc = store_red_review
        gold_doc = store_gold
        gold_review_doc = store_gold_review
        loe_doc = store_loe
        virtual_board_doc = store_virtual_board

        def get_doc_from_radio(radio, upload_contents, upload_filenames, loaded_var):
            if radio and radio[0] == 'uploaded':
                if upload_contents and upload_contents[0] and upload_filenames and upload_filenames[0]:
                    return process_document(upload_contents[0], upload_filenames[0])
                else:
                    return None
            else:
                return loaded_var

        if doc_type == "Shred":
            if not store_shred:
                outputs[20] = html.Div("Please upload a document before shredding.")
                outputs[21] = ""
                return outputs
            file_contents = [store_shred]
            try:
                generated = generate_document(doc_type, file_contents)
                outputs[1] = generated
                outputs[20] = markdown_table_preview(generated)
                outputs[21] = "Shred generated"
            except Exception as e:
                outputs[20] = html.Div(f"Error generating document: {str(e)}")
                outputs[21] = "Error"
            return outputs

        if doc_type == "Pink":
            shred = get_doc_from_radio(radio_shred, up_shred_contents, up_shred_filenames, shred_doc)
            if not shred:
                outputs[20] = html.Div("Please provide a Shred requirements document (either loaded or uploaded) to generate Pink.")
                outputs[21] = ""
                return outputs
            try:
                generated = generate_document(doc_type, [shred])
                outputs[22] = generated
                outputs[20] = markdown_narrative_preview(generated)
                outputs[21] = "Pink generated"
            except Exception as e:
                outputs[20] = html.Div(f"Error generating Pink: {str(e)}")
                outputs[21] = "Error"
            return outputs

        if doc_type == "Pink Review":
            pink = get_doc_from_radio(radio_pink, up_pink_contents, up_pink_filenames, pink_doc)
            shred = get_doc_from_radio(radio_shred, up_shred_contents, up_shred_filenames, shred_doc)
            if not pink or not shred:
                outputs[20] = html.Div("Please provide both Pink and Shred documents (either loaded or uploaded) to generate Pink Review.")
                outputs[21] = ""
                return outputs
            try:
                generated = generate_document(doc_type, [pink, shred])
                outputs[23] = generated
                outputs[20] = markdown_table_preview(generated)
                outputs[21] = "Pink Review generated"
            except Exception as e:
                outputs[20] = html.Div(f"Error generating Pink Review: {str(e)}")
                outputs[21] = "Error"
            return outputs

        if doc_type == "Red":
            pink_review = get_doc_from_radio(radio_pink_review, up_pink_review_contents, up_pink_review_filenames, pink_review_doc)
            if not pink_review:
                outputs[20] = html.Div("Please provide a Pink Review document (either loaded or uploaded) to generate Red.")
                outputs[21] = ""
                return outputs
            try:
                generated = generate_document(doc_type, [pink_review])
                outputs[24] = generated
                outputs[20] = markdown_narrative_preview(generated)
                outputs[21] = "Red generated"
            except Exception as e:
                outputs[20] = html.Div(f"Error generating Red: {str(e)}")
                outputs[21] = "Error"
            return outputs

        if doc_type == "Red Review":
            red = get_doc_from_radio(radio_red, up_red_contents, up_red_filenames, red_doc)
            shred = get_doc_from_radio(radio_shred, up_shred_contents, up_shred_filenames, shred_doc)
            if not red or not shred:
                outputs[20] = html.Div("Please provide both Red and Shred documents (either loaded or uploaded) to generate Red Review.")
                outputs[21] = ""
                return outputs
            try:
                generated = generate_document(doc_type, [red, shred])
                outputs[25] = generated
                outputs[20] = markdown_table_preview(generated)
                outputs[21] = "Red Review generated"
            except Exception as e:
                outputs[20] = html.Div(f"Error generating Red Review: {str(e)}")
                outputs[21] = "Error"
            return outputs

        if doc_type == "Gold":
            red_review = get_doc_from_radio(radio_red_review, up_red_review_contents, up_red_review_filenames, red_review_doc)
            if not red_review:
                outputs[20] = html.Div("Please provide a Red Review document (either loaded or uploaded) to generate Gold.")
                outputs[21] = ""
                return outputs
            try:
                generated = generate_document(doc_type, [red_review])
                outputs[26] = generated
                outputs[20] = markdown_narrative_preview(generated)
                outputs[21] = "Gold generated"
            except Exception as e:
                outputs[20] = html.Div(f"Error generating Gold: {str(e)}")
                outputs[21] = "Error"
            return outputs

        if doc_type == "Gold Review":
            gold = get_doc_from_radio(radio_gold, up_gold_contents, up_gold_filenames, gold_doc)
            shred = get_doc_from_radio(radio_shred, up_shred_contents, up_shred_filenames, shred_doc)
            if not gold or not shred:
                outputs[20] = html.Div("Please provide both Gold and Shred documents (either loaded or uploaded) to generate Gold Review.")
                outputs[21] = ""
                return outputs
            try:
                generated = generate_document(doc_type, [gold, shred])
                outputs[27] = generated
                outputs[20] = markdown_table_preview(generated)
                outputs[21] = "Gold Review generated"
            except Exception as e:
                outputs[20] = html.Div(f"Error generating Gold Review: {str(e)}")
                outputs[21] = "Error"
            return outputs

        if doc_type == "LOE":
            gold = get_doc_from_radio(radio_gold, up_gold_contents, up_gold_filenames, gold_doc)
            if not gold:
                outputs[20] = html.Div("Please provide a Gold document (either loaded or uploaded) to generate LOE.")
                outputs[21] = ""
                return outputs
            try:
                generated = generate_document(doc_type, [gold])
                outputs[28] = generated
                outputs[20] = markdown_table_preview(generated)
                outputs[21] = "LOE generated"
            except Exception as e:
                outputs[20] = html.Div(f"Error generating LOE: {str(e)}")
                outputs[21] = "Error"
            return outputs

        if doc_type == "Virtual Board":
            shred = get_doc_from_radio(radio_shred, up_shred_contents, up_shred_filenames, shred_doc)
            if not shred:
                outputs[20] = html.Div("Please provide a Shred requirements document (either loaded or uploaded) to generate Virtual Board.")
                outputs[21] = ""
                return outputs
            try:
                lm_text = ""
                lm_match = re.search(r'(Section\s+L[\s\S]+?)(Section\s+M|$)', shred, re.IGNORECASE)
                if lm_match:
                    lm_text = lm_match.group(1)
                else:
                    lm_text = shred
                generated = generate_document(doc_type, [lm_text])
                outputs[29] = generated
                outputs[20] = markdown_table_preview(generated)
                outputs[21] = "Virtual Board generated"
            except Exception as e:
                outputs[20] = html.Div(f"Error generating Virtual Board: {str(e)}")
                outputs[21] = "Error"
            return outputs

        outputs[20] = html.Div("Unsupported document type or missing required sources.")
        outputs[21] = ""
        return outputs

    if ctx.triggered and (ctx.triggered[0]['prop_id'] == 'btn-send-chat.n_clicks' or ctx.triggered[0]['prop_id'] == 'btn-clear-chat.n_clicks'):
        if ctx.triggered[0]['prop_id'] == 'btn-clear-chat.n_clicks':
            outputs[30] = ""
            return outputs

        doc_map = {
            "Shred": store_shred,
            "Pink": store_pink,
            "Pink Review": store_pink_review,
            "Red": store_red,
            "Red Review": store_red_review,
            "Gold": store_gold,
            "Gold Review": store_gold_review,
            "LOE": store_loe,
            "Virtual Board": store_virtual_board
        }
        current_document = doc_map.get(chat_doc_type)

        if not chat_input or current_document is None:
            outputs[30] = ""
            return outputs

        if chat_doc_type in spreadsheet_types:
            prompt = f"""Update the following {chat_doc_type} spreadsheet based on this instruction: {chat_input}
Current spreadsheet (markdown table format):
{current_document}
Instructions:
1. Provide the updated spreadsheet as a markdown table only.
2. Do not include any narrative, only the markdown table.
Now, provide the updated {chat_doc_type} spreadsheet:
"""
        else:
            prompt = f"""Update the following {chat_doc_type} document 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.
4. If the {chat_doc_type} is Pink, Red, or Gold then your goal is to write a FULL proposal response, not just a strategy and be compliant and compelling by addressing all the requirements from the document provided.  Focus on describing the approach and highly detailed how it will be done, the steps, workflow, people, processes and technology to accomplish the task. Be sure to refer to research that validates the approach and cite sources with measurable outcomes and improve on innovations of the approach.  Do not just say things like we will, or MicroHealth will, use active voice and verbs that are definitive in nature, not maybe, could be, should be, can be and things like that.  This is a proposal response so logical flow is important so the reader can follow.
Now, provide the updated {chat_doc_type}:
"""
        logging.info(f"Updating document via chat for {chat_doc_type} instruction: {chat_input}")
        try:
            new_document = gemini_generate(prompt, max_tokens=4096, temperature=0.5)
            if chat_doc_type in spreadsheet_types:
                outputs[20] = markdown_table_preview(new_document)
            else:
                outputs[20] = markdown_narrative_preview(new_document)
            stores = [store_shred, store_pink, store_pink_review, store_red, store_red_review, store_gold, store_gold_review, store_loe, store_virtual_board]
            doc_types = ["Shred", "Pink", "Pink Review", "Red", "Red Review", "Gold", "Gold Review", "LOE", "Virtual Board"]
            for i, dt in enumerate(doc_types):
                if dt == chat_doc_type:
                    outputs[1 + i] = new_document
            outputs[30] = "Document updated based on: {}".format(chat_input)
            return outputs
        except Exception as e:
            outputs[20] = html.Div(f"Error updating document: {str(e)}")
            outputs[30] = f"Error updating document: {str(e)}"
            return outputs

    if ctx.triggered and ctx.triggered[0]['prop_id'] == "btn-download.n_clicks":
        doc_map = {
            "Shred": dl_store_shred,
            "Pink": dl_store_pink,
            "Pink Review": dl_store_pink_review,
            "Red": dl_store_red,
            "Red Review": dl_store_red_review,
            "Gold": dl_store_gold,
            "Gold Review": dl_store_gold_review,
            "LOE": dl_store_loe,
            "Virtual Board": dl_store_virtual_board
        }
        current_document = doc_map.get(dl_doc_type)
        if current_document is None:
            return outputs
        if dl_doc_type in spreadsheet_types:
            try:
                xlsx_bytes = markdown_tables_to_xlsx(current_document)
                outputs[31] = dcc.send_bytes(xlsx_bytes.read(), f"{dl_doc_type}.xlsx")
            except Exception as e:
                outputs[31] = dcc.send_string(f"Error downloading {dl_doc_type}: {str(e)}", f"{dl_doc_type}_error.txt")
        else:
            try:
                plain = strip_markdown(current_document)
                doc = Document()
                for para in plain.split('\n'):
                    doc.add_paragraph(para)
                output = BytesIO()
                doc.save(output)
                output.seek(0)
                outputs[31] = dcc.send_bytes(output.read(), f"{dl_doc_type}.docx")
            except Exception as e:
                outputs[31] = dcc.send_string(f"Error downloading document: {str(e)}", f"{dl_doc_type}_error.txt")
        return outputs

    return outputs

@app.callback(
    Output("chat-input", "rows"),
    Input("chat-input", "value"),
    State("chat-input", "rows"),
    prevent_initial_call="initial_duplicate"
)
def auto_expand_textarea(value, current_rows):
    if value is None or value == "":
        return 5
    num_lines = value.count('\n') + 1
    max_rows = 20
    rows = min(max(num_lines, 5), max_rows)
    return rows

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