frontend.py

import streamlit as st
import httpx
import json
import pandas as pd
import plotly.express as px

API_BASE = "http://localhost:8000/api"

st.set_page_config(page_title="Episodic Memory Platform", page_icon="🧠", layout="wide")

# --- Custom CSS ---
st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
    
    html, body, [class*="st-"] { font-family: 'Inter', sans-serif; }
    
    .main { background-color: #0e1117; }
    
    .memory-card {
        background: linear-gradient(135deg, #1a1f2e 0%, #151923 100%);
        border: 1px solid #2a2f3e;
        border-radius: 12px;
        padding: 16px 20px;
        margin-bottom: 10px;
    }
    .memory-card:hover { border-color: #4f8cff; }
    
    .score-badge {
        display: inline-block;
        background: linear-gradient(135deg, #4f8cff 0%, #6c5ce7 100%);
        color: white;
        padding: 2px 10px;
        border-radius: 20px;
        font-size: 0.75rem;
        font-weight: 600;
    }
    .tag-badge {
        display: inline-block;
        background: #2a2f3e;
        color: #8892b0;
        padding: 2px 8px;
        border-radius: 6px;
        font-size: 0.7rem;
        margin-right: 4px;
    }
    
    .episode-box {
        background: #12161f;
        border-left: 3px solid #4f8cff;
        border-radius: 0 8px 8px 0;
        padding: 12px 16px;
        margin-bottom: 12px;
    }
    
    .header-gradient {
        background: linear-gradient(90deg, #4f8cff 0%, #6c5ce7 50%, #a855f7 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-weight: 700;
    }
</style>
""", unsafe_allow_html=True)

# --- Header ---
st.markdown('<h1 class="header-gradient">🧠 Episodic Memory Platform</h1>', unsafe_allow_html=True)
st.caption("Powered by Endee Vector DB • OpenRouter LLM • Episodic RAG")

# --- Sidebar ---
with st.sidebar:
    st.markdown("### 👤 User Identity")
    username = st.text_input("Username", value="arjun", placeholder="Enter your username", key="username_input")
    st.caption(f"Endee index: `emp_{username.lower().replace(' ', '_')}`")
    st.divider()
    
    st.markdown("### Add Memory")
    
    upload_img = st.file_uploader("Upload Image (Optional)", type=["jpg", "jpeg", "png", "webp"])
    if upload_img:
        if st.button("🖼️ Save Visual Memory", use_container_width=True, type="primary"):
            with st.spinner("Analyzing image via OpenRouter & storing..."):
                try:
                    files = {"file": (upload_img.name, upload_img.getvalue(), upload_img.type)}
                    resp = httpx.post(f"{API_BASE}/memory/image", files=files, headers={"X-User-Id": username}, timeout=60.0)
                    if resp.status_code == 200:
                        data = resp.json()
                        st.success(f"Visual Memory Stored! ID: `{data['id'][:8]}...`")
                        st.session_state.upload_img = None
                    else:
                        st.error(f"Error: {resp.text[:200]}")
                except Exception as e:
                    st.error(f"Connection error: {e}")
                    
    st.markdown("---")
    mem_text = st.text_area("Memory text", placeholder="What happened today?", height=100)
    mem_tags = st.text_input("Tags (comma-separated)", placeholder="tech, personal")
    
    if st.button("💾 Save Text Memory", use_container_width=True, type="primary"):
        if mem_text.strip():
            tags = [t.strip() for t in mem_tags.split(",") if t.strip()] if mem_tags else []
            with st.spinner("Embedding & storing in Endee..."):
                try:
                    resp = httpx.post(f"{API_BASE}/memory/add", json={"text": mem_text, "tags": tags}, headers={"X-User-Id": username}, timeout=30.0)
                    if resp.status_code == 200:
                        data = resp.json()
                        st.success(f"Stored! ID: `{data['id'][:8]}...`")
                    else:
                        st.error(f"Error: {resp.text[:200]}")
                except Exception as e:
                    st.error(f"Connection error: {e}")
        else:
            st.warning("Enter some text first.")
    
    st.divider()
    st.markdown("### 🌙 Memory Reflection")
    st.write("Synthesize your short-term memories into profound insights.")
    if st.button("Run Reflection Cycle", use_container_width=True):
        with st.spinner("Initiating background reflection..."):
            try:
                resp = httpx.post(f"{API_BASE}/memory/reflect", headers={"X-User-Id": username}, timeout=10.0)
                if resp.status_code == 200:
                    st.success("Reflection running in the background! Check your memories later to see the new `insight`.")
                else:
                    st.error(f"Error {resp.status_code}: {resp.text[:200]}")
            except Exception as e:
                st.error(f"Connection error: {e}")
                
    st.divider()
    st.markdown("### Settings")
    intent = st.selectbox("Query Intent", ["auto", "recall", "summarize", "recommend"])

# Store headers for use in tabs
USER_HEADERS = {"X-User-Id": username}

# --- Main Area ---
tab1, tab2, tab3, tab4 = st.tabs(["🤖 Agent Query", "🔍 Raw Search", "🧭 Memory Map", "🗺️ Memory Wander"])

with tab1:
    query = st.text_input("Ask your memory agent", placeholder="What did I work on recently?", key="agent_q")
    
    if st.button("Ask Agent", type="primary", key="agent_btn"):
        if query.strip():
            payload = {"user_input": query}
            if intent != "auto":
                payload["intent"] = intent
                
            with st.status("Executing RAG Pipeline...", expanded=True) as status:
                st.write("🕵️‍♂️ **Step 1:** Analyzing user intent...")
                try:
                    resp = httpx.post(f"{API_BASE}/agent/ask", json=payload, headers=USER_HEADERS, timeout=60.0)
                    if resp.status_code == 200:
                        data = resp.json()
                        episodes = data.get("episodes_used", [])
                        
                        st.write("🧮 **Step 2:** Vectorizing intent & querying Endee VDB...")
                        st.write("⏳ **Step 3:** Applying time-decay re-ranking constraints...")
                        st.write(f"📅 **Step 4:** Grouping into {len(episodes)} coherent chronological episodes...")
                        st.write("🧠 **Step 5:** Generating grounded LLM response...")
                        
                        status.update(label="Response Generated Successfully!", state="complete", expanded=False)
                        
                        # Response
                        st.markdown("#### Agent Response")
                        st.markdown(f'<div class="memory-card">{data["response"]}</div>', unsafe_allow_html=True)
                        
                        # Episodes
                        if episodes:
                            st.markdown(f"#### Memory Context Extracted ({len(episodes)} Episodes)")
                            for ep in episodes:
                                with st.expander(f"📖 {ep['time_range']} — {len(ep['memories'])} memories"):
                                    for mem in ep["memories"]:
                                        tags_html = "".join([f'<span class="tag-badge">{t}</span>' for t in mem.get("tags", [])])
                                        st.markdown(f'<div class="memory-card">{mem["text"]}<br><br>{tags_html}</div>', unsafe_allow_html=True)
                    else:
                        status.update(label="Query Failed", state="error", expanded=True)
                        st.error(f"Error {resp.status_code}: {resp.text[:300]}")
                except Exception as e:
                    status.update(label="Connection Error", state="error", expanded=True)
                    st.error(f"Connection error: {e}")

with tab2:
    search_q = st.text_input("Semantic search query", placeholder="Search memories...", key="search_q")
    col1, col2 = st.columns(2)
    with col1:
        limit = st.slider("Results", 1, 20, 5)
    with col2:
        tags_filter = st.text_input("Filter by tags", placeholder="tech, endee", key="search_tags")
    
    if st.button("Search Endee", key="search_btn"):
        if search_q.strip():
            payload = {"query": search_q, "limit": limit}
            if tags_filter.strip():
                payload["tags_filter"] = [t.strip() for t in tags_filter.split(",") if t.strip()]
                
            with st.spinner("Querying Endee with metadata filters + time-decay re-ranking..."):
                try:
                    resp = httpx.post(f"{API_BASE}/memory/query", json=payload, headers=USER_HEADERS, timeout=30.0)
                    if resp.status_code == 200:
                        results = resp.json()
                        if results:
                            for r in results:
                                meta = r.get("metadata", {})
                                sem = r.get("_semantic_score", 0)
                                final = r.get("_final_score", 0)
                                
                                # Enhance badge visualization
                                decay_pct = ((sem - final) / sem * 100) if sem > 0 else 0
                                penalty_color = "#e74c3c" if decay_pct > 50 else "#f39c12" if decay_pct > 20 else "#2ecc71"
                                
                                tags_html = "".join([f'<span class="tag-badge">{t}</span>' for t in meta.get("tags", [])])
                                st.markdown(f"""
                                <div class="memory-card">
                                    <div>{meta.get('text', 'N/A')}</div>
                                    <div style="margin-top:12px; font-size: 0.85rem; color: #a1a1aa;">
                                        🕒 Timestamp: {meta.get('timestamp', 'N/A')}
                                    </div>
                                    <div style="margin-top:8px">
                                        <span class="score-badge" title="Raw Vector Cosine Similarity">🎯 Semantic: {sem:.4f}</span>
                                        <span class="score-badge" style="background: {penalty_color}" title="Score after {{e^(-λ·days_old)}} reduction">⏳ Time-Decayed: {final:.4f} (-{decay_pct:.1f}%)</span>
                                        {tags_html}
                                    </div>
                                </div>
                                """, unsafe_allow_html=True)
                        else:
                            st.info("No results found.")
                    else:
                        st.error(f"Error {resp.status_code}: {resp.text[:300]}")
                except Exception as e:
                    st.error(f"Connection error: {e}")

with tab3:
    st.markdown("### 🧭 Memory Map")
    st.write("An interactive 2D projection of your semantic memory space. Memories that are close together are semantically similar.")
    
    if st.button("Generate Memory Map", key="graph_btn"):
        with st.spinner("Fetching memories & computing PCA (this may take a minute)..."):
            try:
                resp = httpx.get(f"{API_BASE}/memory/all", headers=USER_HEADERS, timeout=120.0)
                if resp.status_code == 200:
                    results = resp.json()
                    if len(results) < 3:
                        st.warning("Not enough memories to generate a map. Add at least 3 distinct memories!")
                    else:
                        plot_data = []
                        for r in results:
                            meta = r.get("metadata", {})
                            text = meta.get("text", "")
                            cut_text = text[:80] + "..." if len(text) > 80 else text
                            tags_list = meta.get("tags", [])
                            primary_tag = tags_list[0] if tags_list else "untagged"
                            all_tags = ", ".join(tags_list) if tags_list else "untagged"
                            sim = r.get("similarity", 0.5)
                            
                            plot_data.append({
                                "text": cut_text,
                                "category": primary_tag,
                                "tags": all_tags,
                                "x": r.get("x", 0.0),
                                "y": r.get("y", 0.0),
                                "size": 12 + sim * 8,
                            })
                            
                        df = pd.DataFrame(plot_data)
                        
                        fig = px.scatter(
                            df, x='x', y='y',
                            color='category',
                            size='size',
                            hover_name='text',
                            hover_data={'x': False, 'y': False, 'size': False, 'tags': True, 'category': False},
                            title="Memory Vector Space (PCA 2D)",
                            template="plotly_dark",
                            color_discrete_sequence=px.colors.qualitative.Set2
                        )
                        
                        fig.update_layout(
                            margin=dict(l=20, r=20, b=20, t=50),
                            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False, title=""),
                            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False, title=""),
                            plot_bgcolor="rgba(0,0,0,0)",
                            paper_bgcolor="rgba(0,0,0,0)",
                            legend=dict(
                                title="Category",
                                bgcolor="rgba(30,30,50,0.7)",
                                bordercolor="#4f8cff",
                                borderwidth=1
                            ),
                            height=550,
                        )
                        
                        fig.update_traces(
                            marker=dict(line=dict(width=1, color='rgba(255,255,255,0.3)')),
                            textposition='top center'
                        )
                        
                        st.plotly_chart(fig, use_container_width=True)
                        st.caption(f"Showing {len(results)} memories • Colored by primary tag • Hover for details")
                else:
                    st.error(f"Error {resp.status_code}: {resp.text[:300]}")
            except Exception as e:
                st.error(f"Connection error: {e}")

with tab4:
    st.markdown("### 🗺️ Memory Wander")
    st.write("Pick a random thought from the database and explore its absolute closest and farthest cognitive neighbors.")
    
    if st.button("Surprise Me", type="primary", key="wander_btn"):
        with st.spinner("Wandering the vector space..."):
            try:
                resp = httpx.get(f"{API_BASE}/memory/wander", headers=USER_HEADERS, timeout=30.0)
                if resp.status_code == 200:
                    data = resp.json()
                    anchor = data.get("anchor")
                    if not anchor:
                        st.warning("Not enough memories yet to wander.")
                    else:
                        st.markdown("#### 🎯 Anchor Memory")
                        a_meta = anchor.get("metadata", {})
                        a_tags = "".join([f'<span class="tag-badge">{t}</span>' for t in a_meta.get("tags", [])])
                        st.markdown(f'<div class="memory-card" style="border-color:#a855f7">{a_meta.get("text", "")}<br><br>{a_tags}</div>', unsafe_allow_html=True)
                        
                        col1, col2 = st.columns(2)
                        with col1:
                            st.markdown("#### 🟢 Closest Neighbors")
                            for c in data.get("closest", []):
                                c_meta = c.get("metadata", {})
                                c_sem = c.get("similarity", 0)
                                c_tags = "".join([f'<span class="tag-badge">{t}</span>' for t in c_meta.get("tags", [])])
                                st.markdown(f'<div class="memory-card">{c_meta.get("text", "")}<br><br><span class="score-badge">sim: {c_sem:.4f}</span> {c_tags}</div>', unsafe_allow_html=True)
                                
                        with col2:
                            st.markdown("#### 🔴 Farthest Neighbors")
                            for f in data.get("farthest", []):
                                f_meta = f.get("metadata", {})
                                f_sem = f.get("similarity", 0)
                                f_tags = "".join([f'<span class="tag-badge">{t}</span>' for t in f_meta.get("tags", [])])
                                st.markdown(f'<div class="memory-card">{f_meta.get("text", "")}<br><br><span class="score-badge">sim: {f_sem:.4f}</span> {f_tags}</div>', unsafe_allow_html=True)
                else:
                    st.error(f"Error {resp.status_code}: {resp.text[:300]}")
            except Exception as e:
                st.error(f"Connection error: {e}")