app.py

# app.py
# Universal AI Data Analyst with:
# - Unchanged analysis & assessment logic
# - Fixed Gradio event wiring (uses gr.State for history)
# - Triple-quoted progress strings (no unterminated literals)
# - Sleek full-width UI and Voice-to-Text (browser Web Speech API)
# - Optional HIPAA flags (fallback defaults if not present in settings.py)
from __future__ import annotations

import io
import json
import os
import traceback
from contextlib import redirect_stdout
from datetime import datetime
from typing import Any, Dict, List

import gradio as gr
import pandas as pd
import regex as re2
import re
from langchain_cohere import ChatCohere  # noqa: F401
from settings import (
    GENERAL_CONVERSATION_PROMPT,
    COHERE_MODEL_PRIMARY,
    COHERE_TIMEOUT_S,   # noqa: F401
    USE_OPEN_FALLBACKS  # noqa: F401
)
# Try to import optional HIPAA flags; fall back to safe defaults if not defined.
try:
    from settings import PHI_MODE, PERSIST_HISTORY, HISTORY_TTL_DAYS, REDACT_BEFORE_LLM, ALLOW_EXTERNAL_PHI
except Exception:
    PHI_MODE = False
    PERSIST_HISTORY = True
    HISTORY_TTL_DAYS = 365
    REDACT_BEFORE_LLM = False
    ALLOW_EXTERNAL_PHI = True

from audit_log import log_event
from privacy import safety_filter, refusal_reply
from llm_router import cohere_chat, _co_client, cohere_embed

# ---------------------- Helpers (analysis logic unchanged) ----------------------
def load_markdown_text(filepath: str) -> str:
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        return f"**Error:** Document `{os.path.basename(filepath)}` not found."

def _sanitize_text(s: str) -> str:
    if not isinstance(s, str):
        return s
    # Remove control characters (except newline and tab)
    return re2.sub(r"[\p{C}--[\n\t]]+", "", s)

# Conservative PHI redaction patterns (only applied if PHI_MODE & REDACT_BEFORE_LLM are enabled)
PHI_PATTERNS = [
    (re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), "[REDACTED_SSN]"),
    (re.compile(r"\b\d{9}\b"), "[REDACTED_MRN]"),
    (re.compile(r"\b\d{3}[-.\s]?\d{3}[-.\s]?\d{4}\b"), "[REDACTED_PHONE]"),
    (re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}"), "[REDACTED_EMAIL]"),
    (re.compile(r"\b(19|20)\d{2}-\d{2}-\d{2}\b"), "[REDACTED_DOB]"),
    (re.compile(r"\b\d{2}/\d{2}/(19|20)\d{2}\b"), "[REDACTED_DOB]"),
    (re.compile(r"\b\d{5}(-\d{4})?\b"), "[REDACTED_ZIP]"),
]

# ------------------------------------------------------------------
# Helper to safely convert pandas scalars → native Python types
# ------------------------------------------------------------------
def to_python(val):
    """Convert pandas/numpy scalars to native Python types for JSON serialization"""
    import numpy as np
    if isinstance(val, (np.integer, np.int64)):
        return int(val)
    if isinstance(val, (np.floating, np.float64)):
        return float(val)
    if hasattr(val, 'item'):
        return val.item()
    return val
    
def redact_phi(text: str) -> str:
    if not isinstance(text, str):
        return text
    t = text
    for pat, repl in PHI_PATTERNS:
        t = pat.sub(repl, t)
    return t

def safe_log(event_name: str, meta: dict | None = None):
    # Avoid logging raw PHI or payloads
    try:
        meta = (meta or {}).copy()
        meta.pop("raw", None)
        log_event(event_name, None, meta)
    except Exception:
        # Never raise from logging
        pass

def _create_python_script(user_scenario: str, schema_context: str) -> str:
    EXPERT_ANALYTICAL_GUIDELINES = """
--- EXPERT ANALYTICAL GUIDELINES ---
When writing your script, you MUST follow these expert business rules:
1.  **Linking Datasets Rule:** If you need to connect facilities to health zones when the 'zone' column is not in the facility list,
    you must first identify the high-priority zone from the beds data, then find the major city (by facility count) in the facility list,
    and *then* assess that city's capacity. Do not try to filter the facility list by a 'zone' column if it does not exist in the schema.
2.  **Prioritization Rule:** To prioritize locations, you MUST combine the most recent population data with specific high-risk health indicators
    to create a multi-factor risk score.
3.  **Capacity Calculation Rule:** For capacity over a 3-month window, assume **60 working days**.
4.  **Cost Calculation Rule:** Sum 'Startup cost' and 'Ongoing cost' per person before multiplying.
"""
    prompt_for_coder = f"""\
You are an expert Python data scientist. Your job is to write a script to extract the data needed to answer the user's request.
You have dataframes in a list `dfs`.

{EXPERT_ANALYTICAL_GUIDELINES}

--- DATA SCHEMA ---
{schema_context}
--- END DATA SCHEMA ---

CRITICAL RULES:
1.  **DO NOT READ FILES:** You MUST NOT include `pd.read_csv`. The data is ALREADY loaded in the `dfs` variable. You MUST use this variable. Failure to do so will cause a fatal error.
2.  **JSON OUTPUT ONLY:** Your script's ONLY output must be a single JSON object printed to stdout containing the raw data findings.
3.  **BE PRECISE:** Use the exact, case-sensitive column names from the schema and robustly clean strings (`re.sub()`) before converting to numbers.
4.  **JSON SERIALIZATION:** Before adding data to your final dictionary for JSON conversion, you MUST convert any pandas-specific types (like `int64`) to standard Python types using `.item()` for single values or `.tolist()` for lists.

--- USER'S SCENARIO ---
{user_scenario}

--- PYTHON SCRIPT ---
Now, write the complete Python script that performs the analysis and prints a single, serializable JSON object.
```python
"""
    generated_text = cohere_chat(prompt_for_coder)
    match = re2.search(r"```python\n(.*?)```", generated_text, re2.DOTALL)
    if match:
        return match.group(1).strip()
    return "print(json.dumps({'error': 'Failed to generate a valid Python script.'}))"


def _generate_long_report(prompt: str) -> str:
    try:
        client = _co_client()
        if not client:
            return "Error: Cohere client not initialized."
        response = client.chat(
            model=COHERE_MODEL_PRIMARY,
            message=prompt,
            max_tokens=4096,
        )
        return response.text
    except Exception as e:
        safe_log("cohere_chat_error", {"err": str(e)})
        return f"Error during final report generation: {e}"


def _generate_final_report(user_scenario: str, raw_data_json: str) -> str:
    prompt_for_writer = f"""\
You are an expert management consultant and data analyst.
A data science script has run to extract key findings. You have the user's original request and the raw JSON data.

Your task is to synthesize these raw findings into a single, comprehensive, and professional report that directly answers all of the user's questions with detailed justifications.

--- USER'S ORIGINAL SCENARIO & DELIVERABLES ---
{user_scenario}
--- END SCENARIO ---

--- RAW DATA FINDINGS (JSON) ---
{raw_data_json}
--- END RAW DATA ---

Now, write the final, polished report. The report MUST:
1.  Follow the "Expected Output Format" requested by the user.
2.  Use tables, bullet points, and DETAILED narrative justifications for each recommendation.
3.  Synthesize the raw data into actionable insights. Do not just copy the raw numbers; interpret them.
4.  Ensure you fully address ALL evaluation questions, especially the final recommendations.
"""
    return _generate_long_report(prompt_for_writer)


def _append_msg(h: List[Dict[str, str]], r: str, c: str) -> List[Dict[str, str]]:
    return (h or []) + [{"role": r, "content": c}]


def ping_cohere() -> str:
    try:
        cli = _co_client()
        if not cli:
            return "Cohere client not initialized."
        vecs = cohere_embed(["hello", "world"])
        return f"Cohere OK ✅ (model={COHERE_MODEL_PRIMARY})" if vecs else "Cohere reachable."
    except Exception as e:
        return f"Cohere ping failed: {e}"


def handle(user_msg: str, files: list, yield_update) -> str:
    try:
        # Safety filter on incoming message
        safe_in, blocked_in, reason_in = safety_filter(user_msg, mode="input")
        if blocked_in:
            return refusal_reply(reason_in)

        # Optional PHI redaction for prompts sent to an external LLM
        redacted_in = safe_in
        if PHI_MODE and REDACT_BEFORE_LLM:
            redacted_in = redact_phi(safe_in)

        file_paths: List[str] = [getattr(f, "name", None) or f for f in (files or [])]

        if file_paths:
            # CSV analysis path (unchanged)
            dataframes, schema_parts = [], []
            for i, p in enumerate(file_paths):
                if p.endswith(".csv"):
                    try:
                        df = pd.read_csv(p)
                    except UnicodeDecodeError:
                        df = pd.read_csv(p, encoding="latin1")
                    dataframes.append(df)
                    schema_parts.append(
                        f"DataFrame `dfs[{i}]` (`{os.path.basename(p)}`):\n{df.head().to_markdown()}\n"
                    )

            if not dataframes:
                return "Please upload at least one CSV file."

            schema_context = "\n".join(schema_parts)

            # If external PHI is not allowed, use redacted prompt; otherwise use original
            prompt_for_code = redacted_in if (PHI_MODE and not ALLOW_EXTERNAL_PHI) else safe_in

            yield_update("""```
🧠 Generating aligned analysis script...
```""")
            analysis_script = _create_python_script(prompt_for_code, schema_context)

            yield_update("""```
⚙️ Executing script to extract raw data...
```""")
            execution_namespace = {"dfs": dataframes, "pd": pd, "re": re, "json": json}
            output_buffer = io.StringIO()

            try:
                with redirect_stdout(output_buffer):
                    exec(analysis_script, execution_namespace)
                raw_data_output = output_buffer.getvalue()
                # ←←← ADD THIS SAFETY WRAPPER
                try:
                    raw_data = json.loads(raw_data_output)
                except json.JSONDecodeError:
                    # Sometimes the model prints extra text → try to extract JSON
                    import re
                    json_match = re.search(r'\{.*\}', raw_data_output, re.DOTALL)
                    if json_match:
                        raw_data = json.loads(json_match.group(0))
                    else:
                        raise ValueError("No valid JSON found in script output")
                # Convert any remaining pandas types safely
                def convert_pandas(obj):
                    if isinstance(obj, dict):
                        return {k: convert_pandas(v) for k, v in obj.items()}
                    elif isinstance(obj, list):
                        return [convert_pandas(v) for v in obj]
                    else:
                        return to_python(obj)
                raw_data = convert_pandas(raw_data)
                raw_data_json = json.dumps(raw_data)
            except Exception as e:
                return (
                    f"An error occurred executing the script: {e}\n\nGenerated Script:\n"
                    f"```python\n{analysis_script}\n```"
                )

            yield_update("""```
✍️ Synthesizing final comprehensive report...```""")
            writer_input = redacted_in if (PHI_MODE and not ALLOW_EXTERNAL_PHI) else safe_in
            final_report = _generate_final_report(writer_input, raw_data_output)
            return _sanitize_text(final_report)
        else:
            # Pure chat path
            chat_input = redacted_in if (PHI_MODE and not ALLOW_EXTERNAL_PHI) else safe_in
            prompt = f"{GENERAL_CONVERSATION_PROMPT}\n\nUser: {chat_input}\nAssistant:"
            return _sanitize_text(cohere_chat(prompt) or "How can I help further?")

    except Exception as e:
        tb = traceback.format_exc()
        safe_log("app_error", {"err": str(e)})
        return "A critical error occurred. Please contact your administrator." if PHI_MODE else f"A critical error occurred: {e}"


PRIVACY_POLICY_TEXT = load_markdown_text("privacy_policy.md")
TERMS_OF_SERVICE_TEXT = load_markdown_text("terms_of_service.md")


# ---------------------- Sleek UI assets (CSS/JS only) ----------------------

SLEEK_CSS = """
/* Full-bleed, modern look */
:root, body, #root, .gradio-container { height: 100%; }
.gradio-container { padding: 0 !important; }
.block { padding: 0 !important; }

/* Header */
.header {
  padding: 20px 28px;
  background: linear-gradient(135deg, #0e1726, #1d2a44 60%, #243a5e);
  color: #fff;
  display: flex; align-items: center; justify-content: space-between;
  gap: 16px;
}
.header h1 { margin: 0; font-size: 22px; letter-spacing: 0.3px; font-weight: 600; }
.header .badge { font-size: 12px; opacity: 0.9; background:#ffffff22; padding:6px 10px; border-radius: 999px; }

/* Main layout */
.main {
  display: grid;
  grid-template-columns: 420px 1fr;
  gap: 16px;
  padding: 16px;
  height: calc(100vh - 72px);
  box-sizing: border-box;
}
.left, .right {
  background: #0b1020;
  color: #e9edf3;
  border-radius: 16px;
  border: 1px solid #1c2642;
}
.left { padding: 16px; display: flex; flex-direction: column; gap: 12px; }
.right { padding: 0; display: flex; flex-direction: column; }

/* Panels */
.panel-title { font-size: 14px; font-weight: 600; color: #aeb8cc; margin-bottom: 6px; }
.helper { font-size: 12px; color: #97a3bb; margin-bottom: 8px; }

/* Sticky actions */
.actions {
  display: flex; gap: 8px; align-items: center; justify-content: stretch;
}
.actions .gr-button { flex: 1; }

/* Tabs full height */
.right .tabs { height: 100%; display: flex; flex-direction: column; }
.right .tabitem { flex: 1; display: flex; flex-direction: column; }
#chatbot_container { flex: 1; }
#chatbot_container .gr-chatbot { height: 100%; }

/* Tiny separators */
.hr { height: 1px; background: #16203b; margin: 10px 0; }

/* Voice hint */
.voice-hint { font-size: 12px; color:#9fb0cc; margin-top: 4px; }
/* ——— MAKE ANALYSIS OUTPUT WINDOW MUCH TALLER & SCROLL-FRIENDLY ——— */
#chatbot_container {
    flex: 1;
    min-height: 0;                    /* Critical for proper flex shrinking */
}

#chatbot_container .gr-chatbot {
    height: 100% !important;
    max-height: none !important;      /* Remove Gradio's artificial cap */
}

#chatbot_container .message-wrap {
    max-width: 100% !important;
}

/* Make the actual message container take full height and scroll nicely */
#chatbot_container .chatbot {
    overflow-y: auto !important;
    overflow-x: hidden;
    padding: 20px !important;
    scrollbar-width: thin;
    scrollbar-color: #3a4a6e #16203b;
}

/* Optional: nicer scrollbar for WebKit browsers */
#chatbot_container .chatbot::-webkit-scrollbar {
    width: 8px;
}
#chatbot_container .chatbot::-webkit-scrollbar-track {
    background: #16203b;
}
#chatbot_container .chatbot::-webkit-scrollbar-thumb {
    background: #3a4a6e;
    border-radius: 4px;
}

/* Make markdown content more readable in long reports */
#chatbot_container .message pre {
    overflow-x: auto;
    background: #0f1629 !important;
    border: 1px solid #2a3755;
}

/* Increase visible height dramatically */
.main {
    height: calc(100vh - 72px) !important;   /* Already good */
    padding: 12px 16px;                      /* Slightly less padding = more space */
}
/* ——— EXPANDED ANALYSIS OUTPUT WINDOW ——— */
#chatbot_container { flex: 1; min-height: 0; }
#chatbot_container .gr-chatbot { height: 100% !important; max-height: none !important; }
#chatbot_container .chatbot {
    overflow-y: auto !important;
    padding: 20px !important;
    scrollbar-width: thin;
    scrollbar-color: #3a4a6e #16203b;
}
#chatbot_container .chatbot::-webkit-scrollbar { width: 8px; }
#chatbot_container .chatbot::-webkit-scrollbar-track { background: #16203b; }
#chatbot_container .chatbot::-webkit-scrollbar-thumb { background: #3a4a6e; border-radius: 4px; }
/* ——— CRITICAL FIX: Make Chatbot fill the entire right panel ——— */
#chatbot_container {
    flex: 1 1 100% !important;
    min-height: 0;
    display: flex !important;
}

#chatbot_container > .wrap {
    flex: 1 !important;
    display: flex !important;
    flex-direction: column !important;
}

/* This is the actual scrolling message area */
#chatbot_container .chatbot {
    flex: 1 !important;
    min-height: 0 !important;
    max-height: none !important;
    overflow-y: auto !important;
    overflow-x: hidden !important;
    padding: 24px !important;
}

/* Remove Gradio’s default max-height caps */
#chatbot_container .gr-chatbot,
#chatbot_container .gr-prose,
#chatbot_container .message-wrap {
    max-height: none !important;
    height: 100% !important;
}

/* Optional: nicer scrollbar */
#chatbot_container .chatbot::-webkit-scrollbar {
    width: 8px;
}
#chatbot_container .chatbot::-webkit-scrollbar-track {
    background: transparent;
}
#chatbot_container .chatbot::-webkit-scrollbar-thumb {
    background: rgba(100, 120, 160, 0.4);
    border-radius: 4px;
}
#chatbot_container .chatbot::-webkit-scrollbar-thumb:hover {
    background: rgba(100, 120, 160, 0.7);
}
/* ──────── FINAL WORKING FIX FOR GRADIO 4+ CHATBOT HEIGHT (2025) ──────── */
#chatbot_container {
    flex: 1 !important;
    min-height: 0;
    display: flex !important;
    flex-direction: column !important;
}

/* This is the real container that holds the messages in Gradio 4+ */
#chatbot_container .svelte-1cea1s5 {
    flex: 1 !important;
    min-height: 0 !important;
    display: flex !important;
    flex-direction: column !important;
}

/* The actual scrollable message area (this is the one that was hidden) */
#chatbot_container .messages {
    flex: 1 !important;
    overflow-y: auto !important;
    overflow-x: hidden !important;
    padding: 24px !important;
    min-height: 0 !important;
}

/* Remove any max-height caps */
#chatbot_container .gr-chatbot,
#chatbot_container .svelte-1cea1s5,
#chatbot_container .messages,
#chatbot_container * {
    max-height: none !important;
}

/* Nice scrollbar */
#chatbot_container .messages::-webkit-scrollbar {
    width: 8px;
}
#chatbot_container .messages::-webkit-scrollbar-track {
    background: transparent;
}
#chatbot_container .messages::-webkit-scrollbar-thumb {
    background: rgba(100, 120, 160, 0.4);
    border-radius: 4px;
}
#chatbot_container .messages::-webkit-scrollbar-thumb:hover {
    background: rgba(100, 120, 160, 0.7);
}

/* Optional: make code blocks look better in long reports */
#chatbot_container pre {
    background: #0f1629 !important;
    border: 1px solid #2a3755 !important;
    border-radius: 8px !important;
}
/* ── GRADIO CHATBOT SCROLL FIX (2025) ── */
/* Adaptive height: Scales to 80% of viewport, min 500px for small screens */
#chatbot_root {
    height: calc(80vh - 50px) !important;  /* Fills most of right panel, minus header/margins */
    min-height: 500px !important;
    max-height: 90vh !important;
    overflow-y: auto !important;  /* FORCE SCROLLBAR WHEN NEEDED */
    overflow-x: hidden !important;
    scrollbar-width: thin !important;
    scrollbar-color: #3a4a6e #16203b !important;
}

/* Target inner messages container (Gradio's scrollable area) */
#chatbot_root .messages,
#chatbot_root [role="log"] {  /* Fallback for type="messages" */
    height: 100% !important;
    overflow-y: auto !important;
    padding: 20px !important;
}

/* WebKit scrollbar (Chrome/Edge/Safari) */
#chatbot_root::-webkit-scrollbar,
#chatbot_root .messages::-webkit-scrollbar {
    width: 8px !important;
}
#chatbot_root::-webkit-scrollbar-track {
    background: #16203b !important;
}
#chatbot_root::-webkit-scrollbar-thumb {
    background: #3a4a6e !important;
    border-radius: 4px !important;
}
#chatbot_root::-webkit-scrollbar-thumb:hover {
    background: rgba(100, 120, 160, 0.7) !important;
}

/* Ensure long markdown/tables don't break layout */
#chatbot_root pre, #chatbot_root table {
    overflow-x: auto !important;
    background: #0f1629 !important;
    border: 1px solid #2a3755 !important;
    border-radius: 8px !important;
}
"""

VOICE_STT_HTML = """
<script>
let __rs_rec = null;
function rs_toggle_stt(elemId){
  const SpeechRecognition = window.SpeechRecognition || window.webkitSpeechRecognition;
  if (!SpeechRecognition){
    alert("This browser does not support Speech Recognition. Try Chrome or Edge.");
    return;
  }
  if (__rs_rec){ __rs_rec.stop(); __rs_rec = null; return; }
  __rs_rec = new SpeechRecognition();
  __rs_rec.lang = "en-US";
  __rs_rec.interimResults = true;
  __rs_rec.continuous = true;

  const box = document.querySelector(`#${elemId} textarea`);
  if (!box){ alert("Prompt box not found."); return; }
  let base = box.value || "";

  __rs_rec.onresult = (ev) => {
    let t = "";
    for (let i = ev.resultIndex; i < ev.results.length; i++){
      t += ev.results[i].transcript;
    }
    box.value = (base + " " + t).trim();
    box.dispatchEvent(new Event("input", { bubbles: true }));
  };
  __rs_rec.onend = () => { __rs_rec = null; };
  __rs_rec.start();
}
</script>
"""


# ---------------------- Sleek UI (with fixed State wiring) ----------------------

with gr.Blocks(theme=gr.themes.Soft(), css=SLEEK_CSS, fill_width=True) as demo:
    # Persistent in-memory history component (fixes list/_id error)
    assessment_history = gr.State([])

    # Header
    with gr.Row(elem_classes=["header"]):
        gr.Markdown("<h1>Clarity Ops Augemented Decision Support</h1>")
        pill = "PHI Mode ON · history off" if (PHI_MODE and not PERSIST_HISTORY) else \
               "PHI Mode ON" if PHI_MODE else "PHI Mode OFF"
        gr.Markdown(f"<span class='badge'>{pill}</span>")

    # Main layout
    with gr.Row(elem_classes=["main"]):
        # Left panel
        with gr.Column(elem_classes=["left"]):
            gr.Markdown("<div class='panel-title'>New Assessment</div>")
            gr.Markdown("<div class='helper'>Upload CSVs for analysis, or enter a prompt. Voice works in modern browsers.</div>")
            files_input = gr.Files(
                label="Upload Data Files (.csv)",
                file_count="multiple",
                type="filepath",
                file_types=[".csv"],
            )
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Paste your scenario or question here...",
                lines=12,
                elem_id="prompt_box",
                autofocus=True,
            )

            with gr.Row(elem_classes=["actions"]):
                send_btn = gr.Button("▶️ Run Analysis", variant="primary")
                clear_btn = gr.Button("🧹 Clear")
                voice_btn = gr.Button("🎙️ Voice")

            gr.Markdown("<div class='voice-hint'>Click Voice to start/stop dictation into the prompt box.</div>")
            ping_btn = gr.Button("🔌 Ping Cohere")
            ping_out = gr.Markdown()

            gr.Markdown("<div class='hr'></div>")
            if PHI_MODE:
                gr.Markdown(
                    "⚠️ **PHI Mode:** History persistence is disabled by default. Avoid unnecessary identifiers."
                )

            with gr.Accordion("Privacy & Terms", open=False):
                gr.Markdown(PRIVACY_POLICY_TEXT)
                gr.Markdown("<div class='hr'></div>")
                gr.Markdown(TERMS_OF_SERVICE_TEXT)

        # Right panel
        with gr.Column(elem_classes=["right"]):
            with gr.Tabs(elem_classes=["tabs"]):
                with gr.TabItem("Current Assessment", id=0, elem_classes=["tabitem"]):
                    with gr.Column(elem_id="chatbot_container"):
                       chat_history_output = gr.Chatbot(
                            label="Analysis Output", 
                            type="messages",
                            height="600",  # ← This removes the 400px cap and lets it fill the parent
                            container=False,
                            autoscroll=True,
                            elem_id="chatbot_root",  # For CSS targeting
                            resizable=True,
                        )
                with gr.TabItem("Assessment History", id=1, elem_classes=["tabitem"]):
                    gr.Markdown("### Review Past Assessments")
                    history_dropdown = gr.Dropdown(label="Select an assessment to review", choices=[])
                    history_display = gr.Markdown(label="Selected Assessment Details")

    # Inject voice-to-text helper
    gr.HTML(VOICE_STT_HTML)

    # --------- Event logic (unchanged analysis flow) ----------

    def run_analysis_wrapper(prompt, files, chat_history_list, history_state_list):
        if not prompt:
            gr.Warning("Please enter a prompt.")
            yield chat_history_list, history_state_list, gr.update()
            return

        # Append user's message
        chat_with_user_msg = _append_msg(chat_history_list, "user", prompt)

        # Optional progress callback (not streaming in this UI)
        def dummy_update(message: str):
            pass

        # Thinking bubble
        thinking_message = _append_msg(
            chat_with_user_msg,
            "assistant",
            """```
🧠 Generating and executing analysis... Please wait.
```""",
        )
        yield thinking_message, history_state_list, gr.update()

        # Run analysis/chat
        ai_response_text = handle(prompt, files, dummy_update)

        # Append final assistant response
        final_chat = _append_msg(chat_with_user_msg, "assistant", ai_response_text)
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        # Capture filenames (if any)
        file_names: List[str] = []
        if files:
            file_names = [
                os.path.basename(f.name if hasattr(f, "name") else f) for f in files
            ]

        # Build history record
        new_entry = {
            "id": timestamp,
            "prompt": prompt,
            "files": file_names,
            "response": ai_response_text,
            "chat_history": final_chat,
        }

        # Respect PHI/history flags
        if PERSIST_HISTORY and (not PHI_MODE or (PHI_MODE and HISTORY_TTL_DAYS > 0)):
            updated_history: List[Dict[str, Any]] = (history_state_list or []) + [new_entry]
        else:
            updated_history = history_state_list or []

        history_labels = [f"{item['id']} - {item['prompt'][:40]}..." for item in updated_history]

        yield final_chat, updated_history, gr.update(choices=history_labels)

    def view_history(selection: str, history_state_list: List[Dict[str, Any]]) -> str:
        if not selection or not history_state_list:
            return ""
        try:
            selected_id = selection.split(" - ", 1)
        except Exception:
            selected_id = selection

        selected_assessment = next(
            (item for item in history_state_list if item.get("id") == selected_id), None
        )
        if not selected_assessment:
            return "Could not find the selected assessment."

        file_list = selected_assessment.get("files", [])
        file_list_md = "\n- ".join(file_list) if file_list else "*(no files uploaded)*"

        chat_entries = selected_assessment.get("chat_history", [])
        chat_md_lines = []
        for msg in chat_entries:
            role = msg.get("role", "").capitalize()
            content = msg.get("content", "")
            chat_md_lines.append(f"**{role}:** {content}")
        chat_md = "\n\n".join(chat_md_lines)

        return f"""### Assessment from: {selected_assessment['id']}
**Files Used:**
- {file_list_md}
---
**Original Prompt:**
> {selected_assessment['prompt']}
---
**AI Generated Response:**
{selected_assessment['response']}
---
**Chat Transcript:**
{chat_md}
"""

    # Wire events (using proper gr.State component for history)
    send_btn.click(
        run_analysis_wrapper,
        inputs=[prompt_input, files_input, chat_history_output, assessment_history],
        outputs=[chat_history_output, assessment_history, history_dropdown],
    )
    history_dropdown.change(
        view_history,
        inputs=[history_dropdown, assessment_history],
        outputs=[history_display],
    )
    clear_btn.click(
        lambda: (None, None, []),
        outputs=[prompt_input, files_input, chat_history_output],
    )
    ping_btn.click(ping_cohere, outputs=[ping_out])
    voice_btn.click(None, [], [], js="rs_toggle_stt('prompt_box')")


if __name__ == "__main__":
    if not os.getenv("COHERE_API_KEY"):
        print("🔴 COHERE_API_KEY environment variable not set. Application may not function correctly.")
    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))