app.py

import os
import base64
import io
import dash
from dash import dcc, html, Input, Output, State, callback_context, ALL
import dash_bootstrap_components as dbc
import dash.dash_table as dt
import pandas as pd
import logging
from docx import Document
import mimetypes
from threading import Lock
import tempfile
import shutil
import uuid
import re

import google.generativeai as genai

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

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

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

MAX_INPUT_TOKENS = 1048576
MAX_OUTPUT_TOKENS = 65536

SESSION_STORE = {}

PROMPT_FILES = {
    'shred': 'prompt_shred',
    'loe': 'prompt_LOE',
    'compliance': 'prompt_compliance',
    'proposal': 'prompt_proposal',
    'recover': 'prompt_recover',
    'virtual_board': 'prompt_virtual_board'
}

def read_prompt_file(prompt_type):
    fname = PROMPT_FILES.get(prompt_type)
    if not fname:
        return f"[Error: No prompt file mapping for {prompt_type}]"
    try:
        with open(fname, 'r', encoding='utf-8') as f:
            return f.read().strip()
    except Exception as e:
        logging.error(f"Could not read prompt file {fname}: {e}")
        return f"[Error: Could not read prompt file for {prompt_type}]"

def get_session_id_from_cookie(cookie_str):
    if not cookie_str:
        return None
    for part in cookie_str.split(";"):
        if part.strip().startswith("dash_session="):
            return part.strip().split("=")[1]
    return None

def get_session_id(session_id=None):
    if session_id and session_id in SESSION_STORE:
        return session_id
    if session_id:
        return session_id
    sid = str(uuid.uuid4())
    return sid

def get_session_data(session_id):
    if session_id not in SESSION_STORE:
        tempdir = tempfile.mkdtemp(prefix="rfp_session_")
        SESSION_STORE[session_id] = {
            "uploaded_documents": {},
            "uploaded_documents_fileid": {},
            "uploaded_documents_bytes": {},
            "proposals": {},
            "proposals_fileid": {},
            "shredded_documents": {},
            "generated_response": None,
            "gemini_lock": Lock(),
            "session_tempdir": tempdir
        }
    return SESSION_STORE[session_id]

def truncate_filename(filename, maxlen=30):
    if len(filename) <= maxlen:
        return filename
    else:
        partlen = (maxlen - 3) // 2
        return filename[:partlen] + "..." + filename[-partlen:]

def decode_document(decoded_bytes):
    try:
        content = decoded_bytes.decode('utf-8')
        logging.info("Document decoded as UTF-8.")
        return content
    except UnicodeDecodeError as e_utf8:
        try:
            content = decoded_bytes.decode('latin-1')
            logging.warning("Document decoded as Latin-1 due to utf-8 decode error: %s", e_utf8)
            return content
        except Exception as e:
            logging.error("Document decode failed for both utf-8 and latin-1: %s", e)
            return None

def guess_mime_type(filename):
    mime_type, _ = mimetypes.guess_type(filename)
    if not mime_type:
        ext = filename.lower().split('.')[-1]
        if ext == "pdf":
            mime_type = "application/pdf"
        elif ext == "docx":
            mime_type = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        elif ext == "doc":
            mime_type = "application/msword"
        elif ext == "xlsx":
            mime_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        elif ext == "xls":
            mime_type = "application/vnd.ms-excel"
        else:
            mime_type = "application/octet-stream"
    return mime_type

def upload_to_gemini_file(decoded_bytes, filename):
    try:
        mime_type = guess_mime_type(filename)
        file_obj = io.BytesIO(decoded_bytes)
        file_obj.name = filename
        myfile = genai.upload_file(file_obj, display_name=filename, mime_type=mime_type)
        if hasattr(myfile, "name"):
            logging.info(f"File uploaded to Gemini: {filename}, file_id: {myfile.name}")
            return myfile.name
        else:
            logging.error(f"Gemini file upload did not return file name/id. Response: {myfile}")
            return None
    except Exception as e:
        logging.error(f"Exception during file upload to Gemini: {e}")
        return None

def gemini_generate_content(prompt, file_id=None, chat_input=None, file_ids=None):
    try:
        files = []
        if file_ids:
            for fid in file_ids:
                try:
                    gemini_file = genai.get_file(fid)
                    files.append(gemini_file)
                except Exception as e:
                    logging.error(f"Could not fetch Gemini file for id {fid}: {e}")
        elif file_id:
            try:
                gemini_file = genai.get_file(file_id)
                files.append(gemini_file)
            except Exception as e:
                logging.error(f"Could not fetch Gemini file for id {file_id}: {e}")
        content_list = []
        if files:
            content_list.extend(files)
            content_list.append("\n\n")
        content_list.append(prompt)
        logging.info(f"Prompt sent to Gemini: {prompt[:500]}...")
        model = genai.GenerativeModel(GEMINI_MODEL)
        response = model.generate_content(
            contents=content_list,
            generation_config=genai.types.GenerationConfig(
                max_output_tokens=MAX_OUTPUT_TOKENS
            )
        )
        result = response.text if hasattr(response, "text") else str(response)
        logging.info(f"Gemini response (first 500 chars): {str(result)[:500]}...")
        return result
    except Exception as e:
        logging.error("Error during Gemini generate_content: %s", e)
        return f"Error during Gemini completion: {e}"

def parse_markdown_table(md):
    lines = md.split('\n')
    table_lines = []
    in_table = False
    for l in lines:
        if l.strip().startswith('|') and l.strip().endswith('|'):
            table_lines.append(l.strip())
            in_table = True
        elif in_table and l.strip() == '':
            break
    if not table_lines:
        raise ValueError("No markdown table found")
    header = table_lines[0].strip('|').split('|')
    header = [h.strip() for h in header]
    rows = []
    for l in table_lines[2:]:
        r = [c.strip() for c in l.strip('|').split('|')]
        if len(r) == len(header):
            rows.append(r)
    df = pd.DataFrame(rows, columns=header)
    return df

def save_markdown_as_xlsx(md_text, base_filename):
    try:
        df = parse_markdown_table(md_text)
        memf = io.BytesIO()
        df.to_excel(memf, index=False, engine='xlsxwriter')
        memf.seek(0)
        return memf.read()
    except Exception as e:
        logging.error(f"Failed to convert markdown to XLSX for {base_filename}: {e}")
        return None

def save_shredded_as_docx(shredded_text, rfp_filename):
    doc = Document()
    doc.add_heading(f"Shredded Requirements for {rfp_filename}", 0)
    for line in shredded_text.split('\n'):
        doc.add_paragraph(line)
    memf = io.BytesIO()
    doc.save(memf)
    memf.seek(0)
    return memf.read()

def save_proposal_as_docx(proposal_text, base_filename):
    doc = Document()
    doc.add_heading(f"Proposal Response for {base_filename}", 0)
    for line in proposal_text.split('\n'):
        doc.add_paragraph(line)
    memf = io.BytesIO()
    doc.save(memf)
    memf.seek(0)
    return memf.read()

def save_compliance_as_docx(compliance_text, rfp_filename):
    doc = Document()
    doc.add_heading(f"Compliance Check for {rfp_filename}", 0)
    for line in compliance_text.split('\n'):
        doc.add_paragraph(line)
    memf = io.BytesIO()
    doc.save(memf)
    memf.seek(0)
    return memf.read()

def save_virtual_board_as_docx(board_text, base_filename):
    doc = Document()
    doc.add_heading(f"Evaluation Board for {base_filename}", 0)
    for line in board_text.split('\n'):
        doc.add_paragraph(line)
    memf = io.BytesIO()
    doc.save(memf)
    memf.seek(0)
    return memf.read()

def save_loe_as_docx(loe_text, proposal_filename):
    doc = Document()
    doc.add_heading(f"Level of Effort for {proposal_filename}", 0)
    for line in loe_text.split('\n'):
        doc.add_paragraph(line)
    memf = io.BytesIO()
    doc.save(memf)
    memf.seek(0)
    return memf.read()

def process_document(sess_data, action, selected_filename=None, chat_input=None, rfp_decoded_bytes=None, selected_proposal_filename=None):
    doc_content = None
    doc_fileid = None
    proposal_content = None
    proposal_fileid = None

    if selected_filename and selected_filename in sess_data["uploaded_documents"]:
        doc_content = sess_data["uploaded_documents"][selected_filename]
        doc_fileid = sess_data["uploaded_documents_fileid"].get(selected_filename)
    if selected_proposal_filename and selected_proposal_filename in sess_data["proposals"]:
        proposal_content = sess_data["proposals"][selected_proposal_filename]
        proposal_fileid = sess_data["proposals_fileid"].get(selected_proposal_filename)

    if action == 'shred':
        if not doc_content:
            logging.warning("No uploaded document found for shredding.")
            return "No document uploaded.", None, None, None, None
        prompt = read_prompt_file('shred')
        if chat_input:
            prompt += f"\nUser additional instructions: {chat_input}\n"
        prompt += f"\nFile Name: {selected_filename}\n\n"
        prompt += doc_content
        logging.info(f"[SHRED] Sending document {selected_filename} to Gemini for shredding.")
        result = gemini_generate_content(prompt, file_id=doc_fileid, chat_input=chat_input)
        if result and not result.startswith("Error"):
            xlsx_bytes = save_markdown_as_xlsx(result, selected_filename)
            generated_xlsx_name = f"{os.path.splitext(selected_filename)[0]}_shredded.xlsx"
            sess_data["uploaded_documents"][generated_xlsx_name] = result
            sess_data["shredded_documents"][generated_xlsx_name] = xlsx_bytes
            return result, generated_xlsx_name, xlsx_bytes, None, None
        else:
            return result, None, None, None, None

    elif action == 'compliance':
        if not proposal_content:
            return "No proposal document selected for compliance.", None, None, None, None
        if not doc_content:
            return "No RFP/SOW/PWS/RFI document selected for compliance.", None, None, None, None

        logging.info(f"[COMPLIANCE] Comparing proposal [{selected_proposal_filename}] to RFP [{selected_filename}]")
        prompt = read_prompt_file('compliance')
        prompt += f"\n---\nRFP/SOW/PWS/RFI ({selected_filename}):\n{doc_content}\n"
        prompt += "---\nGenerated Proposal Document:\n"
        prompt += f"{proposal_content}\n"
        result = gemini_generate_content(prompt, file_id=None, chat_input=None)
        if result and not result.startswith("Error"):
            xlsx_bytes = save_markdown_as_xlsx(result, selected_filename)
            compliance_xlsx_name = f"{os.path.splitext(selected_filename)[0]}_compliance_check.xlsx"
            sess_data["uploaded_documents"][compliance_xlsx_name] = result
            sess_data["shredded_documents"][compliance_xlsx_name] = xlsx_bytes
            return result, compliance_xlsx_name, xlsx_bytes, None, None
        else:
            return result, None, None, None, None

    elif action == 'virtual_board':
        if not proposal_content:
            return "No proposal document selected for evaluation board.", None, None, None, None
        if not doc_content:
            return "No RFP/SOW/PWS/RFI document selected for evaluation board.", None, None, None, None

        logging.info(f"[VIRTUAL_BOARD] Evaluating proposal [{selected_proposal_filename}] against RFP [{selected_filename}]")
        prompt = read_prompt_file('virtual_board')
        prompt += f"\n---\nRFP/SOW/PWS/RFI ({selected_filename}):\n{doc_content}\n"
        prompt += "---\nProposal Document:\n"
        prompt += f"{proposal_content}\n"
        result = gemini_generate_content(prompt, file_id=None, chat_input=None)
        if result and not result.startswith("Error"):
            xlsx_bytes = save_markdown_as_xlsx(result, selected_filename)
            board_xlsx_name = f"{os.path.splitext(selected_filename)[0]}_evaluation_board.xlsx"
            sess_data["uploaded_documents"][board_xlsx_name] = result
            sess_data["shredded_documents"][board_xlsx_name] = xlsx_bytes
            return result, board_xlsx_name, xlsx_bytes, None, None
        else:
            return result, None, None, None, None

    elif action == 'proposal':
        if not doc_content:
            logging.warning("No RFP/SOW/PWS/RFI document selected for proposal action.")
            return "No RFP/SOW/PWS/RFI document selected.", None, None, None, None
        rfp_filename = selected_filename
        rfp_fileid = doc_fileid
        if not rfp_fileid and rfp_filename in sess_data["uploaded_documents_bytes"]:
            try:
                fileid = upload_to_gemini_file(sess_data["uploaded_documents_bytes"][rfp_filename], rfp_filename)
                if fileid:
                    sess_data["uploaded_documents_fileid"][rfp_filename] = fileid
                    rfp_fileid = fileid
                    logging.info(f"RFP file {rfp_filename} uploaded to Gemini for proposal.")
            except Exception as e:
                logging.error(f"Failed to upload RFP file {rfp_filename} for proposal: {e}")
        prompt = read_prompt_file('proposal')
        if chat_input:
            prompt += f"\nUser additional instructions: {chat_input}\n"
        prompt += f"\n---\nRFP/SOW/PWS/RFI ({rfp_filename}):\n{doc_content}\n"
        logging.info(f"[PROPOSAL] Sending document {rfp_filename} to Gemini for proposal generation.")
        result = gemini_generate_content(prompt, file_id=rfp_fileid, chat_input=chat_input)
        sess_data["generated_response"] = result
        if result and not result.startswith("Error"):
            docx_bytes = save_proposal_as_docx(result, rfp_filename)
            generated_docx_name = f"{os.path.splitext(rfp_filename)[0]}_proposal.docx"
            sess_data["proposals"][generated_docx_name] = result
            sess_data["proposals_fileid"][generated_docx_name] = None
            return result, None, None, generated_docx_name, docx_bytes
        else:
            return result, None, None, None, None

    elif action == 'recover':
        if not proposal_content:
            logging.error("No proposal document selected for recovery.")
            return "No proposal document selected for recovery.", None, None, None, None
        if not doc_content:
            logging.error("No compliance check or shredded requirements document selected for recovery.")
            return "No compliance check or shredded requirements document selected for recovery.", None, None, None, None

        findings_content = doc_content
        prompt = read_prompt_file('recover')
        if chat_input:
            prompt += f"\nUser additional instructions: {chat_input}\n"
        prompt += f"\n---\nFindings and Recommendations Table ({selected_filename}):\n{findings_content}\n"
        prompt += f"\n---\nOriginal Proposal Document ({selected_proposal_filename}):\n{proposal_content}\n"
        logging.info(f"[RECOVER] Recovering proposal {selected_proposal_filename} using findings from {selected_filename}.")
        result = gemini_generate_content(prompt, file_id=None, chat_input=chat_input)
        if result and not result.startswith("Error"):
            base_name = os.path.splitext(selected_proposal_filename)[0]
            recovered_docx_name = f"{base_name}_recovered.docx"
            docx_bytes = save_proposal_as_docx(result, base_name)
            sess_data["proposals"][recovered_docx_name] = result
            sess_data["proposals_fileid"][recovered_docx_name] = None
            logging.info(f"Recovered proposal generated and saved as {recovered_docx_name}.")
            return result, None, None, recovered_docx_name, docx_bytes
        else:
            logging.error(f"Error in Gemini recover: {result}")
            return result, None, None, None, None

    elif action == 'loe':
        if not proposal_content:
            logging.warning("No proposal document selected for LOE estimation.")
            return "No proposal document selected for LOE estimation.", None, None, None, None
        proposal_base_name = os.path.splitext(selected_proposal_filename)[0]
        prompt = read_prompt_file('loe')
        if chat_input:
            prompt += f"\nUser additional instructions: {chat_input}\n"
        prompt += f"\n---\nProposal Document ({selected_proposal_filename}):\n{proposal_content}\n"
        logging.info(f"[LOE] Calculating LOE for proposal {selected_proposal_filename}.")
        result = gemini_generate_content(prompt, file_id=None, chat_input=chat_input)
        if result and not result.startswith("Error"):
            loe_xlsx_name = f"{proposal_base_name}_loe.xlsx"
            sess_data["proposals"][loe_xlsx_name] = result
            sess_data["proposals_fileid"][loe_xlsx_name] = None
            xlsx_bytes = save_markdown_as_xlsx(result, proposal_base_name)
            logging.info(f"LOE generated and saved as {loe_xlsx_name}")
            return result, None, None, loe_xlsx_name, xlsx_bytes
        else:
            return result, None, None, None, None

    return "Action not implemented yet.", None, None, None, None

def get_documents_list(docdict, shreddedict):
    all_docs = {}
    for filename, text in docdict.items():
        all_docs[filename] = text
    for filename, doc_bytes in shreddedict.items():
        if filename not in all_docs:
            all_docs[filename] = None
    if not all_docs:
        return html.Div("No documents uploaded or generated.", style={"wordWrap": "break-word"})
    doc_list = []
    for filename in all_docs:
        truncated = truncate_filename(filename)
        ext = filename.lower().split('.')[-1]
        if ext == "xlsx" and filename in shreddedict:
            b64 = base64.b64encode(shreddedict[filename]).decode('utf-8')
            mime = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        elif ext == "docx" and filename in shreddedict:
            b64 = base64.b64encode(shreddedict[filename]).decode('utf-8')
            mime = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        else:
            content = docdict.get(filename, "")
            b64 = base64.b64encode((content.encode('utf-8') if isinstance(content, str) else b"")).decode('utf-8')
            mime = "text/plain"
        download_link = html.A(
            truncated,
            href=f"data:{mime};base64,{b64}",
            download=filename,
            target="_blank",
            style={"wordWrap": "break-word", "marginRight": "10px", "textDecoration": "underline", "maxWidth": "calc(100% - 70px)", "display": "inline-block", "verticalAlign": "middle"}
        )
        doc_list.append(
            dbc.ListGroupItem([
                download_link,
                dbc.Button("Delete", id={'type': 'delete-doc-btn', 'index': filename, 'group': 'doc'}, size="sm", color="danger", className="float-end ms-2")
            ], className="d-flex justify-content-between align-items-center")
        )
    return dbc.ListGroup(doc_list, flush=True)

def get_proposals_list(proposaldict):
    if not proposaldict:
        return html.Div("No proposals uploaded or generated.", style={"wordWrap": "break-word"})
    doc_list = []
    for filename in proposaldict:
        truncated = truncate_filename(filename)
        ext = filename.lower().split('.')[-1]
        file_content = proposaldict[filename]
        try:
            if ext == "xlsx":
                from io import BytesIO
                b64 = base64.b64encode(save_markdown_as_xlsx(file_content, filename)).decode('utf-8')
                mime = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
            elif filename.lower().endswith('_loe.docx'):
                docx_bytes = save_loe_as_docx(file_content, filename)
                b64 = base64.b64encode(docx_bytes).decode('utf-8')
                mime = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
            else:
                docx_bytes = save_proposal_as_docx(file_content, filename)
                b64 = base64.b64encode(docx_bytes).decode('utf-8')
                mime = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        except Exception:
            b64 = base64.b64encode(file_content.encode('utf-8')).decode('utf-8')
            mime = "text/plain"
        download_link = html.A(
            truncated,
            href=f"data:{mime};base64,{b64}",
            download=filename,
            target="_blank",
            style={"wordWrap": "break-word", "marginRight": "10px", "textDecoration": "underline", "maxWidth": "calc(100% - 70px)", "display": "inline-block", "verticalAlign": "middle"}
        )
        doc_list.append(
            dbc.ListGroupItem([
                download_link,
                dbc.Button("Delete", id={'type': 'delete-proposal-btn', 'index': filename, 'group': 'proposal'}, size="sm", color="danger", className="float-end ms-2")
            ], className="d-flex justify-content-between align-items-center")
        )
    return dbc.ListGroup(doc_list, flush=True)

app.layout = dbc.Container([
    dcc.Store(id='preview-window-state', data='expanded'),
    dcc.Store(id='session-id-store', storage_type='session'),
    html.Div(id='set-session-cookie', style={'display': 'none'}),
    dcc.Location(id='dummy-url', refresh=False),
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardHeader(html.H5("Documents")),
                dbc.CardBody([
                    dcc.Upload(
                        id='upload-document',
                        children=html.Div([
                            'Drag and Drop or ',
                            html.A('Select Files')
                        ]),
                        style={
                            'width': '100%',
                            'height': '60px',
                            'lineHeight': '60px',
                            'borderWidth': '1px',
                            'borderStyle': 'dashed',
                            'borderRadius': '5px',
                            'textAlign': 'center',
                            'margin': '10px'
                        },
                        multiple=False
                    ),
                    html.Div(id='documents-list'),
                    dcc.Dropdown(
                        id='select-document-dropdown',
                        options=[],
                        placeholder="Select a document to work with",
                        value=None,
                        style={"marginBottom": "10px"}
                    ),
                ])
            ], className="mb-3"),
            dbc.Card([
                dbc.CardHeader(html.H5("Proposals")),
                dbc.CardBody([
                    dcc.Upload(
                        id='upload-proposal',
                        children=html.Div([
                            'Drag and Drop or ',
                            html.A('Select Files')
                        ]),
                        style={
                            'width': '100%',
                            'height': '60px',
                            'lineHeight': '60px',
                            'borderWidth': '1px',
                            'borderStyle': 'dashed',
                            'borderRadius': '5px',
                            'textAlign': 'center',
                            'margin': '10px'
                        },
                        multiple=False
                    ),
                    html.Div(id='proposals-list'),
                    dcc.Dropdown(
                        id='select-proposal-dropdown',
                        options=[],
                        placeholder="Select a proposal document",
                        value=None,
                        style={"marginBottom": "10px"}
                    ),
                ])
            ], className="mb-3"),
            dbc.Card([
                dbc.CardHeader(html.H5("Instructions")),
                dbc.CardBody([
                    html.Ol([
                        html.Li("Start by uploading a RFP, SOW, PWS or RFI"),
                        html.Li("Next click Shred to get an spreadsheet of the key requirements"),
                        html.Li("Generate a proposal.  You can have it generate the proposal from either the uploaded RFP, SOW, PWS or RFI or Shred.  Select that in the dropdown list under Documents.  If you don't already have a proposal document, just upload it under proposals for the next steps"),
                        html.Li("Check compliance against a document you select in the Proposal section.  Choose the document you want to evaluate by selecting it on the dropdown list.  You must also choose what you want to evaluate it against, from the Documents list like the PWS or the Shred"),
                        html.Li("Recover a document by selecting a document from the proposal and like PWS, shred or compliance check (recommended) under Documents. This will recover any sections that to fix based on the document selected like pws, shred or compliance"),
                        html.Li("Simulate a virtual board evaluation, which you must select the PWS under documents and the document under proposals you want to evaluate.  That will give you representative board ratings.  If you want to recover based on these board ratings, select the document under proposal and the board evaluation document under Documents and then click recover it will recover the problem areas based on the findings."),
                        html.Li("Generate an LOE by selecting the document like PWS, and click LOE, you will get an estimated LOE."),
                    ], style={"wordWrap": "break-word"})
                ])
            ], className="mb-3"),
        ], style={'minWidth': '260px', 'width':'30vw','maxWidth':'30vw'}, width=3),

        dbc.Col([
            dbc.Card([
                dbc.CardHeader(html.H2("RFP Proposal Assistant", style={'wordWrap': 'break-word'})),
                dbc.CardBody([
                    dbc.Form([
                        dbc.Textarea(id="chat-input", placeholder="Enter additional instructions...", style={"width":"100%", "wordWrap": "break-word"}, className="mb-2"),
                    ]),
                    html.Div([
                        dbc.Button("Shred", id="shred-action-btn", className="me-3 mb-2 btn-primary"),
                        dbc.Button("Proposal", id="proposal-action-btn", className="me-3 mb-2 btn-secondary"),
                        dbc.Button("Compliance", id="compliance-action-btn", className="me-3 mb-2 btn-tertiary"),
                        dbc.Button("Recover", id="recover-action-btn", className="me-3 mb-2 btn-tertiary"),
                        dbc.Button("Virtual Board", id="board-action-btn", className="me-3 mb-2 btn-tertiary"),
                        dbc.Button("LOE", id="loe-action-btn", className="mb-2 btn-tertiary"),
                        dbc.Button("Cancel", id="cancel-action-btn", className="ms-3 mb-2 btn-danger", color="danger"),
                    ], className="mt-3 mb-3 d-flex flex-wrap"),
                    dcc.Loading(
                        id="loading",
                        type="default",
                        children=html.Div(
                            id="output-preview-container",
                            children=html.Div(id="output-data-upload"),
                            style={
                                "height": "70vh",
                                "overflowY": "auto",
                                "overflowX": "auto",
                            }
                        ),
                        style={"textAlign": "center"}
                    )
                ])
            ], style={'backgroundColor': 'white'})
        ], style={'width':'70vw','maxWidth':'70vw'}, width=9)
    ], style={'marginTop':'20px'})
], fluid=True)

app.clientside_callback(
    """
    function(n, dummy_url) {
      let sid = window.sessionStorage.getItem('dash_session');
      let store_val = null;
      if(!sid) {
        sid = ([1e7]+-1e3+-4e3+-8e3+-1e11).replace(/[018]/g, c =>
          (c ^ crypto.getRandomValues(new Uint8Array(1))[0] & 15 >> c / 4).toString(16)
        );
        window.sessionStorage.setItem('dash_session', sid);
      }
      document.cookie = "dash_session=" + sid + "; path=/";
      store_val = sid;
      return store_val;
    }
    """,
    Output('session-id-store', 'data'),
    Input('dummy-url', 'pathname')
)

@app.callback(
    Output('output-data-upload', 'children'),
    Output('documents-list', 'children'),
    Output('select-document-dropdown', 'options'),
    Output('select-document-dropdown', 'value'),
    Output('proposals-list', 'children'),
    Output('select-proposal-dropdown', 'options'),
    Output('select-proposal-dropdown', 'value'),
    Output('preview-window-state', 'data'),
    [
        Input('shred-action-btn', 'n_clicks'),
        Input('proposal-action-btn', 'n_clicks'),
        Input('compliance-action-btn', 'n_clicks'),
        Input('recover-action-btn', 'n_clicks'),
        Input('board-action-btn', 'n_clicks'),
        Input('loe-action-btn', 'n_clicks'),
        Input('upload-document', 'contents'),
        State('upload-document', 'filename'),
        Input({'type': 'delete-doc-btn', 'index': ALL, 'group': 'doc'}, 'n_clicks'),
        State('select-document-dropdown', 'value'),
        Input('upload-proposal', 'contents'),
        State('upload-proposal', 'filename'),
        Input({'type': 'delete-proposal-btn', 'index': ALL, 'group': 'proposal'}, 'n_clicks'),
        State('select-proposal-dropdown', 'value'),
        State('chat-input', 'value'),
        Input('cancel-action-btn', 'n_clicks'),
        State('preview-window-state', 'data'),
        State('session-id-store', 'data'),
        Input('session-id-store', 'data')
    ],
    prevent_initial_call=False
)
def unified_master_callback(
    shred_clicks, proposal_clicks, compliance_clicks, recover_clicks, board_clicks, loe_clicks,
    rfp_content, rfp_filename, doc_delete_clicks, selected_doc,
    proposal_content, proposal_filename, proposal_delete_clicks, selected_proposal,
    chat_input, cancel_clicks, preview_window_state,
    session_id_state, session_id_store_input
):
    sid = get_session_id(session_id_state if session_id_state else session_id_store_input)
    sess_data = get_session_data(sid)
    ctx = callback_context
    output_data_upload = html.Div("No action taken yet.", style={"wordWrap": "break-word"})
    uploaded_rfp_decoded_bytes = None

    triggered_id = getattr(ctx, "triggered_id", None)
    is_doc_delete = False
    is_proposal_delete = False
    doc_del_filename = None
    proposal_del_filename = None

    if isinstance(triggered_id, dict):
        if triggered_id.get('type') == 'delete-doc-btn' and triggered_id.get('group') == 'doc':
            is_doc_delete = True
            doc_del_filename = triggered_id.get('index')
        elif triggered_id.get('type') == 'delete-proposal-btn' and triggered_id.get('group') == 'proposal':
            is_proposal_delete = True
            proposal_del_filename = triggered_id.get('index')

    if is_doc_delete and doc_del_filename:
        if doc_del_filename in sess_data["uploaded_documents"]:
            del sess_data["uploaded_documents"][doc_del_filename]
        if doc_del_filename in sess_data["uploaded_documents_fileid"]:
            try:
                genai.delete_file(sess_data["uploaded_documents_fileid"][doc_del_filename])
            except Exception as e:
                logging.warning(f"[{sid}] Failed to delete Gemini file {doc_del_filename}: {e}")
            del sess_data["uploaded_documents_fileid"][doc_del_filename]
        if doc_del_filename in sess_data["uploaded_documents_bytes"]:
            del sess_data["uploaded_documents_bytes"][doc_del_filename]
        if doc_del_filename in sess_data["shredded_documents"]:
            del sess_data["shredded_documents"][doc_del_filename]
        tempdir = sess_data.get("session_tempdir")
        if tempdir and os.path.isdir(tempdir):
            try:
                file_path = os.path.join(tempdir, doc_del_filename)
                if os.path.exists(file_path):
                    os.remove(file_path)
            except Exception:
                pass
        if selected_doc == doc_del_filename:
            selected_doc = None
        logging.info(f"[{sid}] Document deleted: {doc_del_filename}")

    if is_proposal_delete and proposal_del_filename:
        if proposal_del_filename in sess_data["proposals"]:
            del sess_data["proposals"][proposal_del_filename]
        if proposal_del_filename in sess_data["proposals_fileid"]:
            try:
                genai.delete_file(sess_data["proposals_fileid"][proposal_del_filename])
            except Exception as e:
                logging.warning(f"[{sid}] Failed to delete Gemini proposal file {proposal_del_filename}: {e}")
            del sess_data["proposals_fileid"][proposal_del_filename]
        tempdir = sess_data.get("session_tempdir")
        if tempdir and os.path.isdir(tempdir):
            try:
                file_path = os.path.join(tempdir, proposal_del_filename)
                if os.path.exists(file_path):
                    os.remove(file_path)
            except Exception:
                pass
        if selected_proposal == proposal_del_filename:
            selected_proposal = None
        logging.info(f"[{sid}] Proposal deleted: {proposal_del_filename}")

    def safe_get_n_clicks(ctx, idx):
        try:
            return ctx.inputs_list[idx]
        except Exception:
            return []

    doc_delete_clicks = safe_get_n_clicks(ctx, 8)
    proposal_delete_clicks = safe_get_n_clicks(ctx, 12)

    if not ctx.triggered or (getattr(ctx, "triggered_id", None) == 'session-id-store'):
        doc_options = [{'label': truncate_filename(fn), 'value': fn} for fn in sess_data["uploaded_documents"].keys()]
        doc_value = next(iter(sess_data["uploaded_documents"]), None) if sess_data["uploaded_documents"] else None
        proposal_options = [{'label': truncate_filename(fn), 'value': fn} for fn in sess_data["proposals"].keys()]
        proposal_value = next(iter(sess_data["proposals"]), None) if sess_data["proposals"] else None
        documents_list = get_documents_list(sess_data["uploaded_documents"], sess_data["shredded_documents"])
        proposals_list = get_proposals_list(sess_data["proposals"])
        return (
            output_data_upload,
            documents_list, doc_options, doc_value,
            proposals_list, proposal_options, proposal_value,
            "expanded"
        )

    if getattr(ctx, "triggered_id", None) == 'cancel-action-btn':
        lock = sess_data.get("gemini_lock")
        if lock and lock.locked():
            try:
                lock.release()
                logging.info(f"[{sid}] Gemini lock released on cancel.")
            except RuntimeError:
                logging.warning(f"[{sid}] Attempted to release an unlocked Gemini lock.")
        output_data_upload = html.Div("[Cancelled by user]\n", style={"wordWrap": "break-word"})
        doc_options = [{'label': truncate_filename(fn), 'value': fn} for fn in sess_data["uploaded_documents"].keys()]
        doc_value = selected_doc if selected_doc in sess_data["uploaded_documents"] else (next(iter(sess_data["uploaded_documents"]), None) if sess_data["uploaded_documents"] else None)
        proposals_list = get_proposals_list(sess_data["proposals"])
        proposal_options = [{'label': truncate_filename(fn), 'value': fn} for fn in sess_data["proposals"].keys()]
        proposal_value = selected_proposal if selected_proposal in sess_data["proposals"] else (next(iter(sess_data["proposals"]), None) if sess_data["proposals"] else None)
        documents_list = get_documents_list(sess_data["uploaded_documents"], sess_data["shredded_documents"])
        return (
            output_data_upload,
            documents_list, doc_options, doc_value,
            proposals_list, proposal_options, proposal_value,
            "expanded"
        )

    if getattr(ctx, "triggered_id", None) == 'upload-document' and rfp_content is not None and rfp_filename:
        content_type, content_string = rfp_content.split(',')
        decoded = base64.b64decode(content_string)
        uploaded_rfp_decoded_bytes = decoded
        text = decode_document(decoded)
        fileid = None
        if rfp_filename.lower().endswith(('.pdf', '.docx', '.xlsx', '.xls')):
            fileid = upload_to_gemini_file(decoded, rfp_filename)
        if text is not None:
            sess_data["uploaded_documents"][rfp_filename] = text
            sess_data["uploaded_documents_bytes"][rfp_filename] = decoded
            if fileid:
                sess_data["uploaded_documents_fileid"][rfp_filename] = fileid
            logging.info(f"[{sid}] Document uploaded: {rfp_filename}")
        else:
            logging.error(f"[{sid}] Failed to decode uploaded document: {rfp_filename}")

    if getattr(ctx, "triggered_id", None) == 'upload-proposal' and proposal_content is not None and proposal_filename:
        content_type, content_string = proposal_content.split(',')
        decoded = base64.b64decode(content_string)
        text = decode_document(decoded)
        fileid = None
        if proposal_filename.lower().endswith(('.pdf', '.docx', '.xlsx', '.xls')):
            fileid = upload_to_gemini_file(decoded, proposal_filename)
        if text is not None:
            sess_data["proposals"][proposal_filename] = text
            if fileid:
                sess_data["proposals_fileid"][proposal_filename] = fileid
            logging.info(f"[{sid}] Proposal uploaded: {proposal_filename}")
        else:
            logging.error(f"[{sid}] Failed to decode uploaded proposal: {proposal_filename}")

    doc_options = [{'label': truncate_filename(fn), 'value': fn} for fn in sess_data["uploaded_documents"].keys()]
    doc_value = selected_doc if selected_doc in sess_data["uploaded_documents"] else (next(iter(sess_data["uploaded_documents"]), None) if sess_data["uploaded_documents"] else None)
    documents_list = get_documents_list(sess_data["uploaded_documents"], sess_data["shredded_documents"])
    proposals_list = get_proposals_list(sess_data["proposals"])
    proposal_options = [{'label': truncate_filename(fn), 'value': fn} for fn in sess_data["proposals"].keys()]
    proposal_value = selected_proposal if selected_proposal in sess_data["proposals"] else (next(iter(sess_data["proposals"]), None) if sess_data["proposals"] else None)

    action_btns = [
        'shred-action-btn', 'proposal-action-btn', 'compliance-action-btn',
        'recover-action-btn', 'board-action-btn', 'loe-action-btn'
    ]

    if getattr(ctx, "triggered_id", None) in action_btns:
        got_lock = sess_data["gemini_lock"].acquire(blocking=False)
        if not got_lock:
            output_data_upload = html.Div("Another Gemini operation is in progress. Please wait or cancel.", style={"wordWrap": "break-word"})
            return (
                output_data_upload,
                documents_list, doc_options, doc_value,
                proposals_list, proposal_options, proposal_value,
                "expanded"
            )
        try:
            triggered_id = getattr(ctx, "triggered_id", None)
            if triggered_id == "shred-action-btn":
                action_name = "shred"
                result, generated_filename, generated_xlsx_bytes, _, _ = process_document(sess_data, action_name, doc_value, chat_input, uploaded_rfp_decoded_bytes, None)
                # Try to parse as markdown table, show DataTable if possible, else fallback to Markdown
                try:
                    df = parse_markdown_table(result)
                    output_data_upload = dt.DataTable(
                        columns=[{"name": i, "id": i} for i in df.columns],
                        data=df.to_dict('records'),
                        style_table={
                            "overflowX": "auto",
                            "overflowY": "auto",
                            "maxHeight": "60vh",
                            "minWidth": "100%"
                        },
                        style_cell={
                            "whiteSpace": "normal",
                            "height": "auto",
                            "textAlign": "left",
                            "wordBreak": "break-word"
                        },
                        style_header={
                            "fontWeight": "bold"
                        },
                        page_action="none"
                    )
                except Exception as e:
                    logging.warning(f"Failed to parse markdown as table for preview: {e}")
                    output_data_upload = dcc.Markdown(result, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
            elif triggered_id == "compliance-action-btn":
                action_name = "compliance"
                result, generated_filename, generated_xlsx_bytes, _, _ = process_document(
                    sess_data, action_name, doc_value, chat_input, None, proposal_value
                )
                output_data_upload = dcc.Markdown(result, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
            elif triggered_id == "board-action-btn":
                action_name = "virtual_board"
                result, generated_filename, generated_xlsx_bytes, _, _ = process_document(
                    sess_data, action_name, doc_value, chat_input, None, proposal_value
                )
                output_data_upload = dcc.Markdown(result, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
            elif triggered_id == "proposal-action-btn":
                action_name = "proposal"
                result, _, _, generated_filename, generated_docx_bytes = process_document(
                    sess_data, action_name, doc_value, chat_input, None, None
                )
                output_data_upload = dcc.Markdown(result, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
            elif triggered_id == "recover-action-btn":
                action_name = "recover"
                result, _, _, generated_filename, generated_docx_bytes = process_document(
                    sess_data, action_name, doc_value, chat_input, None, proposal_value
                )
                output_data_upload = dcc.Markdown(result, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
            elif triggered_id == "loe-action-btn":
                action_name = "loe"
                result, _, _, generated_filename, generated_xlsx_bytes = process_document(
                    sess_data, action_name, None, chat_input, None, proposal_value
                )
                output_data_upload = dcc.Markdown(result, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
        finally:
            sess_data["gemini_lock"].release()
        doc_options = [{'label': truncate_filename(fn), 'value': fn} for fn in sess_data["uploaded_documents"].keys()]
        doc_value = doc_value if doc_value in sess_data["uploaded_documents"] else (next(iter(sess_data["uploaded_documents"]), None) if sess_data["uploaded_documents"] else None)
        proposal_options = [{'label': truncate_filename(fn), 'value': fn} for fn in sess_data["proposals"].keys()]
        proposal_value = proposal_value if proposal_value in sess_data["proposals"] else (next(iter(sess_data["proposals"]), None) if sess_data["proposals"] else None)
        documents_list = get_documents_list(sess_data["uploaded_documents"], sess_data["shredded_documents"])
        proposals_list = get_proposals_list(sess_data["proposals"])
        return (
            output_data_upload,
            documents_list, doc_options, doc_value,
            proposals_list, proposal_options, proposal_value,
            "shrunk"
        )

    doc_value = doc_value if doc_value in sess_data["uploaded_documents"] else (next(iter(sess_data["uploaded_documents"]), None) if sess_data["uploaded_documents"] else None)
    proposal_value = proposal_value if proposal_value in sess_data["proposals"] else (next(iter(sess_data["proposals"]), None) if sess_data["proposals"] else None)
    return (
        output_data_upload,
        documents_list, doc_options, doc_value,
        proposals_list, proposal_options, proposal_value,
        "expanded"
    )

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