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 openai
import logging
import threading
import re 
import markdown
from bs4 import BeautifulSoup

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

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

openai.api_key = os.environ.get("OPENAI_API_KEY", "")

uploaded_files = {}
current_document = None
document_type = None
shredded_document = None
pink_review_document = None
pink_document = None
red_review_document = None
red_document = None
gold_review_document = None
gold_document = None
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"
}

# Document dependency rules: mapping doc type to allowed sources for generation and review
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 get_right_col_content(selected_type):
    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")]
    ))
    # Shred: can always upload
    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()
            ])
        )
    # For all proposal/response/review types, show only allowed upload controls
    elif selected_type in doc_dependencies:
        sources = doc_dependencies[selected_type]["source"]
        # For each source, show upload and radio if loaded/generated exists
        for src in sources:
            label = ""
            store_var = None
            if src == "shred":
                label = "Shred (Requirements)"
                store_var = shredded_document
            elif src == "pink":
                label = "Pink Document"
                store_var = pink_document
            elif src == "pink_review":
                label = "Pink Review"
                store_var = pink_review_document
            elif src == "red":
                label = "Red Document"
                store_var = red_document
            elif src == "red_review":
                label = "Red Review"
                store_var = red_review_document
            elif src == "gold":
                label = "Gold Document"
                store_var = gold_document
            # Upload control for this source
            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"
                ),
            ], 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))

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-4"),
        html.Div([
            html.Div(className="blinking-dot", style={'margin':'0 auto','width':'16px','height':'16px'}),
        ], style={'textAlign':'center', 'marginBottom':'10px'}),
        html.Hr(),
        html.Div(
            id='doc-type-buttons'
        ),
        chat_card
    ]

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

app.layout = dbc.Container([
    dcc.Store(id='selected-doc-type', data="Shred"),
    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')
)
def update_right_col(selected_type):
    return get_right_col_content(selected_type)

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

@app.callback(
    Output('file-list', '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(
                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]
        shredded_document = None
        logging.info("Documents uploaded and file list updated.")
        return existing_files + new_files
    return existing_files

@app.callback(
    Output('file-list', 'children', allow_duplicate=True),
    Input({'type': 'remove-file', 'index': ALL}, 'n_clicks'),
    State('file-list', 'children'),
    prevent_initial_call=True
)
def remove_file(n_clicks, existing_files):
    global uploaded_files, shredded_document
    ctx = dash.callback_context
    if not ctx.triggered:
        raise dash.exceptions.PreventUpdate
    triggered_id = ctx.triggered[0]['prop_id'].split('.')[0]
    if not triggered_id.startswith("{\"type\":\"remove-file\""):
        raise dash.exceptions.PreventUpdate
    try:
        import json
        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}")
        raise dash.exceptions.PreventUpdate
    uploaded_files.pop(removed_file, None)
    shredded_document = None
    logging.info(f"Removed file: {removed_file}")
    filtered_files = []
    if existing_files:
        for file in 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)
    return filtered_files

# Upload callbacks for all doc dependencies (Pink, Pink Review, Red, Red Review, Gold, Gold Review, LOE, Virtual Board)
def make_upload_callback(src):
    @app.callback(
        Output({'type': f'uploaded-doc-name-{src}', 'index': MATCH}, 'children'),
        Output({'type': f'upload-doc-type-{src}', 'index': MATCH}, 'contents'),
        Output({'type': f'radio-doc-source-{src}', 'index': MATCH}, 'value'),
        Input({'type': f'upload-doc-type-{src}', 'index': MATCH}, 'contents'),
        State({'type': f'upload-doc-type-{src}', 'index': MATCH}, 'filename'),
        State({'type': f'upload-doc-type-{src}', 'index': MATCH}, 'id')
    )
    def update_uploaded_doc_name(contents, filename, id_dict):
        global uploaded_doc_contents
        if contents is not None:
            uploaded_doc_contents[(src, id_dict['index'])] = (contents, filename)
            logging.info(f"{id_dict['index']} {src} file uploaded: {filename}")
            return filename, contents, "uploaded"
        return "", None, "loaded"
for src in ["shred", "pink", "pink_review", "red", "red_review", "gold"]:
    make_upload_callback(src)

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 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}")
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4-1106-preview",
            messages=[
                {"role": "system", "content": "You are a helpful, expert government proposal writer."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=4096,
            temperature=0.25,
        )
        generated_text = response['choices'][0]['message']['content']
        logging.info("Document generated successfully.")
        return generated_text
    except Exception as e:
        logging.error(f"Error generating document: {str(e)}")
        raise

@app.callback(
    Output('document-preview', 'children'),
    Output('loading-output', 'children'),
    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'),
    prevent_initial_call=True
)
def generate_any_doc(
    n_clicks_list, btn_ids,
    radio_shred, upload_shred_contents, upload_shred_filenames,
    radio_pink, upload_pink_contents, upload_pink_filenames,
    radio_pink_review, upload_pink_review_contents, upload_pink_review_filenames,
    radio_red, upload_red_contents, upload_red_filenames,
    radio_red_review, upload_red_review_contents, upload_red_review_filenames,
    radio_gold, upload_gold_contents, upload_gold_filenames
):
    global current_document, document_type, shredded_document, pink_review_document, pink_document, red_review_document, red_document, gold_review_document, gold_document
    ctx = callback_context
    logging.info(f"generate_any_doc triggered: n_clicks_list={n_clicks_list}, btn_ids={btn_ids}")
    if not ctx.triggered:
        raise dash.exceptions.PreventUpdate
    idx = [i for i, x in enumerate(n_clicks_list) if x]
    if not idx:
        raise dash.exceptions.PreventUpdate
    idx = idx[-1]
    doc_type = btn_ids[idx]['index']
    document_type = doc_type

    # SHRED - always allowed
    if doc_type == "Shred":
        if not uploaded_files:
            logging.info("No uploaded files for Shred. Aborting.")
            return html.Div("Please upload a document before shredding."), ""
        file_contents = list(uploaded_files.values())
        try:
            generated = generate_document(doc_type, file_contents)
            current_document = generated
            shredded_document = generated
            preview = markdown_table_preview(generated)
            logging.info("Shred document generated.")
            return preview, "Shred generated"
        except Exception as e:
            logging.error(f"Error generating document: {str(e)}")
            return html.Div(f"Error generating document: {str(e)}"), "Error"

    # Dependency rules
    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

    # Pink: can only use Shred (requirements), loaded or uploaded
    if doc_type == "Pink":
        shred_doc = get_doc_from_radio(radio_shred, upload_shred_contents, upload_shred_filenames, shredded_document)
        if not shred_doc:
            return html.Div("Please provide a Shred requirements document (either loaded or uploaded) to generate Pink."), ""
        try:
            generated = generate_document(doc_type, [shred_doc])
            current_document = generated
            pink_document = generated
            preview = markdown_narrative_preview(generated)
            logging.info("Pink document generated.")
            return preview, "Pink generated"
        except Exception as e:
            logging.error(f"Error generating Pink: {str(e)}")
            return html.Div(f"Error generating Pink: {str(e)}"), "Error"

    # Pink Review: must compare Pink (uploaded or generated) with Shred (uploaded or generated)
    if doc_type == "Pink Review":
        pink_doc = get_doc_from_radio(radio_pink, upload_pink_contents, upload_pink_filenames, pink_document)
        shred_doc = get_doc_from_radio(radio_shred, upload_shred_contents, upload_shred_filenames, shredded_document)
        if not pink_doc or not shred_doc:
            return html.Div("Please provide both Pink and Shred documents (either loaded or uploaded) to generate Pink Review."), ""
        try:
            generated = generate_document(doc_type, [pink_doc, shred_doc])
            current_document = generated
            pink_review_document = generated
            preview = markdown_table_preview(generated)
            logging.info("Pink Review document generated.")
            return preview, "Pink Review generated"
        except Exception as e:
            logging.error(f"Error generating Pink Review: {str(e)}")
            return html.Div(f"Error generating Pink Review: {str(e)}"), "Error"

    # Red: can only use Pink Review (uploaded or generated)
    if doc_type == "Red":
        pink_review_doc = get_doc_from_radio(radio_pink_review, upload_pink_review_contents, upload_pink_review_filenames, pink_review_document)
        if not pink_review_doc:
            return html.Div("Please provide a Pink Review document (either loaded or uploaded) to generate Red."), ""
        try:
            generated = generate_document(doc_type, [pink_review_doc])
            current_document = generated
            red_document = generated
            preview = markdown_narrative_preview(generated)
            logging.info("Red document generated.")
            return preview, "Red generated"
        except Exception as e:
            logging.error(f"Error generating Red: {str(e)}")
            return html.Div(f"Error generating Red: {str(e)}"), "Error"

    # Red Review: must compare Red (uploaded or generated) with Shred (uploaded or generated)
    if doc_type == "Red Review":
        red_doc = get_doc_from_radio(radio_red, upload_red_contents, upload_red_filenames, red_document)
        shred_doc = get_doc_from_radio(radio_shred, upload_shred_contents, upload_shred_filenames, shredded_document)
        if not red_doc or not shred_doc:
            return html.Div("Please provide both Red and Shred documents (either loaded or uploaded) to generate Red Review."), ""
        try:
            generated = generate_document(doc_type, [red_doc, shred_doc])
            current_document = generated
            red_review_document = generated
            preview = markdown_table_preview(generated)
            logging.info("Red Review document generated.")
            return preview, "Red Review generated"
        except Exception as e:
            logging.error(f"Error generating Red Review: {str(e)}")
            return html.Div(f"Error generating Red Review: {str(e)}"), "Error"

    # Gold: can only use Red Review (uploaded or generated)
    if doc_type == "Gold":
        red_review_doc = get_doc_from_radio(radio_red_review, upload_red_review_contents, upload_red_review_filenames, red_review_document)
        if not red_review_doc:
            return html.Div("Please provide a Red Review document (either loaded or uploaded) to generate Gold."), ""
        try:
            generated = generate_document(doc_type, [red_review_doc])
            current_document = generated
            gold_document = generated
            preview = markdown_narrative_preview(generated)
            logging.info("Gold document generated.")
            return preview, "Gold generated"
        except Exception as e:
            logging.error(f"Error generating Gold: {str(e)}")
            return html.Div(f"Error generating Gold: {str(e)}"), "Error"

    # Gold Review: must compare Gold (uploaded or generated) with Shred (uploaded or generated)
    if doc_type == "Gold Review":
        gold_doc = get_doc_from_radio(radio_gold, upload_gold_contents, upload_gold_filenames, gold_document)
        shred_doc = get_doc_from_radio(radio_shred, upload_shred_contents, upload_shred_filenames, shredded_document)
        if not gold_doc or not shred_doc:
            return html.Div("Please provide both Gold and Shred documents (either loaded or uploaded) to generate Gold Review."), ""
        try:
            generated = generate_document(doc_type, [gold_doc, shred_doc])
            current_document = generated
            gold_review_document = generated
            preview = markdown_table_preview(generated)
            logging.info("Gold Review document generated.")
            return preview, "Gold Review generated"
        except Exception as e:
            logging.error(f"Error generating Gold Review: {str(e)}")
            return html.Div(f"Error generating Gold Review: {str(e)}"), "Error"

    # LOE: can only use Gold (uploaded or generated)
    if doc_type == "LOE":
        gold_doc = get_doc_from_radio(radio_gold, upload_gold_contents, upload_gold_filenames, gold_document)
        if not gold_doc:
            return html.Div("Please provide a Gold document (either loaded or uploaded) to generate LOE."), ""
        try:
            generated = generate_document(doc_type, [gold_doc])
            current_document = generated
            preview = markdown_table_preview(generated)
            logging.info("LOE document generated.")
            return preview, "LOE generated"
        except Exception as e:
            logging.error(f"Error generating LOE: {str(e)}")
            return html.Div(f"Error generating LOE: {str(e)}"), "Error"

    # Virtual Board: can only use Shred (L&M evaluation criteria)
    if doc_type == "Virtual Board":
        shred_doc = get_doc_from_radio(radio_shred, upload_shred_contents, upload_shred_filenames, shredded_document)
        if not shred_doc:
            return html.Div("Please provide a Shred requirements document (either loaded or uploaded) to generate Virtual Board."), ""
        try:
            # Extract section L&M if possible
            lm_text = ""
            lm_match = re.search(r'(Section\s+L[\s\S]+?)(Section\s+M|$)', shred_doc, re.IGNORECASE)
            if lm_match:
                lm_text = lm_match.group(1)
            else:
                lm_text = shred_doc
            generated = generate_document(doc_type, [lm_text])
            current_document = generated
            preview = markdown_table_preview(generated)
            logging.info("Virtual Board document generated.")
            return preview, "Virtual Board generated"
        except Exception as e:
            logging.error(f"Error generating Virtual Board: {str(e)}")
            return html.Div(f"Error generating Virtual Board: {str(e)}"), "Error"

    return html.Div("Unsupported document type or missing required sources."), ""

@app.callback(
    Output('chat-output', 'children'),
    Output('document-preview', 'children', allow_duplicate=True),
    Input('btn-send-chat', 'n_clicks'),
    Input('btn-clear-chat', 'n_clicks'),
    State('chat-input', 'value'),
    State('selected-doc-type', 'data'),
    prevent_initial_call=True
)
def update_document_via_chat(btn_send, btn_clear, chat_input, selected_doc_type):
    global current_document, document_type
    ctx = callback_context
    if not ctx.triggered:
        raise dash.exceptions.PreventUpdate
    trigger_id = ctx.triggered[0]['prop_id'].split('.')[0]
    if trigger_id == 'btn-clear-chat':
        return "", dash.no_update
    if not chat_input or current_document is None:
        raise dash.exceptions.PreventUpdate

    if selected_doc_type in spreadsheet_types:
        prompt = f"""Update the following {selected_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 {selected_doc_type} spreadsheet:
"""
    else:
        prompt = f"""Update the following {selected_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.
Now, provide the updated {selected_doc_type}:
"""
    logging.info(f"Updating document via chat for {selected_doc_type} instruction: {chat_input}")
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4-1106-preview",
            messages=[
                {"role": "system", "content": "You are a helpful, expert government proposal writer."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=32768,
            temperature=0.5,
        )
        current_document = response['choices'][0]['message']['content']
        logging.info("Document updated via chat successfully.")
        if selected_doc_type in spreadsheet_types:
            preview = markdown_table_preview(current_document)
        else:
            preview = markdown_narrative_preview(current_document)
        return f"Document updated based on: {chat_input}", preview
    except Exception as e:
        logging.error(f"Error updating document via chat: {str(e)}")
        return f"Error updating document: {str(e)}", html.Div(f"Error updating document: {str(e)}")

@app.callback(
    Output("download-document", "data"),
    Input("btn-download", "n_clicks"),
    State('selected-doc-type', 'data'),
    prevent_initial_call=True
)
def download_document(n_clicks, selected_doc_type):
    global current_document
    if current_document is None:
        raise dash.exceptions.PreventUpdate

    if selected_doc_type in spreadsheet_types:
        try:
            xlsx_bytes = markdown_tables_to_xlsx(current_document)
            logging.info(f"{selected_doc_type} document downloaded as Excel.")
            return dcc.send_bytes(xlsx_bytes.read(), f"{selected_doc_type}.xlsx")
        except Exception as e:
            logging.error(f"Error downloading {selected_doc_type} document: {str(e)}")
            return dcc.send_string(f"Error downloading {selected_doc_type}: {str(e)}", f"{selected_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)
            logging.info(f"{selected_doc_type} document downloaded as Word.")
            output.seek(0)
            return dcc.send_bytes(output.read(), f"{selected_doc_type}.docx")
        except Exception as e:
            logging.error(f"Error downloading document: {str(e)}")
            return dcc.send_string(f"Error downloading document: {str(e)}", f"{selected_doc_type}_error.txt")

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