app.py

import os
import json
import shutil
import gradio as gr
import tempfile
from datetime import datetime
from typing import List, Dict, Any, Optional
from pytube import YouTube
from pathlib import Path # <-- Add this import at the top of your file with the other imports
import re

# --- Agent Imports & Safe Fallbacks ---
try:
    from alz_companion.agent import (
        bootstrap_vectorstore, make_rag_chain, answer_query, synthesize_tts,
        transcribe_audio, detect_tags_from_query, describe_image, build_or_load_vectorstore,
        _default_embeddings
    )
    from alz_companion.prompts import BEHAVIOUR_TAGS, EMOTION_STYLES
    from langchain.schema import Document
    from langchain_community.vectorstores import FAISS
    AGENT_OK = True
except Exception as e:
    AGENT_OK = False
    # Define all fallback functions and classes
    def bootstrap_vectorstore(sample_paths=None, index_path="data/"): return object()
    def build_or_load_vectorstore(docs, index_path, is_personal=False): return object()
    def make_rag_chain(vs_general, vs_personal, **kwargs): return lambda q, **k: {"answer": f"(Demo) You asked: {q}", "sources": []}
    def answer_query(chain, q, **kwargs): return chain(q, **kwargs)
    def synthesize_tts(text: str, lang: str = "en"): return None
    def transcribe_audio(filepath: str, lang: str = "en"): return "This is a transcribed message."
    def detect_tags_from_query(query: str, behavior_options: list, emotion_options: list): return {"detected_behavior": "None", "detected_emotion": "None"}
    def describe_image(image_path: str): return "This is a description of an image."
    def _default_embeddings(): return None
    class Document:
        def __init__(self, page_content, metadata):
            self.page_content = page_content
            self.metadata = metadata
    class FAISS:
        def __init__(self):
            self.docstore = type('obj', (object,), {'_dict': {}})()
    BEHAVIOUR_TAGS = {"None": []}
    EMOTION_STYLES = {"None": {}}
    print(f"WARNING: Could not import from alz_companion ({e}). Running in UI-only demo mode.")

# --- Centralized Configuration ---
CONFIG = {
    "themes": ["All", "The Father", "Still Alice", "Away from Her", "General Caregiving"],
    "roles": ["patient", "caregiver"],
    "behavior_tags": ["None"] + list(BEHAVIOUR_TAGS.keys()),
    "emotion_tags": ["None"] + list(EMOTION_STYLES.keys()),
    "languages": {"English": "en", "Chinese": "zh", "Malay": "ms", "French": "fr", "Spanish": "es"},
    "tones": ["warm", "neutral", "formal", "playful"]
}

# --- File Management & Vector Store Logic ---
INDEX_BASE = os.getenv('INDEX_BASE', 'data')
UPLOADS_BASE = os.path.join(INDEX_BASE, "uploads")

# OPTION A: --- CHANGE THIS LINE ---
PERSONAL_INDEX_PATH = os.path.join(str(Path.home()), "AlzCompanionData", "personal_faiss_index")
# for another space: PERSONAL_INDEX_PATH = os.path.join(str(Path.home()), "AlzPersonalData", "personal_faiss_index")
# changing it from the absolute path (designed for a single local computer) to relative path (perfect for Hugging Face Spaces):
# Does NOT work -> PERSONAL_INDEX_PATH = os.path.join(INDEX_BASE, "personal_faiss_index")
# OPTION A: --- END CHANGE ---
# old code PERSONAL_INDEX_PATH = os.path.join(INDEX_BASE, "personal_faiss_index")
os.makedirs(UPLOADS_BASE, exist_ok=True)
# OPTION A: Also create the parent directory for the personal index
os.makedirs(os.path.dirname(PERSONAL_INDEX_PATH), exist_ok=True)

# OPTION B: --- Example for macOS or Linux ---
# OPTION B: PERSONAL_INDEX_PATH = "/Users/YourUsername/AlzCompanionData/personal_faiss_index"
# OPTION B: Make sure to create the directory
# OPTION B: os.makedirs(os.path.dirname(PERSONAL_INDEX_PATH), exist_ok=True)

THEME_PATHS = {t: os.path.join(INDEX_BASE, f"faiss_index_{t.replace(' ', '').lower()}") for t in CONFIG["themes"]}
vectorstores = {}
personal_vectorstore = None


def canonical_theme(tk: str) -> str: return tk if tk in CONFIG["themes"] else "All"
def theme_upload_dir(theme: str) -> str:
    p = os.path.join(UPLOADS_BASE, f"theme_{canonical_theme(theme).replace(' ', '').lower()}")
    os.makedirs(p, exist_ok=True)
    return p
def load_manifest(theme: str) -> Dict[str, Any]:
    p = os.path.join(theme_upload_dir(theme), "manifest.json")
    if os.path.exists(p):
        try:
            with open(p, "r", encoding="utf-8") as f: return json.load(f)
        except Exception: pass
    return {"files": {}}
def save_manifest(theme: str, man: Dict[str, Any]):
    with open(os.path.join(theme_upload_dir(theme), "manifest.json"), "w", encoding="utf-8") as f: json.dump(man, f, indent=2)
def list_theme_files(theme: str) -> List[tuple[str, bool]]:
    man = load_manifest(theme)
    base = theme_upload_dir(theme)
    found = [(n, bool(e)) for n, e in man.get("files", {}).items() if os.path.exists(os.path.join(base, n))]
    existing = {n for n, e in found}
    for name in sorted(os.listdir(base)):
        if name not in existing and os.path.isfile(os.path.join(base, name)): found.append((name, False))
    man["files"] = dict(found)
    save_manifest(theme, man)
    return found
def copy_into_theme(theme: str, src_path: str) -> str:
    fname = os.path.basename(src_path)
    dest = os.path.join(theme_upload_dir(theme), fname)
    shutil.copy2(src_path, dest)
    return dest
def seed_files_into_theme(theme: str):
    SEED_FILES = [
        ("sample_data/caregiving_tips.txt", True),
        ("sample_data/the_father_segments_tagged_with_emotion_hybrid.jsonl", True),
        ("sample_data/still_alice_segments_tagged_with_emotion_hybrid.jsonl", True),
        ("sample_data/away_from_her_segments_tagged_with_emotion_hybrid.jsonl", True)
    ]
    man, changed = load_manifest(theme), False
    for path, enable in SEED_FILES:
        if not os.path.exists(path): continue
        fname = os.path.basename(path)
        if not os.path.exists(os.path.join(theme_upload_dir(theme), fname)):
            copy_into_theme(theme, path)
            man["files"][fname] = bool(enable)
            changed = True
    if changed: save_manifest(theme, man)

def ensure_index(theme='All'):
    theme = canonical_theme(theme)
    if theme in vectorstores: return vectorstores[theme]
    upload_dir = theme_upload_dir(theme)
    enabled_files = [os.path.join(upload_dir, n) for n, enabled in list_theme_files(theme) if enabled]
    index_path = THEME_PATHS.get(theme)
    vectorstores[theme] = bootstrap_vectorstore(sample_paths=enabled_files, index_path=index_path)
    return vectorstores[theme]

# --- Gradio Callbacks ---
def collect_settings(*args):
    keys = ["role", "patient_name", "caregiver_name", "tone", "language", "tts_lang", "temperature", "behaviour_tag", "emotion_tag", "active_theme", "tts_on", "debug_mode"]
    return dict(zip(keys, args))

def parse_and_tag_entries(text_content: str, source: str) -> List[Document]:
    # Correct separator for the user's file format
    # entries = text_content.split('\n—\n')

    # --- CHANGE START ---
    # Define a regular expression pattern that matches any of the desired separators on their own line.
    # Corrected pattern with '--' added, ordered from longest to shortest
    separator_pattern = r'\n(?:---|--|-|-\.\.-|-\.-)\n'
    
    # Use re.split() to split the text based on the pattern
    entries = re.split(separator_pattern, text_content)
    # --- CHANGE END ---

    docs_to_add = []
    
    for entry in entries:
        if not entry.strip():
            continue
            
        title = "Untitled Text Entry"
        content = entry.strip()

        # Improved parsing logic
        lines = entry.strip().split('\n')
        if lines and "title:" in lines[0].lower():
            title_line = lines[0].split(':', 1)
            title = title_line[1].strip() if len(title_line) > 1 else "Untitled"
            
            content_part = "\n".join(lines[1:])
            if "content:" in content_part.lower():
                content = content_part.split(':', 1)[1].strip()
            else:
                content = content_part.strip()

        full_content = f"Title: {title}\n\nContent: {content}"

        print(f"  - Parsing entry: '{title}'")
        behavior_options = CONFIG.get("behavior_tags", [])
        emotion_options = CONFIG.get("emotion_tags", [])
        detected_tags = detect_tags_from_query(content, behavior_options=behavior_options, emotion_options=emotion_options)
        
        metadata = {"source": source, "title": title}
        if detected_tags.get("detected_behavior") != "None": metadata["behaviors"] = [detected_tags.get("detected_behavior").lower()]
        if detected_tags.get("detected_emotion") != "None": metadata["emotion"] = detected_tags.get("detected_emotion").lower()
        
        docs_to_add.append(Document(page_content=full_content, metadata=metadata))
        
    return docs_to_add

def handle_add_knowledge(title, text_input, file_input, image_input, yt_url):
    global personal_vectorstore
    docs_to_add = []

    # Corrected prioritization of inputs
    if text_input and text_input.strip():
        # Handle manual text input first
        docs_to_add = parse_and_tag_entries(f"Title: {title}\n\nContent: {text_input}", "Text Input")
    elif file_input:
        content_source = os.path.basename(file_input)
        if file_input.lower().endswith('.txt'):
            with open(file_input, 'r', encoding='utf-8') as f:
                file_content = f.read()
            docs_to_add = parse_and_tag_entries(file_content, content_source)
        else: # Handle audio/video
            final_title = title.strip() if title and title.strip() else "Audio/Video Note"
            content_text = transcribe_audio(file_input)
            full_content = f"Title: {final_title}\n\nContent: {content_text}"
            docs_to_add = parse_and_tag_entries(full_content, content_source)
    elif image_input:
        final_title = title.strip() if title and title.strip() else "Image Note"
        content_text = describe_image(image_input)
        full_content = f"Title: {final_title}\n\nContent: {content_text}"
        docs_to_add = parse_and_tag_entries(full_content, "Image Input")
    elif yt_url and ("youtube.com" in yt_url or "youtu.be" in yt_url):
        try:
            yt = YouTube(yt_url)
            video_title = yt.title
            final_title = title.strip() if title and title.strip() else video_title
            audio_stream = yt.streams.get_audio_only()
            with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_audio_file:
                audio_stream.download(filename=temp_audio_file.name)
                temp_audio_path = temp_audio_file.name
            content_text = transcribe_audio(temp_audio_path)
            content_source = f"YouTube: {video_title}"
            os.remove(temp_audio_path)
            full_content = f"Title: {final_title}\n\nContent: {content_text}"
            docs_to_add = parse_and_tag_entries(full_content, content_source)
        except Exception as e:
            return f"Error processing YouTube link: {e}"
    else:
        return "Please provide a title and content, or another input source."

    if not docs_to_add:
        return "No processable content found to add."
    
    if personal_vectorstore is None:
        personal_vectorstore = build_or_load_vectorstore(docs_to_add, PERSONAL_INDEX_PATH, is_personal=True)
    else:
        personal_vectorstore.add_documents(docs_to_add)
    
    personal_vectorstore.save_local(PERSONAL_INDEX_PATH)
    return f"Successfully added {len(docs_to_add)} new memory/memories."

def save_chat_to_memory(chat_history):
    global personal_vectorstore
    if not chat_history: return "Nothing to save."
    formatted_chat = []
    for message in chat_history:
        role = "User" if message["role"] == "user" else "Assistant"
        content = message["content"].strip()
        if content.startswith("*(Auto-detected context:"): continue
        formatted_chat.append(f"{role}: {content}")
    conversation_text = "\n".join(formatted_chat)
    if not conversation_text: return "No conversation content to save."
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    title = f"Conversation from {timestamp}"
    full_content = f"Title: {title}\n\nContent:\n{conversation_text}"
    doc_to_add = Document(page_content=full_content, metadata={"source": "Saved Chat", "title": title})
    if personal_vectorstore is None:
        personal_vectorstore = build_or_load_vectorstore([doc_to_add], PERSONAL_INDEX_PATH, is_personal=True)
    else:
        personal_vectorstore.add_documents([doc_to_add])
    personal_vectorstore.save_local(PERSONAL_INDEX_PATH)
    print(f"Saved conversation to long-term memory.")
    return f"Conversation from {timestamp} saved successfully!"

def list_personal_memories():
    global personal_vectorstore
    if personal_vectorstore is None or not hasattr(personal_vectorstore.docstore, '_dict') or not personal_vectorstore.docstore._dict:
        return gr.update(value=[["No memories to display", "", ""]]), gr.update(choices=["No memories to select"], value=None)
    docs = list(personal_vectorstore.docstore._dict.values())
    dataframe_data = [[doc.metadata.get('title', 'Untitled'), doc.metadata.get('source', 'Unknown'), doc.page_content] for doc in docs]
    dropdown_choices = [doc.page_content for doc in docs]
    return gr.update(value=dataframe_data), gr.update(choices=dropdown_choices)

def delete_personal_memory(memory_to_delete):
    global personal_vectorstore
    if personal_vectorstore is None or not memory_to_delete:
        return "Knowledge base is empty or no memory selected."
    all_docs = list(personal_vectorstore.docstore._dict.values())
    docs_to_keep = [doc for doc in all_docs if doc.page_content != memory_to_delete]
    if len(all_docs) == len(docs_to_keep):
        return "Error: Could not find the selected memory to delete."
    print(f"Deleting memory. {len(docs_to_keep)} memories remaining.")
    if not docs_to_keep:
        if os.path.isdir(PERSONAL_INDEX_PATH):
            shutil.rmtree(PERSONAL_INDEX_PATH)
        personal_vectorstore = build_or_load_vectorstore([], PERSONAL_INDEX_PATH, is_personal=True)
    else:
        new_vs = FAISS.from_documents(docs_to_keep, _default_embeddings())
        new_vs.save_local(PERSONAL_INDEX_PATH)
        personal_vectorstore = new_vs
    return "Successfully deleted memory. The list will now refresh."

def chat_fn(user_text, audio_file, settings, chat_history):
    global personal_vectorstore
    question = (user_text or "").strip()
    if audio_file and not question:
        try:
            voice_lang_name = settings.get("tts_lang", "English")
            voice_lang_code = CONFIG["languages"].get(voice_lang_name, "en")
            question = transcribe_audio(audio_file, lang=voice_lang_code)
        except Exception as e:
            err_msg = f"Audio Error: {e}" if settings.get("debug_mode") else "Sorry, I couldn't understand the audio."
            chat_history.append({"role": "assistant", "content": err_msg})
            return "", None, chat_history
    if not question:
        return "", None, chat_history
    chat_history.append({"role": "user", "content": question})
    manual_behavior_tag = settings.get("behaviour_tag")
    manual_emotion_tag = settings.get("emotion_tag")
    if manual_behavior_tag not in [None, "None"] or manual_emotion_tag not in [None, "None"]:
        scenario_tag, emotion_tag = manual_behavior_tag, manual_emotion_tag
    else:
        behavior_options = CONFIG.get("behavior_tags", [])
        emotion_options = CONFIG.get("emotion_tags", [])
        detected_tags = detect_tags_from_query(question, behavior_options=behavior_options, emotion_options=emotion_options)
        scenario_tag, emotion_tag = detected_tags.get("detected_behavior"), detected_tags.get("detected_emotion")
        if (scenario_tag and scenario_tag != "None") or (emotion_tag and emotion_tag != "None"):
            detected_msg = f"*(Auto-detected context: Behavior=`{scenario_tag}`, Emotion=`{emotion_tag}`)*"
            chat_history.append({"role": "assistant", "content": detected_msg})
    active_theme = settings.get("active_theme", "All")
    vs_general = ensure_index(active_theme)
    if personal_vectorstore is None:
        personal_vectorstore = build_or_load_vectorstore([], PERSONAL_INDEX_PATH, is_personal=True)
    rag_chain_settings = {"role": settings.get("role"), "temperature": settings.get("temperature"), "language": settings.get("language"), "patient_name": settings.get("patient_name"), "caregiver_name": settings.get("caregiver_name"), "tone": settings.get("tone"),}
    chain = make_rag_chain(vs_general, personal_vectorstore, **rag_chain_settings)
    if scenario_tag == "None": scenario_tag = None
    if emotion_tag == "None": emotion_tag = None
    simple_history = chat_history[:-1]
    response = answer_query(chain, question, chat_history=simple_history, scenario_tag=scenario_tag, emotion_tag=emotion_tag)
    answer = response.get("answer", "[No answer found]")
    chat_history.append({"role": "assistant", "content": answer})
    audio_out = None
    if settings.get("tts_on") and answer:
        tts_lang_code = CONFIG["languages"].get(settings.get("tts_lang"), "en")
        audio_out = synthesize_tts(answer, lang=tts_lang_code)
    from gradio import update
    return "", (update(value=audio_out, visible=bool(audio_out))), chat_history

def upload_knowledge(files, current_theme):
    if not files: return "No files were selected to upload."
    added = 0
    for f in files:
        try:
            copy_into_theme(current_theme, f.name); added += 1
        except Exception as e: print(f"Error uploading file {f.name}: {e}")
    if added > 0 and current_theme in vectorstores: del vectorstores[current_theme]
    return f"Uploaded {added} file(s). Refreshing file list..."
def save_file_selection(current_theme, enabled_files):
    man = load_manifest(current_theme)
    for fname in man['files']: man['files'][fname] = fname in enabled_files
    save_manifest(current_theme, man)
    if current_theme in vectorstores: del vectorstores[current_theme]
    return f"Settings saved. Index for theme '{current_theme}' will rebuild on the next query."
def refresh_file_list_ui(current_theme):
    files = list_theme_files(current_theme)
    enabled = [f for f, en in files if en]
    msg = f"Found {len(files)} file(s). {len(enabled)} enabled."
    return gr.update(choices=[f for f, _ in files], value=enabled), msg
def auto_setup_on_load(current_theme):
    theme_dir = theme_upload_dir(current_theme)
    if not os.listdir(theme_dir):
        print("First-time setup: Auto-seeding sample data...")
        seed_files_into_theme(current_theme)
    all_settings = collect_settings("patient", "", "", "warm", "English", "English", 0.7, "None", "None", "All", True, False)
    files_ui, status_msg = refresh_file_list_ui(current_theme)
    return all_settings, files_ui, status_msg

# In app.py, inside the Gradio Callbacks section for debugging
def test_save_file():
    """A simple function to test if we can write a file to the persistent storage."""
    try:
        # Get the directory where the personal index is supposed to be stored
        storage_dir = os.path.dirname(PERSONAL_INDEX_PATH)
        test_file_path = os.path.join(storage_dir, "persistence_test.txt")
        
        # Write the current time to the file
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        content = f"File saved successfully at: {current_time}"
        
        with open(test_file_path, "w", encoding="utf-8") as f:
            f.write(content)
            
        return f"✅ Success! Wrote test file to: {test_file_path}"
    except Exception as e:
        return f"❌ Error! Failed to write file. Reason: {e}"

def check_test_file():
    """A simple function to check if the test file from a previous session exists."""
    try:
        storage_dir = os.path.dirname(PERSONAL_INDEX_PATH)
        test_file_path = os.path.join(storage_dir, "persistence_test.txt")

        if os.path.exists(test_file_path):
            with open(test_file_path, "r", encoding="utf-8") as f:
                content = f.read()
            return f"✅ Success! Found test file. Contents: '{content}'"
        else:
            return f"❌ Failure. Test file not found at: {test_file_path}"
    except Exception as e:
        return f"❌ Error! Failed to check for file. Reason: {e}"



# --- UI Definition ---
CSS = ".gradio-container { font-size: 14px; } #chatbot { min-height: 250px; } #audio_out audio { max-height: 40px; } #audio_in audio { max-height: 40px; padding: 0; }"

with gr.Blocks(theme=gr.themes.Soft(), css=CSS) as demo:
    settings_state = gr.State({})
    with gr.Tab("Chat"):
        user_text = gr.Textbox(show_label=False, placeholder="Type your message here...")
        audio_in = gr.Audio(sources=["microphone"], type="filepath", label="Voice Input", elem_id="audio_in")
        with gr.Row():
            submit_btn = gr.Button("Send", variant="primary")
            save_btn = gr.Button("Save to Memory")
            clear_btn = gr.Button("Clear")
        chat_status = gr.Markdown()
        audio_out = gr.Audio(label="Response Audio", autoplay=True, visible=True, elem_id="audio_out")
        chatbot = gr.Chatbot(elem_id="chatbot", label="Conversation", type="messages")
    with gr.Tab("Personalize"):
        with gr.Accordion("Add to Personal Knowledge Base", open=True):
            gr.Markdown("Add personal notes, memories, or descriptions. A descriptive title helps the AI find memories more accurately.")
            personal_title = gr.Textbox(label="Title / Entry Name", placeholder="e.g., 'Dad's favorite songs'")
            personal_text = gr.Textbox(lines=5, label="Text Content (or use file upload)", placeholder="Type or paste text here. Use '—' on a new line to separate multiple entries.")
            personal_file = gr.File(label="Upload Audio/Video/Text File")
            personal_image = gr.Image(type="filepath", label="Upload Image")
            personal_yt_url = gr.Textbox(label="Or, provide a YouTube URL", placeholder="Paste a YouTube link here...")
            with gr.Row():
                personal_add_btn = gr.Button("Add Knowledge to Memory", variant="primary")
            personal_status = gr.Markdown()
        with gr.Accordion("Manage Personal Knowledge", open=False):
            personal_memory_display = gr.DataFrame(headers=["Title", "Source", "Content"], label="Saved Personal Memories", interactive=False, row_count=(5, "dynamic"))
            with gr.Row():
                personal_refresh_btn = gr.Button("Refresh Memories")
            with gr.Row():
                personal_delete_selector = gr.Dropdown(label="Select a memory to delete (by its full content)", scale=3, interactive=True)
                personal_delete_btn = gr.Button("Delete Selected Memory", variant="stop", scale=1)
            personal_delete_status = gr.Markdown()
    with gr.Tab("Settings"):
        with gr.Group():
            gr.Markdown("## Conversation & Persona Settings")
            with gr.Row():
                role = gr.Radio(CONFIG["roles"], value="caregiver", label="Your Role")
                temperature = gr.Slider(0.0, 1.2, value=0.7, step=0.1, label="Creativity")
                tone = gr.Dropdown(CONFIG["tones"], value="warm", label="Response Tone")
            with gr.Row():
                patient_name = gr.Textbox(label="Patient's Name", placeholder="e.g., 'Dad' or 'John'")
                caregiver_name = gr.Textbox(label="Caregiver's Name", placeholder="e.g., 'me' or 'Jane'")
                behaviour_tag = gr.Dropdown(CONFIG["behavior_tags"], value="None", label="Behaviour Filter (Manual Override)")
                emotion_tag = gr.Dropdown(CONFIG["emotion_tags"], value="None", label="Emotion Filter (Manual Override)")
        with gr.Accordion("Language, Voice & Debugging", open=False):
            language = gr.Dropdown(list(CONFIG["languages"].keys()), value="English", label="Response Language")
            tts_lang = gr.Dropdown(list(CONFIG["languages"].keys()), value="English", label="Voice Language")
            tts_on = gr.Checkbox(True, label="Enable Voice Response (TTS)")
            debug_mode = gr.Checkbox(False, label="Show Debug Info")
        gr.Markdown("--- \n ## General Knowledge Base Management")
        active_theme = gr.Radio(CONFIG["themes"], value="All", label="Active Knowledge Theme")
        with gr.Row():
            with gr.Column(scale=1):
                files_in = gr.File(file_count="multiple", file_types=[".jsonl", ".txt"], label="Upload Knowledge Files")
                upload_btn = gr.Button("Upload to Theme", variant="secondary")
                seed_btn = gr.Button("Import Sample Data", variant="secondary")
            with gr.Column(scale=2):
                mgmt_status = gr.Markdown()
                files_box = gr.CheckboxGroup(choices=[], label="Enable Files for the Selected Theme")
                with gr.Row():
                    save_files_btn = gr.Button("Save Selection", variant="primary")
                    refresh_btn = gr.Button("Refresh List")
        # --- ADD THE NEW DIAGNOSTIC TOOL AT THE BOTTOM ---
        with gr.Accordion("Persistence Test", open=False):
            gr.Markdown("Use this tool to verify that the Hugging Face persistent storage is working correctly. \n1. Click 'Run Test'. \n2. Manually restart the Space. \n3. Click 'Check for File'.")
            with gr.Row():
                test_save_btn = gr.Button("1. Run Persistence Test (Save File)")
                check_save_btn = gr.Button("3. Check for Test File")
            test_status = gr.Markdown()
            
    # --- Event Wiring ---
    all_settings_components = [role, patient_name, caregiver_name, tone, language, tts_lang, temperature, behaviour_tag, emotion_tag, active_theme, tts_on, debug_mode]
    for component in all_settings_components:
        component.change(fn=collect_settings, inputs=all_settings_components, outputs=settings_state)
    
    submit_btn.click(fn=chat_fn, inputs=[user_text, audio_in, settings_state, chatbot], outputs=[user_text, audio_out, chatbot])
    save_btn.click(fn=save_chat_to_memory, inputs=[chatbot], outputs=[chat_status])
    clear_btn.click(lambda: (None, None, [], None, "", ""), outputs=[user_text, audio_out, chatbot, audio_in, user_text, chat_status])
    
    personal_add_btn.click(
        fn=handle_add_knowledge,
        inputs=[personal_title, personal_text, personal_file, personal_image, personal_yt_url], 
        outputs=[personal_status]
    ).then(
        lambda: (None, None, None, None, None), 
        outputs=[personal_title, personal_text, personal_file, personal_image, personal_yt_url]
    )
    
    personal_refresh_btn.click(fn=list_personal_memories, inputs=None, outputs=[personal_memory_display, personal_delete_selector])
    personal_delete_btn.click(fn=delete_personal_memory, inputs=[personal_delete_selector], outputs=[personal_delete_status]).then(fn=list_personal_memories, inputs=None, outputs=[personal_memory_display, personal_delete_selector])
    
    upload_btn.click(upload_knowledge, inputs=[files_in, active_theme], outputs=[mgmt_status]).then(refresh_file_list_ui, inputs=[active_theme], outputs=[files_box, mgmt_status])
    save_files_btn.click(save_file_selection, inputs=[active_theme, files_box], outputs=[mgmt_status])
    seed_btn.click(seed_files_into_theme, inputs=[active_theme]).then(refresh_file_list_ui, inputs=[active_theme], outputs=[files_box, mgmt_status])
    refresh_btn.click(refresh_file_list_ui, inputs=[active_theme], outputs=[files_box, mgmt_status])
    active_theme.change(refresh_file_list_ui, inputs=[active_theme], outputs=[files_box, mgmt_status])
    demo.load(auto_setup_on_load, inputs=[active_theme], outputs=[settings_state, files_box, mgmt_status])
    test_save_btn.click(fn=test_save_file, inputs=None, outputs=[test_status])
    check_save_btn.click(fn=check_test_file, inputs=None, outputs=[test_status])

# --- Startup Logic ---
def pre_load_indexes():
    global personal_vectorstore
    print("Pre-loading all knowledge base indexes at startup...")
    for theme in CONFIG["themes"]:
        print(f"  - Loading general index for theme: '{theme}'")
        try:
            ensure_index(theme)
            print(f"    ...'{theme}' theme loaded successfully.")
        except Exception as e:
            print(f"    ...Error loading theme '{theme}': {e}")
    print("  - Loading personal knowledge index...")
    try:
        personal_vectorstore = build_or_load_vectorstore([], PERSONAL_INDEX_PATH, is_personal=True)
        print("    ...Personal knowledge loaded successfully.")
    except Exception as e:
        print(f"    ...Error loading personal knowledge: {e}")
    print("All indexes loaded. Application is ready.")

if __name__ == "__main__":
    pre_load_indexes()
    demo.queue().launch(debug=True)