import streamlit as st
import os
import glob
import re
import base64
import pytz
from datetime import datetime
from transformers.agents import CodeAgent, ReactCodeAgent, ReactJsonAgent, load_tool

# Page config
st.set_page_config(page_title="🌳✨ AI Knowledge Tree Builder 🛠️🤓", page_icon="🌳", layout="wide")
st.sidebar.markdown("""
🌳🤖 **AI Knowledge Tree Builder**
1. 📱 Universal access
2. ⚡ Rapid builds (<2min)
3. 🔗 Linked AI sources
4. 🎯 Lean core
5. 🧠 Shareable memory
6. 👤 Personalized
7. 🐦 Cloneable brevity
""")

# Initialize Agents
tools = [load_tool("text-to-speech"), load_tool("image_question_answering")]
agents = {
    "CodeCrafter": CodeAgent(tools=tools, system_prompt="Craft code like a pro! 🖥️"),
    "StepSage": ReactCodeAgent(tools=tools, system_prompt="Step-by-step wisdom! 🧠"),
    "JsonJugger": ReactJsonAgent(tools=tools, system_prompt="JSON-powered antics! 🤡"),
    "OutlineOracle": ReactCodeAgent(tools=tools, system_prompt="Outline everything! 📋"),
    "ToolTitan": CodeAgent(tools=tools, system_prompt="List tools with swagger! 🔧"),
    "SpecSpinner": ReactJsonAgent(tools=tools, system_prompt="Spin specs with style! 📜"),
    "ImageImp": CodeAgent(tools=tools, system_prompt="Dream up image prompts! 🎨"),
    "VisualVortex": ReactCodeAgent(tools=tools, system_prompt="Generate visuals! 🖼️"),
    "GlossGuru": ReactJsonAgent(tools=tools, system_prompt="Define terms with pizzazz! 📖"),
}

# MoE Prompts mapped to agents
moe_prompts = {
    "Create a python streamlit app.py...": "CodeCrafter",
    "Create a python gradio app.py...": "CodeCrafter",
    "Create a mermaid model...": "OutlineOracle",
    "Create a top three list of tools...": "ToolTitan",
    "Create a specification in markdown...": "SpecSpinner",
    "Create an image generation prompt...": "ImageImp",
    "Generate an image which describes...": "VisualVortex",
    "List top ten glossary terms...": "GlossGuru",
    "": "StepSage"  # Default agent for blank selection
}

# Session state
if 'selected_file' not in st.session_state:
    st.session_state.selected_file = None
if 'view_mode' not in st.session_state:
    st.session_state.view_mode = 'view'
if 'files' not in st.session_state:
    st.session_state.files = []

# MoE Selection
moe_options = list(moe_prompts.keys())
selected_moe = st.selectbox("Pick an MoE Adventure! 🎲", moe_options, index=0, key="selected_moe")

# Utility Functions (simplified for brevity, reusing yours where unchanged)
def sanitize_filename(text):
    return re.sub(r'[^\w\s-]', ' ', text.strip())[:50]

def generate_timestamp_filename(query):
    central = pytz.timezone('US/Central')
    now = datetime.now(central)
    return f"{now.strftime('%I%M%p %m%d%Y')} ({sanitize_filename(query)}).md"

def save_ai_interaction(query, result, is_rerun=False):
    filename = generate_timestamp_filename(query)
    content = f"# {'Rerun' if is_rerun else 'Query'}: {query}\n\n## AI Response\n{result}"
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(content)
    return filename

def run_agent(task, agent_name):
    agent = agents[agent_name]
    if isinstance(agent, CodeAgent):
        return agent.run(task, return_generated_code=True)
    return agent.run(task)

# Main App
def main():
    st.markdown("### 🌳 AI Knowledge Tree Builder 🧠🌱 Let’s Grow Some Smarts!")
    query = st.text_input("Ask Away! 🤔", placeholder="E.g., 'Explain transformers!'")
    
    if query:
        agent_name = moe_prompts[selected_moe]
        try:
            result = run_agent(f"{selected_moe} {query}" if selected_moe else query, agent_name)
            st.markdown(f"🎉 {agent_name} says: {result}")
            saved_file = save_ai_interaction(query, result)
            st.success(f"Saved to {saved_file}")
            st.session_state.selected_file = saved_file
        except Exception as e:
            st.error(f"😱 Oops! {e}")

    # File Management (simplified)
    st.sidebar.title("📁 Files")
    md_files = sorted([f for f in glob.glob("*.md") if f.lower() != 'readme.md'])
    for file in md_files:
        st.sidebar.write(file)
        if st.sidebar.button(f"View {file}", key=f"v_{file}"):
            st.session_state.selected_file = file
    if st.session_state.selected_file:
        with open(st.session_state.selected_file, 'r') as f:
            st.markdown(f.read())

if __name__ == "__main__":
    main()