src/streamlit_app.py

import streamlit as st

import altair as alt
import numpy as np
import pandas as pd
import rag
import os

cache_dir_base = "/app/.cache_app"
if not os.path.exists(cache_dir_base):
    try:
        os.makedirs(cache_dir_base, exist_ok=True)
    except Exception as e:
        st.error(f"Failed to create base cache directory {cache_dir_base}: {e}")

os.environ['HF_HOME'] = os.environ.get('HF_HOME', os.path.join(cache_dir_base, 'huggingface_hub'))
os.environ['TRANSFORMERS_CACHE'] = os.environ.get('TRANSFORMERS_CACHE', os.path.join(cache_dir_base, 'transformers'))
os.environ['SENTENCE_TRANSFORMERS_HOME'] = os.environ.get('SENTENCE_TRANSFORMERS_HOME', os.path.join(cache_dir_base, 'sentence_transformers'))
os.environ['PIP_CACHE_DIR'] = os.environ.get('PIP_CACHE_DIR', os.path.join(cache_dir_base, 'pip'))

st.set_page_config(layout="wide", page_title="Communication Board")

if 'assistant_initialized_flag' not in st.session_state:
    st.session_state.current_page = "main"
    st.session_state.show_custom_words = False
    st.session_state.custom_words = []
    st.session_state.text_size = 22
    st.session_state.theme = "Default"
    st.session_state.speech_rate = 1.0
    st.session_state.voice_option = "Default Voice"
    st.session_state.messages = []
    st.session_state.assistant = None
    st.session_state.text_output = ""
    st.session_state.assistant_initialized_flag = True

theme_colors = {
    "Default": {"bg": "#FFFFFF", "text": "#000000", "pronoun": "#FFFF99", "verb": "#CCFFCC", "question": "#CCCCFF", "common": "#FFCC99", "preposition": "#99CCFF", "descriptive": "#CCFF99", "misc": "#FFB6C1"},
    "High Contrast": {"bg": "#FFFFFF", "text": "#000000", "pronoun": "#FFFF00", "verb": "#00FF00", "question": "#0000FF", "common": "#FF6600", "preposition": "#00CCFF", "descriptive": "#66FF33", "misc": "#FF3366"},
    "Pastel": {"bg": "#F8F8FF", "text": "#333333", "pronoun": "#FFEFD5", "verb": "#E0FFFF", "question": "#D8BFD8", "common": "#FFE4B5", "preposition": "#B0E0E6", "descriptive": "#F0FFF0", "misc": "#FFF0F5"},
    "Dark Mode": {"bg": "#333333", "text": "#FFFFFF", "pronoun": "#8B8B00", "verb": "#006400", "question": "#00008B", "common": "#8B4500", "preposition": "#00688B", "descriptive": "#698B22", "misc": "#8B1A1A"}
}
colors = theme_colors[st.session_state.theme]

@st.cache_resource
def initialize_assistant(doc_path):
    if not os.path.exists(doc_path):
        st.sidebar.warning(f"Doc '{os.path.basename(doc_path)}' not found at '{doc_path}'. Creating dummy.")
        try:
            parent_dir = os.path.dirname(doc_path)
            if parent_dir: # Ensure parent_dir is not empty if doc_path is just a filename
                 os.makedirs(parent_dir, exist_ok=True)
            with open(doc_path, "w") as f:
                f.write("""Elliot Harper is a 29‐year‐old man whose life story is defined by both challenges and triumphs. Born with a congenital motor impairment that necessitates the use of a power wheelchair and an augmentative and alternative communication (AAC) device, Elliot has spent his life navigating a world that often sees only his physical limitations rather than his vibrant mind and resilient spirit.

Early Life and Education
From his earliest school days, Elliot experienced the sting of exclusion. Group projects, team sports, and school social events frequently left him feeling isolated and overlooked. Although his teachers recognized his quick wit and robust intellectual abilities—sometimes marveling at how he effortlessly absorbed new concepts using his AAC device—social stereotypes persisted. The very same classroom that fostered his love for learning also reinforced the painful notion that he did not quite fit the conventional mold. Despite these challenges, Elliot excelled academically and learned to use computers and other assistive technologies with remarkable speed, impressing both peers and educators alike.

Personal Relationships and Social Life
As Elliot grew older, his journey into adulthood brought new hurdles in the realm of personal relationships. His experiences with friendship and love have been as diverse as they have been challenging. He forged a deep bond with a former coworker—a friend who eventually became one of his most cherished confidants. Their relationship started in the workplace but blossomed into a genuine friendship marked by moments of humor, shared struggles, and subtle heartbreak. Though he once nurtured romantic feelings for her, Elliot eventually came to understand that differences in lifestyle—she loved outdoor adventures like running, biking, and exploring the trails—posed significant challenges in the dating arena.

Elliot's life is also enriched by the companionship of his attentive pet cat, whose gentle presence in the quiet moments of his day offers him profound comfort. Amid the long, solitary hours—sometimes marked by the anxiety of a brewing storm or the vulnerability of a cold night—both human connection and the soothing purr of his pet have been lifelines that remind him he is never truly alone.

Emotional Resilience and Coping Strategies
Deeply introspective, Elliot has developed a rich inner world to cope with the weight of societal misconceptions and moments of loneliness. His AAC device not only serves as a bridge to the outside world but also channels his creativity. Elliot finds solace and empowerment in writing—he often spends late nights crafting memoir notes that capture the nuances of his lived experience. Whether he is meticulously typing out thoughts on his computer or air-writing words with his arm in a burst of patient creativity, every communication is a testament to his determination to be heard and understood on his own terms.

Humor is another tool in his resilient arsenal. In moments of vulnerability—such as the awkward challenges of navigating social norms or the struggle to find the right words when discussing personal matters—Elliot turns to humor, sometimes researching witty pickup lines late at night to keep his spirits high. His ability to laugh in the face of hardship underlines the strength of his character.

Aspirations and Future Goals
Looking forward, Elliot dreams of a future where technology bridges the gap between physical challenges and boundless opportunities. He hopes to one day upgrade his mobility device—a power wheelchair, perhaps with advanced track capabilities—that would allow him greater independence and access to experiences that many take for granted. Professionally, his passion for technology and data has led him to master complex tools like spreadsheets, reinforcing his belief that his intellect is as dynamic as it is misunderstood.

Elliot's ultimate aspiration is to write a memoir that tells his true story—a narrative of perseverance and authenticity that not only reclaims his identity but also serves as an inspiration to others living with disabilities. He envisions his book as a beacon for those who have been sidelined by societal biases, proving that a vibrant life transcends physical limitations and that every individual has a unique, invaluable story to tell.

Navigating Love and Intimacy
In the realm of dating, Elliot’s journey has been bittersweet. The superficial judgments that often focus solely on his wheelchair have made it incredibly difficult to find a partner who sees past his disability and cherishes his intellect and kind-hearted nature. Elliot longs for a relationship where he is valued for his genuine self rather than merely his visible differences. Despite numerous setbacks and disappointments, his enduring hope for meaningful love reflects an unwavering belief in his own worth—a belief that continues to propel him forward.

In Summary
Elliot Harper’s life is a rich tapestry of struggle, strength, and quiet defiance. His story—marked by both moments of deep sorrow and bursts of unexpected joy—reminds us all that the measure of a person lies not in their physical limitations, but in the strength of their character, the depth of their compassion, and the courage with which they strive to realize their dreams. Through his journey with his AAC device, his cherished friendships, and his indomitable spirit, Elliot continues to challenge societal misconceptions and advocate for a world where everyone is seen and celebrated for who they truly are.
""")
        except Exception as e:
            st.sidebar.error(f"Failed to create dummy doc at '{doc_path}': {e}")
            return None
    try:
        st.sidebar.info("Initializing AAC Assistant... This may take some time.")
        assistant = rag.AACAssistant(doc_path)
        st.sidebar.success("AAC Assistant Initialized and ready!")
        return assistant
    except Exception as e:
        st.sidebar.error(f"Fatal Error Initializing AAC Assistant: {e}")
        st.sidebar.error("The assistant could not be started. Please check the logs for details.")
        return None

DEFAULT_DOCUMENT_PATH = "src/aac_user_experiences.txt"

if st.session_state.assistant is None:
    st.session_state.assistant = initialize_assistant(DEFAULT_DOCUMENT_PATH)

css = f"""
<style>
.big-font {{ font-size:{st.session_state.text_size}px !important; text-align: center; }}
.output-box {{ border: 2px solid #ddd; border-radius: 5px; padding: 15px; min-height: 100px; background-color: white; margin-bottom: 15px; font-size: {st.session_state.text_size}px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }}
div[data-testid="stHorizontalBlock"] {{ gap: 5px; }}
section[data-testid="stSidebar"] {{ width: 20rem !important; background-color: {colors["bg"]}; }}
body {{ background-color: {colors["bg"]}; color: {colors["text"]}; }}
.stButton>button {{ width: 100%; height: 60px; font-size: {max(16, st.session_state.text_size - 6)}px; font-weight: bold; white-space: normal; word-wrap: break-word; padding: 0px; transition: transform 0.1s ease; }}
.stButton>button:hover {{ filter: brightness(95%); transform: scale(1.03); box-shadow: 0 2px 3px rgba(0,0,0,0.1); }}
.control-button {{ height: 80px !important; font-size: {max(18, st.session_state.text_size - 4)}px !important; }}
.btn-pronoun {{ background-color: {colors["pronoun"]} !important; border: 1px solid #888 !important; }}
.btn-verb {{ background-color: {colors["verb"]} !important; border: 1px solid #888 !important; }}
.btn-question {{ background-color: {colors["question"]} !important; border: 1px solid #888 !important; }}
.btn-common {{ background-color: {colors["common"]} !important; border: 1px solid #888 !important; }}
.btn-preposition {{ background-color: {colors["preposition"]} !important; border: 1px solid #888 !important; }}
.btn-descriptive {{ background-color: {colors["descriptive"]} !important; border: 1px solid #888 !important; }}
.btn-misc {{ background-color: {colors["misc"]} !important; border: 1px solid #888 !important; }}
.sidebar .stChatMessage {{ background-color: {colors.get('bg', '#FFFFFF')}; border-radius: 8px; }}
</style>
"""
st.markdown(css, unsafe_allow_html=True)

js = """
<script>
function colorButtons() {
    const buttons = document.querySelectorAll('button[id^="key_"]');
    buttons.forEach(button => {
        const id = button.id;
        const parts = id.split('_');
        if (parts.length >= 4) {
            const category = parts[3];
            button.className = button.className.replace(/btn-[a-z]+/g, '').trim();
            button.classList.add('btn-' + category);
        }
    });
}
document.addEventListener('DOMContentLoaded', colorButtons);
const streamlitDoc = window.parent.document;
const observer = new MutationObserver((mutationsList, observer) => {
    for(const mutation of mutationsList) {
        if (mutation.type === 'childList' && mutation.addedNodes.length > 0) {
            colorButtons();
        }
    }
});
observer.observe(streamlitDoc.body, { childList: true, subtree: true });
colorButtons();
</script>
"""
st.markdown(js, unsafe_allow_html=True)

st.title("Communication Board")

st.session_state.text_output = st.text_area(
    "Compose Message:", value=st.session_state.get("text_output", ""), height=100, key="main_text_output_area"
)

layout = [
    [{"word": "I", "category": "pronoun"}, {"word": "am", "category": "verb"}, {"word": "how", "category": "question"}, {"word": "what", "category": "question"}, {"word": "when", "category": "question"}, {"word": "where", "category": "question"}, {"word": "who", "category": "question"}, {"word": "why", "category": "question"}, {"word": "That", "category": "pronoun"}, {"word": "Please", "category": "common"}],
    [{"word": "me", "category": "pronoun"}, {"word": "are", "category": "verb"}, {"word": "is", "category": "verb"}, {"word": "was", "category": "verb"}, {"word": "will", "category": "verb"}, {"word": "help", "category": "verb"}, {"word": "need", "category": "verb"}, {"word": "want", "category": "verb"}, {"word": "thank you", "category": "common"}, {"word": "sorry", "category": "common"}],
    [{"word": "my", "category": "pronoun"}, {"word": "can", "category": "verb"}, {"word": "A", "category": "misc"}, {"word": "B", "category": "misc"}, {"word": "C", "category": "misc"}, {"word": "D", "category": "misc"}, {"word": "E", "category": "misc"}, {"word": "F", "category": "misc"}, {"word": "G", "category": "misc"}, {"word": "H", "category": "misc"}],
    [{"word": "it", "category": "pronoun"}, {"word": "did", "category": "verb"}, {"word": "letter_I", "category": "misc", "display": "I"}, {"word": "J", "category": "misc"}, {"word": "K", "category": "misc"}, {"word": "L", "category": "misc"}, {"word": "M", "category": "misc"}, {"word": "N", "category": "misc"}, {"word": "O", "category": "misc"}, {"word": "P", "category": "misc"}],
    [{"word": "they", "category": "pronoun"}, {"word": "do", "category": "verb"}, {"word": "Q", "category": "misc"}, {"word": "R", "category": "misc"}, {"word": "S", "category": "misc"}, {"word": "T", "category": "misc"}, {"word": "U", "category": "misc"}, {"word": "V", "category": "misc"}, {"word": "W", "category": "misc"}, {"word": "X", "category": "misc"}],
    [{"word": "we", "category": "pronoun"}, {"word": "Y", "category": "misc"}, {"word": "Z", "category": "misc"}, {"word": "1", "category": "misc"}, {"word": "2", "category": "misc"}, {"word": "3", "category": "misc"}, {"word": "4", "category": "misc"}, {"word": "5", "category": "misc"}, {"word": ".", "category": "misc"}, {"word": "?", "category": "misc"}],
    [{"word": "you", "category": "pronoun"}, {"word": "6", "category": "misc"}, {"word": "7", "category": "misc"}, {"word": "8", "category": "misc"}, {"word": "9", "category": "misc"}, {"word": "0", "category": "misc"}, {"word": "-", "category": "misc"}, {"word": "!", "category": "misc"}, {"word": ",", "category": "misc"}, {"word": "SPACE", "category": "misc"}],
    [{"word": "your", "category": "pronoun"}, {"word": "like", "category": "verb"}, {"word": "to", "category": "preposition"}, {"word": "with", "category": "preposition"}, {"word": "in", "category": "preposition"}, {"word": "the", "category": "misc"}, {"word": "and", "category": "misc"}, {"word": "but", "category": "misc"}, {"word": "not", "category": "descriptive"}, {"word": "yes", "category": "common"}]
]

dynamic_layout = list(layout) # Make a mutable copy for adding custom words

if st.session_state.custom_words:
    current_custom_row = []
    for idx, word_info in enumerate(st.session_state.custom_words):
        word = word_info["word"]
        category = word_info["category"]
        current_custom_row.append({"word": f"custom_{idx}_{word}", "display": word, "category": category})
        if len(current_custom_row) == 10:
            dynamic_layout.append(list(current_custom_row))
            current_custom_row = []
    if current_custom_row:
        while len(current_custom_row) < 10:
            current_custom_row.append({"word": "", "category": "misc"})
        dynamic_layout.append(list(current_custom_row))

def add_to_output(word_to_add):
    current_text = st.session_state.get("text_output", "")
    if word_to_add == "SPACE":
        current_text += " "
    elif word_to_add in [".", "?", "!", ",", "-"]:
        current_text += word_to_add
    elif word_to_add.isdigit() or (len(word_to_add) == 1 and word_to_add.isalpha()):
        current_text += word_to_add
    else:
        if current_text and not current_text.endswith(" "):
            current_text += " "
        current_text += word_to_add
    st.session_state.text_output = current_text

st.markdown("### Communication Keyboard")
for row_idx, row_items in enumerate(dynamic_layout):
    cols = st.columns(len(row_items))
    for col_idx, item_info in enumerate(row_items):
        if not isinstance(item_info, dict) or "word" not in item_info or item_info["word"] == "":
            continue
        word = item_info["word"]
        category = item_info.get("category", "misc")
        display_text = item_info.get("display", word)
        button_key = f"key_{row_idx}_{col_idx}_{category}_{word.replace(' ', '_').replace('.', 'dot').replace('?', 'qmark')}" # Make key even more robust

        if cols[col_idx].button(
                display_text if word != "SPACE" else "␣",
                key=button_key,
                use_container_width=True
            ):
            word_to_add_on_click = display_text if word.startswith("custom_") or word.startswith("letter_") else word
            add_to_output(word_to_add_on_click)
            st.rerun()

col1, col2, col3, col4 = st.columns(4)

with col1:
    if st.button("CLEAR", key="clear_btn", help="Clear all text", use_container_width=True, type="secondary"):
        st.session_state.text_output = ""
        st.rerun()
with col2:
    if st.button("SPEAK", key="speak_btn", help="Speak the current text (feature placeholder)", use_container_width=True, type="secondary"):
        if st.session_state.text_output:
            st.toast(f"Speaking (simulated): {st.session_state.text_output}", icon="🔊")
    
    if st.button("⌫ DEL", key="backspace_btn", help="Delete last character", use_container_width=True):
        if st.session_state.text_output:
            st.session_state.text_output = st.session_state.text_output[:-1]
            st.rerun()
with col4:
    if st.button("⌫ WORD", key="backspace_word_btn", help="Delete last word", use_container_width=True):
        current_text = st.session_state.text_output.rstrip()
        if ' ' in current_text:
            st.session_state.text_output = current_text.rsplit(' ', 1)[0] + " "
        else:
            st.session_state.text_output = ""
        st.rerun()

with col3:
    if st.button("SEND 🗣️", key="send_btn", help="Send message to assistant", use_container_width=True, type="primary"):
        if 'assistant' not in st.session_state or st.session_state.assistant is None:
            st.error("AAC Assistant is not initialized. Please check sidebar for errors or wait for initialization to complete.")
        elif not st.session_state.text_output.strip():
            st.toast("Please compose a message first.", icon="✍️")
        else:
            user_message_to_send = st.session_state.text_output
            st.session_state.messages.append({"role": "user", "content": user_message_to_send})
            st.session_state.text_output = ""
            
            with st.spinner("Elliot is thinking... please wait a moment..."):
                try:
                    print(f"DEBUG [Streamlit UI]: User typed: '{user_message_to_send}'")
                    print(f"DEBUG [Streamlit UI]: Calling assistant.process_query for '{user_message_to_send}'")
                    assistant_response = st.session_state.assistant.process_query(user_message_to_send)
                    print(f"DEBUG [Streamlit UI]: Received assistant response: '{assistant_response}'")
                    st.session_state.messages.append({"role": "assistant", "content": assistant_response})
                except Exception as e:
                    error_message_display = f"An error occurred while processing your message: {e}"
                    st.error(error_message_display)
                    st.session_state.messages.append({"role": "assistant", "content": f"*Error: Could not get a response. Details: {e}*"})
            st.rerun()

with st.sidebar:
    st.divider()
    st.subheader("Conversation Log")
    chat_container = st.container(height=400 if st.session_state.messages else 100)
    with chat_container:
        if not st.session_state.messages:
            st.caption("No messages yet...")
        else:
            for message in st.session_state.messages:
                with st.chat_message(message["role"]):
                    st.markdown(message["content"])