import streamlit as st
import pandas as pd
import rag  # Import the rag module
import os   # Import os for file path handling

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

    # --- Session State Initialization ---
    # Check specifically for 'assistant' to ensure it's initialized
    if 'assistant' 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
        # Initialize text_output only if assistant is not initialized, assuming they go together
        st.session_state.text_output = ""

    # --- Theme Colors ---
    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]

    # --- Helper Function to Initialize Assistant (Adapted from previous main.py) ---
    @st.cache_resource # Cache the assistant
    def initialize_assistant(doc_path):
        """Initializes the AACAssistant."""
        # Create a dummy document if it doesn't exist for demonstration
        if not os.path.exists(doc_path):
            st.sidebar.warning(f"Doc '{os.path.basename(doc_path)}' not found. Creating dummy.")
            try:
                with open(doc_path, "w") as f:
                    f.write("""
                    I grew up in Seattle and love the rain.
                    My favorite hobby is playing chess.
                    I have a dog named Max.
                    I studied computer science.
                    I enjoy sci-fi movies.
                    """)
            except Exception as e:
                st.sidebar.error(f"Failed to create dummy doc: {e}")
                return None
        try:
            assistant = rag.AACAssistant(doc_path)
            st.sidebar.success("AAC Assistant Initialized.")
            return assistant
        except Exception as e:
            st.sidebar.error(f"Error initializing AAC Assistant: {e}")
            st.sidebar.error("Ensure Ollama/LM Studio running.")
            return None

    DEFAULT_DOCUMENT_PATH = "aac_user_experiences.txt"

    # --- CSS Styling ---
    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;
            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;
            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 chat styling */
        .sidebar .stChatMessage {{
            background-color: {colors.get('bg', '#FFFFFF')}; /* Use theme background */
            border-radius: 8px;
        }}
    </style>
    """

    # --- JS for Button Coloring (no delay, no setTimeout) ---
    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.classList.add('btn-' + category);
                }
            });
        }
        document.addEventListener('DOMContentLoaded', colorButtons);
        new MutationObserver(colorButtons).observe(document.body, { childList: true, subtree: true });
    </script>
    """

    st.markdown(css, unsafe_allow_html=True)
    st.markdown(js, unsafe_allow_html=True)

    # --- Keyboard Layout ---
    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"}
        ]
    ]

    # --- Add Custom Words ---
    if st.session_state.custom_words:
        custom_row = []
        for idx, word_info in enumerate(st.session_state.custom_words):
            word = word_info["word"]
            category = word_info["category"]
            custom_row.append({"word": f"custom_{idx}_{word}", "display": word, "category": category})
            if len(custom_row) == 10:
                layout.append(custom_row)
                custom_row = []
        if custom_row:
            while len(custom_row) < 10:
                custom_row.append({"word": "", "category": "misc"})
            layout.append(custom_row)

    # --- Initialize Assistant ---
    # Attempt initialization only once or if it failed previously
    if st.session_state.assistant is None: # This check is now safe
        st.session_state.assistant = initialize_assistant(DEFAULT_DOCUMENT_PATH)

    # --- Output Box (move to top, before keyboard) ---
    st.title("Communication Board")
    # Use st.text_area directly for input and display, bound to session state
    st.session_state.text_output = st.text_area(
        "Compose Message:", value=st.session_state.text_output, height=100, key="main_text_output"
    )

    # --- Keyboard Rendering (no delay, instant update) ---
    send_action = False # Flag for SEND button
    def add_to_output(word):
        if word == "SPACE":
            st.session_state.text_output += " "
        elif word in [".", "?", "!", ",", "-"]:
            st.session_state.text_output += word
        elif word.isdigit() or (len(word) == 1 and word.isalpha()):
            st.session_state.text_output += word
        else:
            if st.session_state.text_output and not st.session_state.text_output.endswith(" "):
                st.session_state.text_output += " "
            st.session_state.text_output += word

    st.markdown("### Communication Keyboard")
    for row_idx, row in enumerate(layout):
        cols = st.columns(len(row))
        for col_idx, item in enumerate(cols):
            word_info = row[col_idx]
            if "word" not in word_info or word_info["word"] == "":
                continue
            word = word_info["word"]
            category = word_info["category"]
            display = word_info.get("display", word)
            key = f"key_{row_idx}_{col_idx}_{category}_{word}"
            def make_callback(w=word, d=display):
                def cb():
                    if w.startswith("custom_") or w.startswith("letter_"):
                        add_to_output(d)
                    else:
                        add_to_output(w)
                return cb
            with cols[col_idx]:
                st.button(
                    display if word != "SPACE" else "␣",
                    key=key,
                    on_click=make_callback()
                )

    # --- Control Buttons ---
    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="primary"):
            st.session_state.text_output = ""
            st.rerun() # Rerun to reflect change
    with col2:
        if st.button("SPEAK", key="speak_btn", help="Speak the current text", use_container_width=True, type="primary"):
            if st.session_state.text_output:
                st.toast(f"Speaking: {st.session_state.text_output}", icon="🔊")
        if st.button("⌫ DEL", key="backspace", 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]
    with col4:
        if st.button("⌫ WORD", key="backspace_word", help="Delete last word", use_container_width=True):
            if st.session_state.text_output:
                words = st.session_state.text_output.rstrip().split()
                if words:
                    words.pop()
                    st.session_state.text_output = " ".join(words)
                    if words:
                        st.session_state.text_output += " "
    with col3: # Use the 3rd column for SEND
        if st.button("SEND", key="send_btn", help="Send message to assistant", use_container_width=True, type="primary"):
             if st.session_state.text_output:
                 send_action = True # Set flag to process sending


    # --- Settings Sidebar ---
    with st.sidebar:
        # --- Settings Section Commented Out ---
        # st.title("Settings")
        # st.subheader("Interface")
        # theme_options = list(theme_colors.keys())
        # theme_index = theme_options.index(st.session_state.theme)
        # new_theme = st.selectbox("Theme", theme_options, index=theme_index)
        # new_text_size = st.slider("Text Size", 12, 36, st.session_state.text_size)
        # if st.button("Apply Settings Changes", type="primary"):
        #     changed = False
        #     if new_theme != st.session_state.theme:
        #         st.session_state.theme = new_theme
        #         changed = True
        #     if new_text_size != st.session_state.text_size:
        #         st.session_state.text_size = new_text_size
        #         changed = True
        #     if changed:
        #         st.rerun()
        # with st.expander("Speech Settings"):
        #     speech_rate = st.slider("Rate", 0.5, 2.0, st.session_state.speech_rate, 0.1)
        #     voice_options = ["Default Voice", "Female Voice", "Male Voice", "Child Voice"]
        #     voice_index = voice_options.index(st.session_state.voice_option)
        #     voice = st.selectbox("Voice", voice_options, index=voice_index)
        #     if st.button("Apply Speech Settings"):
        #         st.session_state.speech_rate = speech_rate
        #         st.session_state.voice_option = voice
        # st.subheader("Custom Words")
        # with st.expander("Add New Word"):
        #     word = st.text_input("Word")
        #     cat = st.selectbox("Category", list(colors.keys()))
        #     col1, col2 = st.columns(2)
        #     with col1:
        #         if st.button("Add", key="add_word"):
        #             if word and cat:
        #                 st.session_state.custom_words.append({"word": word, "category": cat})
        #                 st.success(f"Added '{word}'")
        #                 st.rerun()
        #     if st.session_state.custom_words:
        #         st.write("Current custom words:")
        #         words_df = pd.DataFrame([
        #             {"Word": w["word"], "Category": w["category"]}
        #             for w in st.session_state.custom_words
        #         ])
        #         st.dataframe(words_df, hide_index=True)
        #         with col2:
        #             if st.button("Clear All", key="clear_words"):
        #                 st.session_state.custom_words = []
        #                 st.success("Words cleared")
        #                 st.rerun()
        # --- End of Settings Section Commented Out ---

        st.divider()
        # --- Chat History Display in Sidebar ---
        st.subheader("Conversation")
        chat_container = st.container(height=400) # Fixed height container
        with chat_container:
            for message in st.session_state.messages:
                with st.chat_message(message["role"]):
                    st.markdown(message["content"])

    # --- Process SEND Action ---
    if send_action and st.session_state.assistant:
        user_message = st.session_state.text_output
        st.session_state.messages.append({"role": "user", "content": user_message})

        # Process with AACAssistant
        try:
            # Get the response from the AACAssistant
            response = st.session_state.assistant.process_query(user_message)
            # Add assistant response to chat history
            st.session_state.messages.append({"role": "assistant", "content": response})
        except Exception as e:
            error_message = f"An error occurred: {e}"
            st.error(error_message) # Show error in main area
            st.session_state.messages.append({"role": "assistant", "content": f"*Error processing: {error_message}*"})

        # Clear the board's text area after sending
        st.session_state.text_output = ""
        st.rerun() # Rerun to update chat and clear board

if __name__ == "__main__":
    main()