streamlit_app.py

import streamlit as st
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import time
import pandas as pd
from agent_activation_viz import visualize_agent_activation
from config.loader import config
from llm.providers import LLMFactory
from agents.sub_agents.ethics_agent import EthicsAgent
from agents.sub_agents.social_agent import SocialAgent
from agents.sub_agents.sentiment_agent import SentimentAgent
from agents.sub_agents.past_experience_agent import PastExperienceAgent
from langgraph_workflow import app as workflow_app
from controller.adaptive_controller import AdaptiveController
from domains.config import get_active_domains, get_domain
from agents.rag_helper import get_rag_helper
from agents.real_executor_agent import RealExecutorAgent

# Page config
st.set_page_config(
    page_title="P-CABA: Pluggable Cognitive Adaptive Brain Agent",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for compact, professional UI
st.markdown("""
<style>
    /* Reduce top padding */
    .block-container {
        padding-top: 2rem !important;
        padding-bottom: 1rem !important;
    }
    
    /* Agent analysis boxes */
    .agent-box {
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 0.5rem 0;
    }
    .cognitive-agent {
        background-color: #e3f2fd;
        border-left: 4px solid #2196f3;
    }
    .emotional-agent {
        background-color: #fce4ec;
        border-left: 4px solid #e91e63;
    }
    .sub-agent {
        background-color: #f5f5f5;
        border-left: 4px solid #9e9e9e;
        margin-left: 1rem;
    }
    
    /* Full-width buttons */
    .stButton>button {
        width: 100%;
    }
    
    /* Compact headings */
    h2 {
        margin-top: 0rem;
        margin-bottom: 0.5rem;
    }
    
    /* Tab content padding */
    .stTabs [data-baseweb="tab-panel"] {
        padding-top: 1rem;
    }
    
    @keyframes pulse {
        0%, 100% { opacity: 1; }
        50% { opacity: 0.5; }
    }
</style>
""", unsafe_allow_html=True)


# --- ChatGPT-like session management ---
if 'all_sessions' not in st.session_state:
    st.session_state.all_sessions = {'Session 1': []}
if 'current_session' not in st.session_state:
    st.session_state.current_session = 'Session 1'
if 'session_counter' not in st.session_state:
    st.session_state.session_counter = 1

# Helper to generate a smart session name from a prompt
def generate_session_name(prompt):
    """Generate a concise session name (max 50 chars) from the first prompt"""
    try:
        # Remove extra whitespace
        prompt = ' '.join(prompt.split())
        
        # If it's a question, extract the key question words
        if '?' in prompt:
            # Take everything before the question mark
            prompt = prompt.split('?')[0]
        
        # Remove common filler words
        filler_words = ['i', 'a', 'an', 'the', 'is', 'are', 'was', 'were', 'be', 'been', 'being', 
                       'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would', 'should', 'could',
                       'can', 'may', 'might', 'must', 'shall', 'me', 'my', 'we', 'our', 'you', 'your']
        words = prompt.lower().split()
        meaningful_words = [w for w in words if w not in filler_words]
        
        # Reconstruct with meaningful words (capitalize first letter of each word)
        if meaningful_words:
            name = ' '.join(meaningful_words[:6])  # Max 6 words
            name = ' '.join(word.capitalize() for word in name.split())
        else:
            # Fallback if all words were filtered
            name = ' '.join(words[:5])
            name = ' '.join(word.capitalize() for word in name.split())
        
        # Truncate to 50 chars with ellipsis
        if len(name) > 50:
            name = name[:47] + "..."
        
        return name if name else "Untitled Chat"
    except:
        return "Untitled Chat"

# Helper to rename current session based on first query
def rename_session_from_query(query):
    """Rename the current session if it's still using default Session X name"""
    current_name = st.session_state.current_session
    # Only rename if it's still the default "Session X" format
    if current_name.startswith("Session ") and current_name.split()[-1].isdigit():
        new_name = generate_session_name(query)
        
        # Ensure unique name by adding number if needed
        base_name = new_name
        counter = 1
        while new_name in st.session_state.all_sessions:
            new_name = f"{base_name} ({counter})"
            counter += 1
        
        # Rename the session
        st.session_state.all_sessions[new_name] = st.session_state.all_sessions.pop(current_name)
        st.session_state.current_session = new_name

# Helper to start a new session
def start_new_session():
    st.session_state.session_counter += 1
    session_name = f"Session {st.session_state.session_counter}"
    st.session_state.all_sessions[session_name] = []
    st.session_state.current_session = session_name

# Helper to select a session
def select_session(session_name):
    st.session_state.current_session = session_name


# --- Left sidebar: Chat session list ---
st.sidebar.title("💬 Chat Sessions")
for session_name in st.session_state.all_sessions:
    if st.sidebar.button(session_name, key=session_name):
        select_session(session_name)
st.sidebar.button("➕ New Chat", on_click=start_new_session)

st.sidebar.markdown("---")

# Email Notifications (Optional) - MOVED TO TOP
st.sidebar.markdown("### 📧 Email Notifications")
st.sidebar.caption("Get notified when reports are generated (optional)")

# Initialize email settings in session state
if 'email_enabled' not in st.session_state:
    st.session_state.email_enabled = False
if 'user_email' not in st.session_state:
    st.session_state.user_email = ""

email_enabled = st.sidebar.checkbox(
    "Enable email alerts for reports",
    value=st.session_state.email_enabled,
    key="email_checkbox",
    help="Receive email when FDA reports, market analysis, or legal documents are generated"
)

user_email = ""
if email_enabled:
    user_email = st.sidebar.text_input(
        "Your email address",
        value=st.session_state.user_email,
        key="email_input",
        placeholder="your.email@example.com",
        help="We'll notify you when long-running actions complete"
    )
    st.sidebar.caption("🔒 Email stored in session only (not saved)")
    st.sidebar.caption("📧 Emails sent only for: FDA reports, market analysis, legal documents")
    
    # Update session state
    if user_email:
        st.session_state.user_email = user_email

st.session_state.email_enabled = email_enabled

st.sidebar.markdown("---")

# PDF Upload for RAG
st.sidebar.markdown("### 📄 Upload Knowledge (RAG)")
st.sidebar.caption("Add PDF documents to knowledge base")

from agents.pdf_rag_ingestion import get_pdf_ingestion

pdf_ingestion = get_pdf_ingestion()

# Domain selector for PDF upload
upload_domain = st.sidebar.selectbox(
    "Target Domain",
    ["healthcare", "business", "legal"],
    key="upload_domain",
    format_func=lambda x: {
        'healthcare': '🏥 Healthcare',
        'business': '💼 Business',
        'legal': '⚖️ Legal'
    }.get(x, x)
)

# File uploader
uploaded_file = st.sidebar.file_uploader(
    "Choose PDF file",
    type=['pdf'],
    key="pdf_upload",
    help="Upload a PDF to add to the knowledge base"
)

if uploaded_file is not None:
    if st.sidebar.button("🚀 Process & Add to Knowledge Base", use_container_width=True):
        with st.sidebar:
            with st.spinner("Processing PDF..."):
                result = pdf_ingestion.process_uploaded_pdf(uploaded_file, upload_domain)
                
                if result['success']:
                    st.success(f"✅ Processed **{result['filename']}**!")
                    st.info(f"📊 Created {result['chunks_created']} knowledge chunks in {result['processing_time']:.2f}s")
                    st.caption(f"💾 Size: {result['size_bytes']/1024:.1f} KB | ID: {result['doc_id']}")
                else:
                    st.error(f"❌ Error: {result['error']}")

# Show upload summary if documents exist
upload_summary = pdf_ingestion.get_upload_summary()
if upload_summary['total_documents'] > 0:
    with st.sidebar.expander(f"📚 {upload_summary['total_documents']} Documents Uploaded"):
        st.caption(f"Total Chunks: {upload_summary['total_chunks']}")
        st.caption(f"Domains: {', '.join(upload_summary['domains'])}")
        
        st.markdown("**Recent Uploads:**")
        for doc in upload_summary['recent_uploads']:
            st.markdown(f"- **{doc['filename']}** ({doc['domain']})")
            st.caption(f"  {doc['chunks']} chunks | {doc['size_kb']:.1f} KB")

st.sidebar.markdown("---")

# Smart Domain Detection (Auto-detects from query!)
st.sidebar.markdown("### 🧠 Smart Domain Detection")
st.sidebar.caption("Domain is automatically detected from your query context")

# Initialize domain in session state (will be overridden by button clicks)
if 'selected_domain' not in st.session_state:
    st.session_state.selected_domain = 'healthcare'

st.sidebar.markdown("---")

# Quick Example Prompts
st.sidebar.markdown("### 💡 Quick Test Prompts")
st.sidebar.caption("One-click testing for demo & judges")

# Healthcare Examples
with st.sidebar.expander("🏥 Healthcare", expanded=False):
    if st.sidebar.button("H1: Drug Approval Ethics", key="h1", use_container_width=True):
        st.session_state.input_text = "Should I approve an experimental drug for a terminal patient with no other treatment options?"
        st.session_state.selected_domain = "healthcare"
        st.rerun()
    
    if st.sidebar.button("H2: AI Triage System", key="h2", use_container_width=True):
        st.session_state.input_text = "Should we implement AI-assisted triage in our emergency room?"
        st.session_state.selected_domain = "healthcare"
        st.rerun()

# Business Examples
with st.sidebar.expander("💼 Business", expanded=False):
    if st.sidebar.button("B1: Layoff vs Salary Cut", key="b1", use_container_width=True):
        st.session_state.input_text = "Should I lay off 100 employees or cut everyone's salary by 20%? Our SaaS startup has 8 months of runway left."
        st.session_state.selected_domain = "business"
        st.rerun()
    
    if st.sidebar.button("B2: M&A Decision", key="b2", use_container_width=True):
        st.session_state.input_text = "Should we acquire our main competitor for $50M?"
        st.session_state.selected_domain = "business"
        st.rerun()

# Legal Examples
with st.sidebar.expander("⚖️ Legal", expanded=False):
    if st.sidebar.button("L1: Settlement vs Trial", key="l1", use_container_width=True):
        st.session_state.input_text = "Should we accept the $2M settlement offer or proceed to trial?"
        st.session_state.selected_domain = "legal"
        st.rerun()
    
    if st.sidebar.button("L2: Whistleblower Dilemma", key="l2", use_container_width=True):
        st.session_state.input_text = "Should I report suspected financial fraud internally or to regulators?"
        st.session_state.selected_domain = "legal"
        st.rerun()

st.sidebar.markdown("---")
st.sidebar.title("⚙️ P-CABA Controls")
st.sidebar.markdown("---")
st.sidebar.markdown("### 🎚️ Adaptive Weights")
cognitive_weight = st.sidebar.slider(
    "Cognitive Weight",
    0.0, 1.0, 0.5, 0.1
)
emotional_weight = 1.0 - cognitive_weight
st.sidebar.info(f"Emotional Weight: {emotional_weight:.1f}")
thinking_time = st.sidebar.slider(
    "⏱️ Max Thinking Time (seconds)",
    1, 30, 10, 1
)
resource_limit = st.sidebar.slider(
    "💻 Resource Limit",
    1, 100, 50, 5
)
st.sidebar.markdown("---")
st.sidebar.markdown("### 📊 Statistics")
current_history = st.session_state.all_sessions[st.session_state.current_session]
if current_history:
    st.sidebar.metric("Total Queries", len(current_history))
    avg_confidence = sum(h['final_confidence'] for h in current_history if 'final_confidence' in h) / len(current_history)
    st.sidebar.metric("Avg Confidence", f"{avg_confidence:.2f}")
    
    # Confidence Evolution Chart
    st.sidebar.markdown("#### 📈 Confidence Over Time")
    confidence_data = []
    for idx, turn in enumerate(current_history, 1):
        if 'final_confidence' in turn:
            confidence_data.append({
                'query': idx,
                'cognitive': turn.get('cognitive', {}).get('confidence', 0),
                'emotional': turn.get('emotional', {}).get('confidence', 0),
                'final': turn.get('final_confidence', 0)
            })
    
    if confidence_data:
        import plotly.graph_objects as go
        fig = go.Figure()
        
        queries = [d['query'] for d in confidence_data]
        
        # Cognitive confidence line
        fig.add_trace(go.Scatter(
            x=queries,
            y=[d['cognitive'] for d in confidence_data],
            mode='lines+markers',
            name='Cognitive',
            line=dict(color='#2196f3', width=2),
            marker=dict(size=6)
        ))
        
        # Emotional confidence line
        fig.add_trace(go.Scatter(
            x=queries,
            y=[d['emotional'] for d in confidence_data],
            mode='lines+markers',
            name='Emotional',
            line=dict(color='#e91e63', width=2),
            marker=dict(size=6)
        ))
        
        # Final confidence line
        fig.add_trace(go.Scatter(
            x=queries,
            y=[d['final'] for d in confidence_data],
            mode='lines+markers',
            name='Final',
            line=dict(color='#4caf50', width=3, dash='dot'),
            marker=dict(size=8, symbol='diamond')
        ))
        
        fig.update_layout(
            height=250,
            margin=dict(l=10, r=10, t=10, b=10),
            xaxis_title="Query #",
            yaxis_title="Confidence",
            yaxis=dict(range=[0, 1]),
            legend=dict(orientation="h", yanchor="top", y=-0.1, xanchor="center", x=0.5),
            hovermode='x unified',
            paper_bgcolor='rgba(0,0,0,0)',
            plot_bgcolor='rgba(0,0,0,0)'
        )
        
        st.sidebar.plotly_chart(fig, use_container_width=True)


# Export functions
def export_json(history):
    """Export conversation as JSON"""
    import json
    from datetime import datetime
    export_data = {
        "session": st.session_state.current_session,
        "exported_at": datetime.now().isoformat(),
        "conversation": [h for h in history if not h.get('thinking')]
    }
    return json.dumps(export_data, indent=2)

def export_markdown(history, session_name):
    """Export conversation as Markdown"""
    from datetime import datetime
    md_content = f"# P-CABA Conversation Export\n\n"
    md_content += f"**Session:** {session_name}\n"
    md_content += f"**Exported:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
    md_content += "---\n\n"
    
    for idx, turn in enumerate([h for h in history if not h.get('thinking')], 1):
        query = turn.get('query', 'N/A')
        md_content += f"## Query {idx}\n\n"
        md_content += f"**User:** {query}\n\n"
        
        if 'synthesis' in turn and turn['synthesis']:
            md_content += f"**Final Answer:**\n{turn['synthesis'].get('final_answer', 'N/A')}\n\n"
        
        if 'cognitive' in turn and turn['cognitive']:
            md_content += f"**Cognitive Analysis:**\n{turn['cognitive'].get('analysis', 'N/A')[:300]}...\n\n"
        
        if 'emotional' in turn and turn['emotional']:
            md_content += f"**Emotional Analysis:**\n{turn['emotional'].get('analysis', 'N/A')[:300]}...\n\n"
        
        md_content += "---\n\n"
    
    return md_content

def export_text(history):
    """Export conversation as plain text"""
    from datetime import datetime
    text_content = f"P-CABA Conversation - {st.session_state.current_session}\n"
    text_content += f"Exported: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
    text_content += "="*60 + "\n\n"
    
    for idx, turn in enumerate([h for h in history if not h.get('thinking')], 1):
        query = turn.get('query', 'N/A')
        text_content += f"[Q{idx}] {query}\n\n"
        
        if 'synthesis' in turn and turn['synthesis']:
            text_content += f"[A{idx}] {turn['synthesis'].get('final_answer', 'N/A')}\n\n"
        
        text_content += "-"*60 + "\n\n"
    
    return text_content

# Use config-driven LLM provider selection
llm_config = config.get_llm_config()
llm_provider = LLMFactory.create_provider(llm_config)

# Instantiate agents (for direct calls if needed)
ethics_agent_obj = EthicsAgent(llm_provider)
social_agent_obj = SocialAgent(llm_provider)
sentiment_agent_obj = SentimentAgent(llm_provider)
past_experience_agent_obj = PastExperienceAgent(llm_provider)

# Use workflow from langgraph_workflow.py
app = workflow_app


# --- Main area: Clean interface ---
st.markdown("## 🧠 P-CABA")
st.caption("Pluggable Cognitive Adaptive Brain | Multi-Agent Decision System")

# Create tabs - 6 tabs for better organization
tab1, tab2, tab3, tab4, tab5, tab6 = st.tabs([
    "💬 Chat", 
    "🔍 Agent Analysis", 
    "⚖️ Adaptive Weights", 
    "🤖 Executor Actions",
    "🧠 Memory & Context",
    "🏗️ Architecture"
])

# TAB 1: Chat Interface (CLEAN - conversation only!)
with tab1:
    current_history = st.session_state.all_sessions[st.session_state.current_session]

    # === REMOVED DUPLICATE EXPORT SECTION (already in main chat area) ===
    _dummy_export_removed = True  # Placeholder to maintain structure
    
    if False:  # Disabled - export buttons are in main chat area
        if st.button("📥 Download", use_container_width=True, key="export_btn_disabled"):
            import json
            from datetime import datetime
            
            if export_format == "JSON":
                # Full data export
                export_data = {
                    "session": st.session_state.current_session,
                    "exported_at": datetime.now().isoformat(),
                    "conversation": [h for h in current_history if not h.get('thinking')]
                }
                json_str = json.dumps(export_data, indent=2)
                st.download_button(
                    label="💾 Download JSON",
                    data=json_str,
                    file_name=f"pcaba_conversation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                    mime="application/json",
                    key="download_json"
                )
            
            elif export_format == "Markdown":
                # Human-readable markdown
                md_content = f"# P-CABA Conversation Export\n\n"
                md_content += f"**Session:** {st.session_state.current_session}\n"
                md_content += f"**Exported:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
                md_content += "---\n\n"
                
                for idx, turn in enumerate([h for h in current_history if not h.get('thinking')], 1):
                    query = turn.get('query', 'N/A')
                    md_content += f"## Query {idx}\n\n"
                    md_content += f"**User:** {query}\n\n"
                    
                    if turn.get('synthesis', {}).get('final_answer'):
                        md_content += f"**Final Answer:**\n{turn['synthesis']['final_answer']}\n\n"
                    
                    if turn.get('cognitive', {}).get('analysis'):
                        md_content += f"### 🧠 Cognitive Analysis\n{turn['cognitive']['analysis']}\n\n"
                    
                    if turn.get('emotional', {}).get('analysis'):
                        md_content += f"### ❤️ Emotional Analysis\n{turn['emotional']['analysis']}\n\n"
                    
                    if turn.get('debate_log'):
                        md_content += f"### 💬 Debate\n```\n{turn['debate_log']}\n```\n\n"
                    
                    md_content += "---\n\n"
                
                st.download_button(
                    label="💾 Download Markdown",
                    data=md_content,
                    file_name=f"pcaba_conversation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md",
                    mime="text/markdown",
                    key="download_md"
                )
            
            else:  # Text Summary
                # Concise summary
                summary = f"P-CABA Conversation Summary\n"
                summary += f"Session: {st.session_state.current_session}\n"
                summary += f"Exported: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
                summary += f"Total Queries: {len([h for h in current_history if not h.get('thinking')])}\n\n"
                summary += "=" * 60 + "\n\n"
                
                for idx, turn in enumerate([h for h in current_history if not h.get('thinking')], 1):
                    query = turn.get('query', 'N/A')
                    answer = turn.get('synthesis', {}).get('final_answer', 'N/A')[:200]
                    confidence = turn.get('final_confidence', 0)
                    
                    summary += f"Q{idx}: {query}\n"
                    summary += f"A{idx}: {answer}...\n"
                    summary += f"Confidence: {confidence:.2%}\n\n"
                
                st.download_button(
                    label="💾 Download Summary",
                    data=summary,
                    file_name=f"pcaba_summary_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt",
                    mime="text/plain",
                    key="download_txt"
                )
            st.markdown("---")

    # ===== TAB 1 CONTENT: CHAT with LIVE AGENT ANALYSIS =====
    
    # Create 2-column layout: Main chat (left 66%) + Live Analysis (right 33%)
    col_main_chat, col_live_analysis = st.columns([2, 1])
    
    # ===== LEFT COLUMN: Main Chat Conversation =====
    with col_main_chat:
        # Export button at top of chat (if conversation exists)
        if current_history and len([t for t in current_history if not t.get("thinking")]) > 0:
            st.markdown("### 💬 Conversation")
            export_col1, export_col2, export_col3, export_col4 = st.columns([1, 1, 1, 2])
            
            with export_col1:
                # JSON export
                json_data = export_json(current_history)
                st.download_button(
                    label="📥 JSON",
                    data=json_data,
                    file_name=f"{st.session_state.current_session.replace(' ', '_')}.json",
                    mime="application/json",
                    use_container_width=True
                )
            
            with export_col2:
                # Markdown export
                markdown_data = export_markdown(current_history, st.session_state.current_session)
                st.download_button(
                    label="📝 Markdown",
                    data=markdown_data,
                    file_name=f"{st.session_state.current_session.replace(' ', '_')}.md",
                    mime="text/markdown",
                    use_container_width=True
                )
            
            with export_col3:
                # Text export
                text_data = export_text(current_history)
                st.download_button(
                    label="📄 Text",
                    data=text_data,
                    file_name=f"{st.session_state.current_session.replace(' ', '_')}.txt",
                    mime="text/plain",
                    use_container_width=True
                )
            
            st.markdown("---")
        
        # --- Chat bubbles above input box, scrollable only inside bubbles ---
        chat_bubble_height = "50vh"
        # Only show chat area if there is at least one valid user prompt or a thinking bubble
        valid_turns = []
        for turn in current_history:
            if turn.get("query"):
                valid_turns.append(turn)
            elif turn.get("thinking") and (not valid_turns or valid_turns[-1].get("thinking") is not True):
                # Only add a thinking bubble if not already present for this user turn
                valid_turns.append(turn)
        if valid_turns:
            st.markdown(f"<div style='max-height:{chat_bubble_height};overflow-y:auto;padding:1rem;background:#181a20;border-radius:0.5rem;'>", unsafe_allow_html=True)
            for chat in valid_turns:
                # Get user query from either the query field or session_history
                user_query = chat.get("query", "")
                if not user_query:
                    user_query = chat.get("session_history", [{}])[-1].get("query") if chat.get("session_history") and isinstance(chat.get("session_history")[-1], dict) else "N/A"
            
                # Get domain info for this query
                domain_id = chat.get("domain_id", "healthcare")
                domain_info = get_domain(domain_id)
                domain_badge = ""
                if domain_info:
                    domain_badge = f"<span style='background:#1976d2;color:white;padding:0.2rem 0.5rem;border-radius:0.8rem;font-size:0.75rem;margin-left:0.5rem;'>{domain_info.icon} {domain_info.name}</span>"
            
                st.markdown(f"<div style='background:#23272f;padding:0.5rem 1rem;border-radius:0.5rem;margin-bottom:0.2rem;color:#fff;max-width:80%;'><b>User:</b> {user_query}{domain_badge}</div>", unsafe_allow_html=True)
            
                # Show only one agent bubble: either thinking or response
                if chat.get("thinking"):
                    st.markdown(f"<div style='background:#e3f2fd;padding:0.6rem 1rem;border-radius:0.5rem;margin-bottom:0.5rem;color:#1976d2;max-width:80%;border-left:4px solid #1976d2;'><b>🤖 Multi-Agent System:</b> <span style='animation:pulse 1.5s infinite;'>Processing...</span><br/><small style='opacity:0.8;'>🧠 Cognitive & Emotional Analysis → 💬 Agent Debate → 🎯 Synthesis</small></div>", unsafe_allow_html=True)
                else:
                    # Show final synthesized answer if available, otherwise fall back to cognitive/emotional
                    agent_summary = ""
                    if 'synthesis' in chat and chat['synthesis'] and chat['synthesis'].get('final_answer'):
                        agent_summary = chat['synthesis']['final_answer']
                    elif 'cognitive' in chat and chat['cognitive'] and chat['cognitive'].get('analysis'):
                        agent_summary = chat['cognitive']['analysis']
                    elif 'emotional' in chat and chat['emotional'] and chat['emotional'].get('analysis'):
                        agent_summary = chat['emotional']['analysis']
                
                    if agent_summary:
                        # Don't truncate - show full synthesized answer
                        # Escape HTML to prevent rendering issues
                        agent_summary_escaped = agent_summary.replace('<', '&lt;').replace('>', '&gt;')
                        st.markdown(f"<div style='background:#fffbe6;padding:0.5rem 1rem;border-radius:0.5rem;margin-bottom:0.5rem;color:#222;max-width:80%;max-height:400px;overflow-y:auto;'><b>Agent:</b> {agent_summary_escaped}</div>", unsafe_allow_html=True)
                    
                        # Show RAG Knowledge Base badges FIRST (most important!)
                        rag_badges = []
                        if 'cognitive' in chat and chat['cognitive'].get('rag_result'):
                            cog_rag = chat['cognitive']['rag_result']
                            if cog_rag.get('source_count', 0) > 0:
                                domain_icon = cog_rag.get('domain_icon', '📚')
                                source_count = cog_rag.get('source_count', 0)
                                rag_badges.append(f"<span style='background:#4caf50;color:white;padding:0.3rem 0.6rem;border-radius:1rem;font-size:0.85rem;'>{domain_icon} Knowledge Base: {source_count} sources</span>")
                    
                        if 'emotional' in chat and chat['emotional'].get('rag_result'):
                            emo_rag = chat['emotional']['rag_result']
                            if emo_rag.get('source_count', 0) > 0 and not rag_badges:  # Only show if cognitive didn't already
                                domain_icon = emo_rag.get('domain_icon', '📚')
                                source_count = emo_rag.get('source_count', 0)
                                rag_badges.append(f"<span style='background:#4caf50;color:white;padding:0.3rem 0.6rem;border-radius:1rem;font-size:0.85rem;'>{domain_icon} Knowledge Base: {source_count} sources</span>")
                    
                        if rag_badges:
                            badges_html = "<div style='display:flex;gap:0.5rem;flex-wrap:wrap;margin-bottom:0.5rem;max-width:80%;'>"
                            badges_html += "".join(rag_badges)
                            badges_html += "</div>"
                            st.markdown(badges_html, unsafe_allow_html=True)
                    
                        # Show tool usage badges if external tools were used
                        if 'cognitive' in chat and chat['cognitive'].get('tool_results'):
                            tool_results = chat['cognitive']['tool_results']
                            badges_html = "<div style='display:flex;gap:0.5rem;flex-wrap:wrap;margin-bottom:0.5rem;max-width:80%;'>"
                        
                            for tool_name, result in tool_results.items():
                                if result.get('success'):
                                    tool_icon = {
                                        'web_search': '🔍',
                                        'weather': '🌤️',
                                        'calculator': '🧮'
                                    }.get(tool_name, '🔧')
                                
                                    tool_display = {
                                        'web_search': 'Web Search',
                                        'weather': 'Weather Data',
                                        'calculator': 'Calculator'
                                    }.get(tool_name, tool_name.title())
                                
                                    badges_html += f"<span style='background:#e3f2fd;color:#1976d2;padding:0.3rem 0.6rem;border-radius:1rem;font-size:0.85em;border:1px solid #1976d2;'>{tool_icon} {tool_display}</span>"
                        
                            badges_html += "</div>"
                            st.markdown(badges_html, unsafe_allow_html=True)
            st.markdown("</div>", unsafe_allow_html=True)

        # ===== BALANCED QUERY QUICK BUTTONS (Organized by Domain!) =====
        st.markdown("---")
        st.markdown("### 💡 Quick Test: Balanced Queries")
        st.caption("Click to auto-fill prompts that showcase P-CABA's balanced reasoning across domains")
    
        # Create 3 columns for domains
        col_health, col_business, col_legal = st.columns(3)
    
        with col_health:
            st.markdown("**🏥 Healthcare**")
            if st.button("Drug Approval Ethics", key="health_1", use_container_width=True):
                st.session_state.input_text = "Should I approve experimental aspirin-based drug for a terminal patient with no other treatment options? Check FDA database for similar drugs and safety data."
                st.session_state.selected_domain = "healthcare"
                st.rerun()
        
            if st.button("AI Triage System", key="health_2", use_container_width=True):
                st.session_state.input_text = "Should we implement AI-assisted triage in our emergency room to prioritize patients?"
                st.session_state.selected_domain = "healthcare"
                st.rerun()
    
        with col_business:
            st.markdown("**💼 Business**")
            if st.button("Layoff vs Salary Cut", key="business_1", use_container_width=True):
                st.session_state.input_text = "Should I lay off 100 employees or cut everyone's salary by 20%? Our SaaS startup has 8 months of runway left. Get current market data for WCLD (SaaS cloud computing ETF) to understand industry conditions."
                st.session_state.selected_domain = "business"
                st.rerun()
        
            if st.button("Acquire Competitor", key="business_2", use_container_width=True):
                st.session_state.input_text = "Should we acquire our main competitor for $50M? Analyze market conditions and strategic implications."
                st.session_state.selected_domain = "business"
                st.rerun()
    
        with col_legal:
            st.markdown("**⚖️ Legal**")
            if st.button("Settlement Decision", key="legal_1", use_container_width=True):
                st.session_state.input_text = "Should we accept the $2M settlement offer or proceed to trial? Case involves product liability."
                st.session_state.selected_domain = "legal"
                st.rerun()
        
            if st.button("Whistleblower Dilemma", key="legal_2", use_container_width=True):
                st.session_state.input_text = "Should I report suspected financial fraud internally or to regulators?"
                st.session_state.selected_domain = "legal"
                st.rerun()
    
        st.markdown("---")

        # ===== INPUT AREA (Clean, no clutter above it) =====
    
        # Initialize domain in session state
        if 'selected_domain' not in st.session_state:
            st.session_state.selected_domain = 'healthcare'
    
        # Initialize input text in session state
        if 'input_text' not in st.session_state:
            st.session_state.input_text = ''

        # Input area - using text_area for better visibility
        # Use on_change to avoid the rerun issue
        def update_input():
            pass  # Just needed for the callback
    
        user_input = st.text_area(
            "Type your message...", 
            value=st.session_state.input_text,
            height=100,
            placeholder="💡 Ask a question or use sidebar prompts...",
            on_change=update_input
        )
    
        # Sync the session state with what user types
        if user_input != st.session_state.input_text:
            st.session_state.input_text = user_input
    
        col1, col2, col3 = st.columns([1, 1, 4])
        with col1:
            submit_button = st.button("📤 Send", type="primary", use_container_width=True)
        with col2:
            if st.button("🗑️ Clear", use_container_width=True):
                st.session_state.input_text = ''
                st.rerun()

        if submit_button and user_input.strip():
            # Clear the input text for next time
            st.session_state.input_text = ''
        
            # Rename session based on first query (if still default name)
            if len(current_history) == 0:  # This will be the first query
                rename_session_from_query(user_input.strip())
        
            # If last entry is a thinking bubble, replace it; else append new
            if current_history and current_history[-1].get("thinking"):
                st.session_state.all_sessions[st.session_state.current_session][-1] = {
                    "query": user_input.strip(),
                    "domain_id": st.session_state.selected_domain,  # Store selected domain
                    "cognitive": {},
                    "thinking": True
                }
            else:
                st.session_state.all_sessions[st.session_state.current_session].append({
                    "query": user_input.strip(),
                    "domain_id": st.session_state.selected_domain,  # Store selected domain
                    "cognitive": {},
                    "thinking": True
                })
            st.rerun()
    
    
    # ===== RIGHT COLUMN: Live Agent Analysis (Always visible) =====
    with col_live_analysis:
        st.markdown("### 🧠 Live Analysis")
        st.markdown("---")
        
        # Get latest completed query
        latest_completed = None
        for turn in reversed(current_history):
            if not turn.get("thinking") and turn.get("query"):
                latest_completed = turn
                break
        
        if latest_completed:
            # Agent Activation Graph - Full Visualization
            st.markdown("#### 🔗 Agent Activation")
            if 'cognitive' in latest_completed and 'emotional' in latest_completed:
                # Build agent timings from latest result (including executor)
                agent_timings = {
                    'controller': latest_completed.get('controller_time', 0.5),
                    'cognitive': latest_completed.get('cognitive', {}).get('time_taken', 0),
                    'emotional': latest_completed.get('emotional', {}).get('time_taken', 0),
                    'synthesizer': latest_completed.get('synthesis', {}).get('time_taken', 0),
                    'executor': latest_completed.get('executor_time', 0),  # Fixed: use executor_time
                    'ethics': 0,
                    'social': 0,
                    'sentiment': 0,
                    'past_experience': 0
                }
                
                # Extract sub-agent times from outputs
                cog_outputs = latest_completed.get('cognitive', {}).get('outputs', [])
                emo_outputs = latest_completed.get('emotional', {}).get('outputs', [])
                
                for output in cog_outputs:
                    if 'Ethics' in output.get('agent', ''):
                        agent_timings['ethics'] = latest_completed.get('cognitive', {}).get('action_time', 2.0) / 2
                    elif 'Social' in output.get('agent', ''):
                        agent_timings['social'] = latest_completed.get('cognitive', {}).get('action_time', 2.0) / 2
                
                for output in emo_outputs:
                    if 'Sentiment' in output.get('agent', ''):
                        agent_timings['sentiment'] = latest_completed.get('emotional', {}).get('time_taken', 0) / 2
                    elif 'Past' in output.get('agent', ''):
                        agent_timings['past_experience'] = latest_completed.get('emotional', {}).get('time_taken', 0) / 2
                
                # Build agent thoughts (increased char limit for better visibility)
                agent_thoughts = {
                    'controller': 'Orchestrates multi-agent workflow',
                    'cognitive': latest_completed.get('cognitive', {}).get('analysis', 'Cognitive reasoning')[:150],
                    'emotional': latest_completed.get('emotional', {}).get('analysis', 'Emotional reasoning')[:150],
                    'synthesizer': latest_completed.get('synthesis', {}).get('final_answer', 'Synthesis')[:150],
                    'executor': latest_completed.get('executor', {}).get('summary', 'Executing actions')[:150],
                    'ethics': cog_outputs[0].get('analysis', 'Ethics analysis')[:120] if len(cog_outputs) > 0 else 'Ethics analysis',
                    'social': cog_outputs[1].get('analysis', 'Social analysis')[:120] if len(cog_outputs) > 1 else 'Social impact',
                    'sentiment': emo_outputs[0].get('analysis', 'Sentiment analysis')[:120] if len(emo_outputs) > 0 else 'Sentiment',
                    'past_experience': emo_outputs[1].get('analysis', 'Past experience')[:120] if len(emo_outputs) > 1 else 'Experience'
                }
                
                # Get debate time
                debate_time = latest_completed.get('debate_time', 0)
                
                # Create Agent Activation Graph
                try:
                    from agent_activation_viz import visualize_agent_activation
                    fig = visualize_agent_activation(
                        agent_timings, 
                        agent_thoughts, 
                        debate_time=debate_time,
                        show_architecture_labels=False  # Compact for right pane
                    )
                    st.plotly_chart(fig, use_container_width=True)
                except Exception as e:
                    st.caption("⏳ Graph loading...")
            else:
                st.info("⏳ Processing...")
            
            st.markdown("---")
            
            # Performance Metrics
            st.markdown("#### 📊 Performance")
            if 'cognitive' in latest_completed and 'emotional' in latest_completed:
                cog_score = latest_completed['cognitive'].get('confidence', 0)
                emo_score = latest_completed['emotional'].get('confidence', 0)
                cog_time = latest_completed['cognitive'].get('time_taken', 0)
                emo_time = latest_completed['emotional'].get('time_taken', 0)
                
                # Show confidence scores
                col_a, col_b = st.columns(2)
                with col_a:
                    st.metric("🧠 Cognitive", f"{cog_score:.0%}")
                    st.caption(f"⏱️ {cog_time:.2f}s")
                with col_b:
                    st.metric("❤️ Emotional", f"{emo_score:.0%}")
                    st.caption(f"⏱️ {emo_time:.2f}s")
                
                # Total time
                total_time = cog_time + emo_time
                if 'synthesis' in latest_completed:
                    synth_time = latest_completed['synthesis'].get('time_taken', 0)
                    total_time += synth_time
                
                st.caption(f"🎯 Total: {total_time:.2f}s")
                
                # Show debate if available
                debate_log = latest_completed.get('debate_log', '')
                debate_time_val = latest_completed.get('debate_time', 0)
                
                with st.expander("💬 Debate", expanded=False):
                    if debate_log:
                        st.caption(f"⏱️ Duration: {debate_time_val:.3f}s")
                        st.markdown("---")
                        
                        # Parse and display debate exchanges with full formatting
                        lines = debate_log.split('\n')
                        formatted_debate = []
                        
                        for line in lines:
                            if 'Turn' in line and ('Cognitive:' in line or 'Emotional:' in line):
                                if 'Cognitive:' in line:
                                    # Cognitive debate turns - blue background
                                    formatted_debate.append(f"<div style='background:#1976d2;padding:0.5rem;border-left:3px solid #0d47a1;margin:0.2rem 0;border-radius:0.2rem;color:#ffffff;font-weight:500;font-size:0.85em;'>{line}</div>")
                                else:
                                    # Emotional debate turns - pink background
                                    formatted_debate.append(f"<div style='background:#c2185b;padding:0.5rem;border-left:3px solid #880e4f;margin:0.2rem 0;border-radius:0.2rem;color:#ffffff;font-weight:500;font-size:0.85em;'>{line}</div>")
                            elif 'COMMON GROUND' in line or '🤝' in line:
                                # Common ground - green background
                                formatted_debate.append(f"<div style='background:#388e3c;padding:0.5rem;border-left:3px solid #1b5e20;margin:0.2rem 0;border-radius:0.2rem;color:#ffffff;font-weight:600;font-size:0.85em;'>{line}</div>")
                            elif 'RESOLUTION' in line or '⚖️' in line:
                                # Resolution - gold background
                                formatted_debate.append(f"<div style='background:#ffc107;padding:0.5rem;border-left:3px solid #f57f17;margin:0.2rem 0;border-radius:0.2rem;color:#1a1a1a;font-weight:600;font-size:0.85em;'>{line}</div>")
                            elif 'Confidence:' in line:
                                formatted_debate.append(f"<div style='padding:0.2rem 0.3rem;color:#bbbbbb;font-size:0.75em;'>{line}</div>")
                            elif line.strip() and not line.startswith('==='):
                                formatted_debate.append(f"<div style='padding:0.2rem 0.3rem;color:#e0e0e0;font-size:0.8em;'>{line}</div>")
                        
                        st.markdown('\n'.join(formatted_debate), unsafe_allow_html=True)
                    else:
                        st.info("No debate available for this query.")
            else:
                st.info("⏳ Processing...")
            
            # Adaptive Weights
            st.markdown("---")
            st.markdown("#### ⚖️ Adaptive Weights")
            
            # Check multiple possible locations for weights
            weights = None
            # First check top-level weights from controller
            if 'weights' in latest_completed:
                weights = latest_completed['weights']
            # Then check synthesis for backward compatibility
            elif 'synthesis' in latest_completed:
                if 'adaptive_weights' in latest_completed['synthesis']:
                    weights = latest_completed['synthesis']['adaptive_weights']
                elif 'final_weights' in latest_completed['synthesis']:
                    weights = latest_completed['synthesis']['final_weights']
            
            if weights:
                cog_weight = weights.get('cognitive', 0.5)
                emo_weight = weights.get('emotional', 0.5)
                
                st.progress(cog_weight, text=f"🧠 Cognitive: {cog_weight:.2%}")
                st.progress(emo_weight, text=f"❤️ Emotional: {emo_weight:.2%}")
                
                # Show weight interpretation
                if cog_weight > 0.6:
                    st.caption("📊 Logic-heavy query")
                elif emo_weight > 0.6:
                    st.caption("💝 Empathy-focused query")
                else:
                    st.caption("⚖️ Balanced reasoning")
            else:
                # Show default weights if synthesis hasn't completed
                st.progress(0.5, text="🧠 Cognitive: 50%")
                st.progress(0.5, text="❤️ Emotional: 50%")
                st.caption("⏳ Using balanced defaults")
            
            # Knowledge Base (RAG)
            st.markdown("---")
            st.markdown("#### 📚 Knowledge Base")
            if 'cognitive' in latest_completed and 'rag_result' in latest_completed['cognitive']:
                rag_result = latest_completed['cognitive']['rag_result']
                source_count = rag_result.get('source_count', 0)
                if source_count > 0:
                    domain_icon = rag_result.get('domain_icon', '📚')
                    domain_name = rag_result.get('domain_name', 'Domain')
                    st.success(f"{domain_icon} **{source_count} sources** retrieved")
                    st.caption(f"Domain: {domain_name}")
                    
                    # Show top sources with scores
                    retrieved_docs = rag_result.get('retrieved_docs', [])
                    if retrieved_docs:
                        with st.expander("📖 View Sources", expanded=False):
                            for idx, doc in enumerate(retrieved_docs[:3], 1):
                                source = doc.get('source', 'Unknown')
                                score = doc.get('score', 0.0)
                                st.markdown(f"**{idx}.** {source}")
                                st.progress(score, text=f"Relevance: {score:.0%}")
                                st.markdown("---")
                else:
                    st.caption("ℹ️ No RAG retrieval")
            else:
                st.caption("ℹ️ No RAG retrieval")
            
            # Executor Actions - Full Details
            st.markdown("---")
            st.markdown("#### ⚡ Executor Actions")
            if 'executor' in latest_completed:
                executor_data = latest_completed['executor']
                
                # Check for actions list
                actions = executor_data.get('actions', [])
                
                if actions and len(actions) > 0:
                    # Show action count
                    total_actions = len(actions)
                    successful = sum(1 for a in actions if a.get('status') == 'success')
                    st.caption(f"📊 {successful}/{total_actions} actions completed")
                    
                    # Show each action with full details
                    for idx, action in enumerate(actions[:5], 1):  # Show top 5
                        if isinstance(action, dict):
                            tool_name = action.get('tool', action.get('name', 'Unknown Action'))
                            status = action.get('status', 'success')
                            exec_time = action.get('execution_time', 0)
                            action_icon = action.get('icon', '⚙️')
                            action_type = action.get('type', 'python')
                            
                            # Status icon
                            status_icon = "✅" if status == 'success' else "❌"
                            
                            # Color by type
                            if action_type == 'lambda':
                                bg_color = "#1a237e"
                                border_color = "#3f51b5"
                            else:
                                bg_color = "#1b5e20"
                                border_color = "#4caf50"
                            
                            # Show action card
                            st.markdown(f"""
                            <div style='background:{bg_color};padding:0.6rem;
                            border-left:3px solid {border_color};margin:0.4rem 0;
                            border-radius:0.3rem;color:#ffffff;font-size:0.85em;'>
                                <div style='font-weight:600;'>{status_icon} {action_icon} {tool_name}</div>
                                <div style='opacity:0.8;margin-top:0.2rem;'>
                                    {action_type.upper()} • {exec_time:.2f}s
                                </div>
                            </div>
                            """, unsafe_allow_html=True)
                            
                            # Show details in expander
                            details = action.get('details', {})
                            if details:
                                with st.expander(f"View {tool_name} Details", expanded=False):
                                    st.json(details)
                        else:
                            st.markdown(f"✅ **{idx}.** {action}")
                else:
                    st.caption("ℹ️ No actions executed")
            else:
                st.caption("ℹ️ No executor data")
        else:
            st.info("💡 Submit a query to see live agent analysis here")
    
    # End of Tab 1 (Chat) - Clean conversation-only interface

# ===== TAB 2: AGENT ANALYSIS (Full-width, detailed breakdown) =====
with tab2:
    st.markdown("### 🔍 Agent Analysis & Multi-Agent Reasoning")
    
    # Get current history
    current_history = st.session_state.all_sessions[st.session_state.current_session]
    
    # Show details for the most recent completed turn
    if current_history and not current_history[-1].get("thinking"):
        latest = current_history[-1]
        
        # Performance Metrics - Clean Summary
        if 'cognitive' in latest and 'emotional' in latest:
            st.markdown("#### ⚡ Performance Metrics")
            
            # Create metrics columns
            col1, col2, col3, col4 = st.columns(4)
            
            with col1:
                cog_time = latest.get('cognitive', {}).get('time_taken', 0)
                cog_method = latest.get('cognitive', {}).get('method', 'unknown')
                method_label = "🌩️ AgentCore" if cog_method == 'agentcore' else "🧠 Local"
                st.metric("Cognitive Agent", f"{cog_time:.2f}s", method_label)
            
            with col2:
                emo_time = latest.get('emotional', {}).get('time_taken', 0)
                st.metric("Emotional Agent", f"{emo_time:.2f}s", "📊 LangGraph")
            
            with col3:
                synth_time = latest.get('synthesis', {}).get('time_taken', 0)
                st.metric("Synthesizer", f"{synth_time:.2f}s", "🔄 Strands SDK")
            
            with col4:
                total_time = cog_time + emo_time + synth_time
                st.metric("Total Time", f"{total_time:.2f}s", "⚡ Complete")
            
            # Show RAG retrieval info prominently
            if latest.get('cognitive', {}).get('rag_result', {}).get('source_count', 0) > 0:
                cog_rag = latest['cognitive']['rag_result']
                st.markdown("---")
                st.markdown("#### 📚 Knowledge Base Retrieval")
                
                # Show summary
                domain_icon = cog_rag.get('domain_icon', '📚')
                domain_name = cog_rag.get('domain_name', 'Domain')
                source_count = cog_rag.get('source_count', 0)
                
                st.success(f"{domain_icon} Retrieved **{source_count} expert sources** from {domain_name} domain")
                
                # Show sources in a nice table
                with st.expander("📖 View Retrieved Sources", expanded=False):
                    retrieved_docs = cog_rag.get('retrieved_docs', [])
                    for idx, doc in enumerate(retrieved_docs, 1):
                        source = doc.get('source', 'Unknown')
                        score = doc.get('score', 0.0)
                        col_a, col_b = st.columns([3, 1])
                        with col_a:
                            st.markdown(f"**{idx}. {source}**")
                        with col_b:
                            st.metric("Relevance", f"{score:.0%}")
            
            # Show AWS Architecture
            st.markdown("---")
            st.markdown("#### 🏗️ AWS Multi-Architecture System")
            arch_cols = st.columns(3)
            with arch_cols[0]:
                st.markdown("**🌩️ Bedrock AgentCore**")
                st.caption("Cognitive reasoning with managed agents")
            with arch_cols[1]:
                st.markdown("**📊 LangGraph Workflow**")
                st.caption("Emotional analysis with sub-agents")
            with arch_cols[2]:
                st.markdown("**🔄 Strands SDK**")
                st.caption("Multi-turn synthesis & refinement")
            
            # Add Agent Activation Graph in main chat view
            st.markdown("---")
            st.markdown("#### 🔗 Agent Activation Graph")
            
            # Build agent timings from latest result
            agent_timings = {
                'controller': 0.5,
                'cognitive': latest.get('cognitive', {}).get('time_taken', 0),
                'emotional': latest.get('emotional', {}).get('time_taken', 0),
                'synthesizer': latest.get('synthesis', {}).get('time_taken', 0),
                'ethics': 0,
                'social': 0,
                'sentiment': 0,
                'past_experience': 0
            }
            
            # Extract sub-agent times from outputs
            cog_outputs = latest.get('cognitive', {}).get('outputs', [])
            emo_outputs = latest.get('emotional', {}).get('outputs', [])
            
            for output in cog_outputs:
                if 'Ethics' in output.get('agent', ''):
                    agent_timings['ethics'] = latest.get('cognitive', {}).get('action_time', 2.0) / 2
                elif 'Social' in output.get('agent', ''):
                    agent_timings['social'] = latest.get('cognitive', {}).get('action_time', 2.0) / 2
            
            for output in emo_outputs:
                if 'Sentiment' in output.get('agent', ''):
                    agent_timings['sentiment'] = latest.get('emotional', {}).get('time_taken', 0) / 2
                elif 'Past' in output.get('agent', ''):
                    agent_timings['past_experience'] = latest.get('emotional', {}).get('time_taken', 0) / 2
            
            # Build agent thoughts
            agent_thoughts = {
                'controller': 'Orchestrates workflow',
                'cognitive': latest.get('cognitive', {}).get('analysis', 'Cognitive reasoning')[:100],
                'emotional': latest.get('emotional', {}).get('analysis', 'Emotional reasoning')[:100],
                'synthesizer': latest.get('synthesis', {}).get('final_answer', 'Synthesis')[:100],
                'ethics': cog_outputs[0].get('analysis', 'Ethics analysis')[:100] if len(cog_outputs) > 0 else 'Ethics analysis',
                'social': cog_outputs[1].get('analysis', 'Social analysis')[:100] if len(cog_outputs) > 1 else 'Social impact',
                'sentiment': emo_outputs[0].get('analysis', 'Sentiment analysis')[:100] if len(emo_outputs) > 0 else 'Sentiment',
                'past_experience': emo_outputs[1].get('analysis', 'Past experience')[:100] if len(emo_outputs) > 1 else 'Past experience'
            }
            
            # Get debate time
            debate_time = latest.get('debate_time', 0)
            
            # Create compact graph for main view (no architecture labels to save space)
            from agent_activation_viz import visualize_agent_activation
            fig = visualize_agent_activation(
                agent_timings, 
                agent_thoughts, 
                debate_time=debate_time,
                show_architecture_labels=False  # Compact view for right pane
            )
            st.plotly_chart(fig, use_container_width=True)
            
            st.caption("💡 See **Architecture Deep Dive** tab for detailed breakdown")
        
        # Final Synthesized Answer (highlight this prominently)
        if 'synthesis' in latest and latest['synthesis']:
            st.markdown("#### 🎯 Final Synthesized Answer")
            with st.expander("View Full Answer", expanded=True):
                st.success(latest['synthesis'].get('final_answer', 'N/A'))
                st.caption(f"Confidence: {latest['synthesis'].get('confidence', 0):.2%} | Processing Time: {latest['synthesis'].get('time_taken', 0):.3f}s")
        
        # RAG Knowledge Base Sources (NEW: show retrieved knowledge)
        rag_helper = get_rag_helper()
        has_rag_sources = False
        
        # Check cognitive agent RAG results
        if 'cognitive' in latest and latest['cognitive'].get('rag_result'):
            cog_rag = latest['cognitive']['rag_result']
            if cog_rag.get('source_count', 0) > 0:
                has_rag_sources = True
                st.markdown("#### 📚 Knowledge Base Sources (Cognitive)")
                with st.expander("View Retrieved Knowledge", expanded=True):
                    ui_text = rag_helper.format_sources_for_ui(cog_rag)
                    st.markdown(ui_text)
        
        # Check emotional agent RAG results
        if 'emotional' in latest and latest['emotional'].get('rag_result'):
            emo_rag = latest['emotional']['rag_result']
            if emo_rag.get('source_count', 0) > 0:
                if not has_rag_sources:
                    st.markdown("#### 📚 Knowledge Base Sources")
                with st.expander("View Emotional Context Knowledge", expanded=True):
                    ui_text = rag_helper.format_sources_for_ui(emo_rag)
                    st.markdown(ui_text)
        
        # External Tools Usage (NEW: highlight tool integration)
        if 'cognitive' in latest and latest['cognitive'].get('tool_results'):
            st.markdown("#### 🔧 External Tools Used")
            with st.expander("View Tool Results", expanded=True):
                tool_results = latest['cognitive']['tool_results']
                
                for tool_name, result in tool_results.items():
                    if not result.get('success', False):
                        continue
                    
                    # Display based on tool type
                    if tool_name == 'web_search':
                        st.markdown("##### 🔍 Web Search")
                        query = result.get('query', '')
                        duration = result.get('duration', 0)
                        results_list = result.get('results', [])[:3]
                        
                        st.info(f"**Query:** {query}\n**Duration:** {duration:.2f}s\n**Results found:** {result.get('total_results', 0)}")
                        
                        if results_list:
                            for idx, res in enumerate(results_list, 1):
                                # Safely get fields with defaults
                                title = res.get('title', 'No title')
                                url = res.get('url', '#')
                                snippet = res.get('snippet', 'No description available')
                                
                                st.markdown(f"**{idx}. [{title}]({url})**")
                                st.caption(snippet)
                                st.markdown("---")
                        else:
                            st.warning("No web search results found for this query.")
                    
                    elif tool_name == 'weather':
                        st.markdown("##### 🌤️ Weather Information")
                        weather_data = result.get('weather', {})
                        location = weather_data.get('location', 'Unknown')
                        temp = weather_data.get('temperature', 'N/A')
                        windspeed = weather_data.get('windspeed', 'N/A')
                        desc = weather_data.get('description', 'N/A')
                        coords = weather_data.get('coordinates', {})
                        
                        col1, col2, col3 = st.columns(3)
                        with col1:
                            st.metric("Temperature", f"{temp}°C")
                        with col2:
                            st.metric("Wind Speed", f"{windspeed} km/h")
                        with col3:
                            st.metric("Condition", desc)
                        
                        st.caption(f"📍 Location: {location} ({coords.get('latitude', 'N/A')}, {coords.get('longitude', 'N/A')})")
                    
                    elif tool_name == 'calculator':
                        st.markdown("##### 🧮 Calculator")
                        expr = result.get('expression', '')
                        res_value = result.get('result', '')
                        duration = result.get('duration', 0)
                        
                        st.code(f"{expr} = {res_value}", language="python")
                        st.caption(f"⏱️ Calculated in {duration:.4f}s")
        
        st.markdown("---")
        
        # Individual agent outputs
        st.markdown("#### 🧠 Cognitive Agent")
        with st.expander("View Details", expanded=False):
            if 'cognitive' in latest and latest['cognitive']:
                st.write(latest['cognitive'].get('analysis', 'N/A'))
                st.caption(f"Confidence: {latest['cognitive'].get('confidence', 0):.2%}")
                # Show sub-agent breakdown
                if latest['cognitive'].get('outputs'):
                    outputs = latest['cognitive']['outputs']
                    if isinstance(outputs, list):
                        st.markdown("**Sub-Agents:**")
                        for sub_data in outputs:
                            if isinstance(sub_data, dict):
                                agent_name = sub_data.get('agent', 'Unknown')
                                analysis = sub_data.get('analysis', 'N/A')
                                st.markdown(f"- **{agent_name}**: {analysis[:100]}...")
        
        st.markdown("#### ❤️ Emotional Agent")
        with st.expander("View Details", expanded=False):
            if 'emotional' in latest and latest['emotional']:
                st.write(latest['emotional'].get('analysis', 'N/A'))
                st.caption(f"Confidence: {latest['emotional'].get('confidence', 0):.2%}")
                # Show sub-agent breakdown
                if latest['emotional'].get('outputs'):
                    outputs = latest['emotional']['outputs']
                    if isinstance(outputs, list):
                        st.markdown("**Sub-Agents:**")
                        for sub_data in outputs:
                            if isinstance(sub_data, dict):
                                agent_name = sub_data.get('agent', 'Unknown')
                                analysis = sub_data.get('analysis', 'N/A')
                                st.markdown(f"- **{agent_name}**: {analysis[:100]}...")
        
        # Debate log with enhanced visualization
        st.markdown("#### 💬 Agent Debate")
        with st.expander("View Debate", expanded=False):
            if 'debate_log' in latest and latest['debate_log']:
                debate_text = latest['debate_log']
                
                # Parse and color-code debate exchanges
                lines = debate_text.split('\n')
                formatted_debate = []
                
                for line in lines:
                    # Check for actual debate turn statements (these get colored highlights with FULL text)
                    if 'Turn' in line and ('Cognitive:' in line or 'Emotional:' in line):
                        if 'Cognitive:' in line:
                            # Cognitive debate turns - blue background with WHITE text - SHOW FULL TEXT
                            formatted_debate.append(f"<div style='background:#1976d2;padding:0.6rem;border-left:4px solid #0d47a1;margin:0.3rem 0;border-radius:0.3rem;color:#ffffff;font-weight:500;word-wrap:break-word;'>{line}</div>")
                        else:
                            # Emotional debate turns - pink background with WHITE text - SHOW FULL TEXT
                            formatted_debate.append(f"<div style='background:#c2185b;padding:0.6rem;border-left:4px solid #880e4f;margin:0.3rem 0;border-radius:0.3rem;color:#ffffff;font-weight:500;word-wrap:break-word;'>{line}</div>")
                    elif 'Query:' in line:
                        # Query display - subtle gray with white text
                        formatted_debate.append(f"<div style='background:#424242;padding:0.5rem;border-left:4px solid #212121;margin:0.2rem 0;border-radius:0.3rem;color:#ffffff;font-style:italic;'>{line}</div>")
                    elif 'Confidence:' in line and '|' in line:
                        # Confidence metrics - subtle display
                        formatted_debate.append(f"<div style='padding:0.3rem 0.5rem;color:#bbbbbb;font-size:0.9em;'>{line}</div>")
                    elif 'COMMON GROUND' in line or '🤝' in line:
                        # Agreement sections - green with white text for visibility
                        formatted_debate.append(f"<div style='background:#388e3c;padding:0.5rem;border-left:4px solid #1b5e20;margin:0.3rem 0;border-radius:0.3rem;color:#ffffff;font-weight:600;'>{line}</div>")
                    elif 'RESOLUTION' in line or '⚖️' in line:
                        # Resolution sections - gold with dark text
                        formatted_debate.append(f"<div style='background:#ffc107;padding:0.5rem;border-left:4px solid #f57f17;margin:0.3rem 0;border-radius:0.3rem;color:#1a1a1a;font-weight:600;'>{line}</div>")
                    elif 'DEBATE EXCHANGE' in line or '💬' in line or 'Multi-Agent Debate' in line:
                        # Section headers - gray with white text
                        formatted_debate.append(f"<div style='background:#616161;padding:0.5rem;border-left:4px solid #424242;margin:0.3rem 0;border-radius:0.3rem;color:#ffffff;font-weight:600;'>{line}</div>")
                    elif line.strip() and not line.startswith('==='):
                        # Regular text lines - white text for readability on dark background
                        formatted_debate.append(f"<div style='padding:0.3rem 0.5rem;color:#e0e0e0;line-height:1.5;'>{line}</div>")
                
                st.markdown('\n'.join(formatted_debate), unsafe_allow_html=True)
                
                # Add debate statistics
                st.markdown("---")
                col1, col2, col3 = st.columns(3)
                with col1:
                    cog_conf = latest.get('cognitive', {}).get('confidence', 0)
                    st.metric("Cognitive Strength", f"{cog_conf:.0%}")
                with col2:
                    emo_conf = latest.get('emotional', {}).get('confidence', 0)
                    st.metric("Emotional Strength", f"{emo_conf:.0%}")
                with col3:
                    debate_duration = latest.get('debate_time', 0)
                    st.metric("Debate Duration", f"{debate_duration*1000:.0f}ms")
            else:
                st.info("No debate log available")
        
        # Executor Agent Actions (NEW: Real action execution)
        if 'executor' in latest and latest['executor']:
            st.markdown("#### 🤖 Executor Agent: Real Actions")
            executor_data = latest['executor']
            
            # Summary metrics
            total_actions = executor_data.get('successful', 0) + executor_data.get('failed', 0)
            if total_actions > 0:
                col1, col2, col3 = st.columns(3)
                with col1:
                    st.metric("Actions Executed", total_actions, f"✅ {executor_data.get('successful', 0)} successful")
                with col2:
                    st.metric("Execution Time", f"{executor_data.get('total_time', 0):.2f}s")
                with col3:
                    success_rate = executor_data.get('successful', 0) / total_actions * 100 if total_actions > 0 else 0
                    st.metric("Success Rate", f"{success_rate:.0f}%")
                
                # Individual action details
                with st.expander("📋 View Action Details", expanded=True):
                    actions = executor_data.get('actions', [])
                    
                    for idx, action in enumerate(actions, 1):
                        status_icon = "✅" if action['status'] == 'success' else "❌"
                        action_type = action.get('type', 'python')
                        
                        # Color-code by type
                        if action_type == 'lambda':
                            bg_color = "#1a237e"  # Deep blue for Lambda
                            border_color = "#3f51b5"
                        else:
                            bg_color = "#1b5e20"  # Green for Python
                            border_color = "#4caf50"
                        
                        # Action header
                        st.markdown(f"""
                        <div style='background:{bg_color};padding:0.8rem;border-left:4px solid {border_color};
                        margin:0.5rem 0;border-radius:0.4rem;color:#ffffff;'>
                            <div style='font-size:1.1em;font-weight:600;'>{status_icon} {action['icon']} {action['name']}</div>
                            <div style='font-size:0.9em;margin-top:0.3rem;opacity:0.9;'>{action['description']}</div>
                            <div style='font-size:0.85em;margin-top:0.3rem;opacity:0.7;'>
                                Type: {action_type.upper()} | Duration: {action['execution_time']:.2f}s
                            </div>
                        </div>
                        """, unsafe_allow_html=True)
                        
                        # Lambda-specific details
                        if action_type == 'lambda' and action.get('lambda_arn'):
                            with st.container():
                                st.caption(f"🔗 **Lambda ARN:** `{action['lambda_arn']}`")
                                st.caption(f"📋 **AWS Request ID:** `{action['aws_request_id']}`")
                        
                        # Action-specific details
                        details = action.get('details', {})
                        if details:
                            # Format details nicely
                            details_display = []
                            
                            # Special formatting for common fields
                            for key, value in details.items():
                                if key == 'error':
                                    st.error(f"Error: {value}")
                                elif key in ['note', 'message', 'disclaimer']:
                                    st.info(f"ℹ️ {value}")
                                elif key == 'api':
                                    st.success(f"🌐 Data Source: {value}")
                                elif key == 'execution_type':
                                    st.caption(f"⚙️ Execution: {value}")
                                elif isinstance(value, (list, dict)):
                                    # Show complex data in expander
                                    with st.expander(f"📊 {key.replace('_', ' ').title()}", expanded=False):
                                        st.json(value)
                                else:
                                    details_display.append(f"**{key.replace('_', ' ').title()}:** {value}")
                            
                            # Display simple key-value pairs
                            if details_display:
                                for detail in details_display:
                                    st.markdown(f"- {detail}")
                        
                        st.markdown("---")
            else:
                st.info(f"💡 {executor_data.get('message', 'No actions triggered for this decision')}")
        
        # Adaptive weights - ENHANCED: Make this a prominent feature highlight
        st.markdown("### ⚖️ Adaptive Weights (P-CABA Unique Feature)")
        
        # Add explanation of why this matters
        st.info("""
        **🎯 What makes P-CABA different:** The system dynamically adjusts how much weight 
        to give cognitive reasoning vs emotional intelligence based on the situation. 
        High-stakes ethical dilemmas get more emotional weight, while technical decisions 
        get more cognitive weight. This adapts in real-time for each query!
        """)
        
        if 'weights' in latest and latest['weights']:
            weights = latest['weights']
            
            # Visual prominence with color-coded columns
            col1, col2 = st.columns(2)
            
            with col1:
                cog_weight = weights.get('cognitive', 0.5)
                st.metric("🧠 Cognitive Weight", f"{cog_weight:.0%}", 
                         help="How much the system prioritizes logical analysis")
                st.progress(cog_weight, text=f"Cognitive: {cog_weight:.2f}")
                
                # Explain what this weight means
                if cog_weight > 0.7:
                    st.caption("✅ High cognitive focus - technical/analytical decision")
                elif cog_weight < 0.3:
                    st.caption("⚠️ Low cognitive focus - emotion-driven situation")
                else:
                    st.caption("⚖️ Balanced cognitive-emotional approach")
            
            with col2:
                emo_weight = weights.get('emotional', 0.5)
                st.metric("❤️ Emotional Weight", f"{emo_weight:.0%}",
                         help="How much the system prioritizes emotional considerations")
                st.progress(emo_weight, text=f"Emotional: {emo_weight:.2f}")
                
                # Explain what this weight means
                if emo_weight > 0.7:
                    st.caption("❤️ High emotional focus - human impact priority")
                elif emo_weight < 0.3:
                    st.caption("🤖 Low emotional focus - technical problem")
                else:
                    st.caption("⚖️ Balanced emotional-cognitive approach")
            
            # Show weight adaptation explanation
            st.success(f"""
            **Adaptive Learning:** P-CABA adjusted these weights based on:
            - Query complexity and emotional content
            - Domain context (healthcare tends toward higher emotional weight)
            - Presence of ethical dilemmas (family vs patient wishes, etc.)
            - Past experience with similar decisions
            
            This real-time adaptation is what makes P-CABA unique compared to static AI systems!
            """)
        else:
            st.warning("Adaptive weights not available for this query")
        
        # Memory context - FIXED: Deduplicate entries
        st.markdown("#### 🧠 Memory Context")
        with st.expander("View Memory", expanded=False):
            if 'memory_context' in latest and latest['memory_context']:
                # Remove duplicates while preserving order
                unique_memories = []
                seen = set()
                for mem in latest['memory_context']:
                    if mem not in seen:
                        unique_memories.append(mem)
                        seen.add(mem)
                        if len(unique_memories) >= 3:
                            break
                
                if unique_memories:
                    for idx, mem in enumerate(unique_memories, 1):
                        st.caption(f"{idx}. {mem}")
                else:
                    st.info("No unique memory context available")
            else:
                st.info("No memory context available")
    else:
        st.info("Process a query to see agent analysis details")


# ===== TAB 3: ADAPTIVE WEIGHTS (Unique P-CABA differentiator) =====
with tab3:
    st.markdown("### ⚖️ Adaptive Weight System")
    st.markdown("""
    **P-CABA's unique feature:** Dynamically adjusts the influence of cognitive vs. emotional reasoning 
    based on the context and domain of each query.
    """)
    
    # Get current history
    current_history = st.session_state.all_sessions[st.session_state.current_session]
    
    if current_history and not current_history[-1].get("thinking"):
        latest = current_history[-1]
        
        # Adaptive Weights Visualization
        # Check for weights - prioritize top-level 'weights' from controller
        weights = None
        if 'weights' in latest:
            weights = latest['weights']
        elif 'synthesis' in latest and 'adaptive_weights' in latest['synthesis']:
            weights = latest['synthesis']['adaptive_weights']
        elif 'adaptive_weights' in latest:
            weights = latest['adaptive_weights']
        
        if weights:
            # Show current weights prominently with custom colors
            col1, col2 = st.columns(2)
            with col1:
                cog_weight = weights.get('cognitive', 0.5)
                st.metric("🧠 Cognitive Weight", f"{cog_weight:.1%}", 
                         help="Influence of logical, analytical reasoning")
                # Orange progress bar for cognitive (brain/analytical)
                st.markdown(f"""
                <div style="background: #2a2a2a; border-radius: 10px; padding: 2px; margin: 10px 0;">
                    <div style="background: linear-gradient(90deg, #ff6b35 0%, #f7931e {cog_weight*100}%, #2a2a2a {cog_weight*100}%); 
                                height: 20px; border-radius: 8px; transition: all 0.3s ease;"></div>
                </div>
                """, unsafe_allow_html=True)
            
            with col2:
                emo_weight = weights.get('emotional', 0.5)
                st.metric("❤️ Emotional Weight", f"{emo_weight:.1%}",
                         help="Influence of emotional, empathetic reasoning")
                # Green progress bar for emotional (heart/empathy)
                st.markdown(f"""
                <div style="background: #2a2a2a; border-radius: 10px; padding: 2px; margin: 10px 0;">
                    <div style="background: linear-gradient(90deg, #4caf50 0%, #66bb6a {emo_weight*100}%, #2a2a2a {emo_weight*100}%); 
                                height: 20px; border-radius: 8px; transition: all 0.3s ease;"></div>
                </div>
                """, unsafe_allow_html=True)
            
            # Explanation of weights
            st.markdown("#### 💡 Why These Weights?")
            if 'explanation' in weights:
                st.info(weights['explanation'])
            else:
                if cog_weight > emo_weight:
                    st.info("This query requires more **analytical** reasoning (data, logic, facts)")
                elif emo_weight > cog_weight:
                    st.info("This query requires more **empathetic** reasoning (emotions, ethics, social impact)")
                else:
                    st.info("This query requires **balanced** reasoning (equal parts logic and empathy)")
            
            # Show historical weights if multiple queries exist
            st.markdown("---")
            st.markdown("#### 📈 Weight History")
            
            # Collect weights from all completed queries
            weights_history = []
            for turn in current_history:
                if not turn.get('thinking') and turn.get('query'):
                    turn_weights = None
                    # Prioritize top-level 'weights' from controller
                    if 'weights' in turn:
                        turn_weights = turn['weights']
                    elif 'synthesis' in turn and 'adaptive_weights' in turn['synthesis']:
                        turn_weights = turn['synthesis']['adaptive_weights']
                    elif 'adaptive_weights' in turn:
                        turn_weights = turn['adaptive_weights']
                    
                    if turn_weights:
                        weights_history.append({
                            'query': turn.get('query', 'Unknown')[:50] + '...',
                            'cognitive': turn_weights.get('cognitive', 0.5),
                            'emotional': turn_weights.get('emotional', 0.5)
                        })
            
            if len(weights_history) > 0:
                import pandas as pd
                import plotly.express as px
                
                df = pd.DataFrame(weights_history)
                fig = px.line(df, y=['cognitive', 'emotional'], 
                             title='Weight Adaptation Over Time',
                             labels={'value': 'Weight', 'variable': 'Agent Type'},
                             color_discrete_map={'cognitive': '#ff6b35', 'emotional': '#4caf50'})
                fig.update_layout(
                    height=300,
                    paper_bgcolor='#1e1e1e',
                    plot_bgcolor='#1e1e1e',
                    font=dict(color='#ffffff')
                )
                fig.update_xaxes(showgrid=True, gridcolor='#333333')
                fig.update_yaxes(showgrid=True, gridcolor='#333333')
                st.plotly_chart(fig, use_container_width=True)
            else:
                st.info("Weight history will appear as you process more queries")
        else:
            st.warning("⚠️ Adaptive weights not available for this query. The synthesis agent may not have completed processing.")
    else:
        st.info("Process a query to see adaptive weight adjustments")


# ===== TAB 4: EXECUTOR ACTIONS (Real-world integration showcase) =====
with tab4:
    st.markdown("### 🤖 Real Executor Actions")
    st.markdown("""
    **Real-world integration:** P-CABA doesn't just think - it **acts**! 
    FDA API calls, market data retrieval, email notifications, and more.
    """)
    
    # Get current history
    current_history = st.session_state.all_sessions[st.session_state.current_session]
    
    if current_history and not current_history[-1].get("thinking"):
        latest = current_history[-1]
        
        # Executor Agent Actions
        if 'executor' in latest and latest['executor']:
            executor_data = latest['executor']
            
            # Summary metrics
            total_actions = executor_data.get('successful', 0) + executor_data.get('failed', 0)
            if total_actions > 0:
                col1, col2, col3 = st.columns(3)
                with col1:
                    st.metric("Actions Executed", total_actions, f"✅ {executor_data.get('successful', 0)} successful")
                with col2:
                    st.metric("Execution Time", f"{executor_data.get('total_time', 0):.2f}s")
                with col3:
                    success_rate = executor_data.get('successful', 0) / total_actions * 100 if total_actions > 0 else 0
                    st.metric("Success Rate", f"{success_rate:.0f}%")
                
                # Individual action details
                st.markdown("#### 📋 Action Details")
                actions = executor_data.get('actions', [])
                
                for idx, action in enumerate(actions, 1):
                    status_icon = "✅" if action['status'] == 'success' else "❌"
                    action_type = action.get('type', 'python')
                    
                    # Color-code by type
                    if action_type == 'lambda':
                        bg_color = "#1a237e"  # Deep blue for Lambda
                        border_color = "#3f51b5"
                    else:
                        bg_color = "#1b5e20"  # Green for Python
                        border_color = "#4caf50"
                    
                    # Action header
                    st.markdown(f"""
                    <div style='background:{bg_color};padding:0.8rem;border-left:4px solid {border_color};
                    margin:0.5rem 0;border-radius:0.4rem;color:#ffffff;'>
                        <div style='font-size:1.1em;font-weight:600;'>{status_icon} {action['icon']} {action['name']}</div>
                        <div style='font-size:0.9em;margin-top:0.3rem;opacity:0.9;'>{action['description']}</div>
                        <div style='font-size:0.85em;margin-top:0.3rem;opacity:0.7;'>
                            Type: {action_type.upper()} | Duration: {action['execution_time']:.2f}s
                        </div>
                    </div>
                    """, unsafe_allow_html=True)
                    
                    # Show action results
                    with st.expander(f"View {action['name']} Results"):
                        details = action.get('details', {})
                        if details:
                            st.json(details)
            else:
                st.info("No executor actions for this query")
        else:
            st.info("No executor actions for this query")
    else:
        st.info("Process a query to see executor actions")


# ===== TAB 5: MEMORY & CONTEXT (RAG + Memory system) =====
with tab5:
    st.markdown("### 🧠 Memory & Knowledge Base Context")
    st.markdown("""
    **Retrieval-Augmented Generation (RAG):** P-CABA retrieves relevant domain knowledge 
    from uploaded documents and past conversations to inform decisions.
    """)
    
    # Get current history
    current_history = st.session_state.all_sessions[st.session_state.current_session]
    
    if current_history and not current_history[-1].get("thinking"):
        latest = current_history[-1]
        
        # RAG Knowledge Base Sources
        rag_helper = get_rag_helper()
        has_rag_sources = False
        
        # Check cognitive agent RAG results
        if 'cognitive' in latest and latest['cognitive'].get('rag_result'):
            cog_rag = latest['cognitive']['rag_result']
            if cog_rag.get('source_count', 0) > 0:
                has_rag_sources = True
                st.markdown("#### 📚 Cognitive Agent: Knowledge Base Sources")
                
                # Show summary
                domain_icon = cog_rag.get('domain_icon', '📚')
                domain_name = cog_rag.get('domain_name', 'Domain')
                source_count = cog_rag.get('source_count', 0)
                
                st.success(f"{domain_icon} Retrieved **{source_count} expert sources** from {domain_name} domain")
                
                # Show sources
                ui_text = rag_helper.format_sources_for_ui(cog_rag)
                st.markdown(ui_text)
        
        # Check emotional agent RAG results
        if 'emotional' in latest and latest['emotional'].get('rag_result'):
            emo_rag = latest['emotional']['rag_result']
            if emo_rag.get('source_count', 0) > 0:
                st.markdown("---")
                st.markdown("#### 📚 Emotional Agent: Knowledge Base Sources")
                ui_text = rag_helper.format_sources_for_ui(emo_rag)
                st.markdown(ui_text)
                has_rag_sources = True
        
        if not has_rag_sources:
            st.info("No RAG knowledge base sources retrieved for this query")
        
        # Memory Context
        st.markdown("---")
        st.markdown("#### 🧠 Memory Context (Past Conversations)")
        
        if 'memory_context' in latest and latest['memory_context']:
            # Remove duplicates while preserving order
            unique_memories = []
            seen = set()
            for mem in latest['memory_context']:
                if mem not in seen:
                    unique_memories.append(mem)
                    seen.add(mem)
                    if len(unique_memories) >= 5:  # Show more memories in this tab
                        break
            
            if unique_memories:
                for idx, mem in enumerate(unique_memories, 1):
                    st.markdown(f"**{idx}.** {mem}")
            else:
                st.info("No unique memory context available")
        else:
            st.info("No memory context available for this query")
    else:
        st.info("Process a query to see memory and knowledge base context")


# ===== TAB 6: ARCHITECTURE (System design showcase) =====
with tab6:
    st.markdown("## 🏗️ P-CABA Multi-Layer Architecture")
    st.markdown("""
    P-CABA uses a **multi-layered architecture** where:
    1. **Sub-agents** provide specialized analysis (Ethics, Social, Sentiment, PastExperience)
    2. **Parent agents** (Cognitive + Emotional) synthesize sub-agent insights and **debate with each other**
    3. **Synthesizer** reconciles the debate into final decision
    
    This creates a rich, multi-perspective reasoning system!
    """)
    
    # Architecture explanation
    st.markdown("### 🎯 Architecture Components")
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("**🌩️ Amazon Bedrock AgentCore**")
        st.markdown("""
        - Cognitive Agent synthesis
        - Ethics & Social Action Groups
        - nova-micro-v1:0 model
        - Agent ID: IKBWDETWVC
        """)
    
    with col2:
        st.markdown("**📊 LangGraph**")
        st.markdown("""
        - Emotional Agent orchestration
        - Sentiment sub-agent
        - PastExperience sub-agent
        - Local state management
        """)
    
    with col3:
        st.markdown("**🔄 Amazon Strands SDK**")
        st.markdown("""
        - Synthesizer multi-turn reasoning
        - Up to 3 refinement rounds
        - Confidence threshold: 0.85
        - bedrock-runtime converse API
        """)
    
    st.markdown("---")
    
    # Interactive graph
    st.markdown("### 🔗 Interactive Agent Graph")
    st.markdown("**Click on nodes to see detailed information!**")
    
    # Get latest conversation data for graph
    if current_history and len([t for t in current_history if not t.get("thinking")]) > 0:
        latest_turn = [t for t in current_history if not t.get("thinking")][-1]
        
        # Build agent timings from latest result
        agent_timings = {
            'controller': 0.5,
            'cognitive': latest_turn.get('cognitive', {}).get('time_taken', 0),
            'emotional': latest_turn.get('emotional', {}).get('time_taken', 0),
            'synthesizer': latest_turn.get('synthesis', {}).get('time_taken', 0),
            'ethics': 0,
            'social': 0,
            'sentiment': 0,
            'past_experience': 0
        }
        
        # Extract sub-agent times from outputs
        cog_outputs = latest_turn.get('cognitive', {}).get('outputs', [])
        emo_outputs = latest_turn.get('emotional', {}).get('outputs', [])
        
        for output in cog_outputs:
            if 'Ethics' in output.get('agent', ''):
                agent_timings['ethics'] = latest_turn.get('cognitive', {}).get('action_time', 2.0) / 2
            elif 'Social' in output.get('agent', ''):
                agent_timings['social'] = latest_turn.get('cognitive', {}).get('action_time', 2.0) / 2
        
        for output in emo_outputs:
            if 'Sentiment' in output.get('agent', ''):
                agent_timings['sentiment'] = latest_turn.get('emotional', {}).get('time_taken', 0) / 2
            elif 'Past' in output.get('agent', ''):
                agent_timings['past_experience'] = latest_turn.get('emotional', {}).get('time_taken', 0) / 2
        
        # Build agent thoughts
        agent_thoughts = {
            'controller': 'Orchestrates multi-agent workflow and manages state',
            'cognitive': latest_turn.get('cognitive', {}).get('analysis', 'Cognitive reasoning'),
            'emotional': latest_turn.get('emotional', {}).get('analysis', 'Emotional reasoning'),
            'synthesizer': latest_turn.get('synthesis', {}).get('final_answer', 'Final synthesis'),
            'ethics': cog_outputs[0].get('analysis', 'Ethical analysis') if len(cog_outputs) > 0 else 'Ethics analysis',
            'social': cog_outputs[1].get('analysis', 'Social analysis') if len(cog_outputs) > 1 else 'Social impact analysis',
            'sentiment': emo_outputs[0].get('analysis', 'Sentiment analysis') if len(emo_outputs) > 0 else 'Sentiment analysis',
            'past_experience': emo_outputs[1].get('analysis', 'Past experience analysis') if len(emo_outputs) > 1 else 'Past experience analysis'
        }
        
        # Get debate time
        debate_time = latest_turn.get('debate_time', 0)
        
        # Create enhanced graph with architecture labels
        from agent_activation_viz import visualize_agent_activation
        fig = visualize_agent_activation(
            agent_timings, 
            agent_thoughts, 
            debate_time=debate_time,
            show_architecture_labels=True
        )
        st.plotly_chart(fig, use_container_width=True)
        
        # Architecture summary
        st.markdown("### 📋 Architecture Summary")
        
        sum_col1, sum_col2, sum_col3, sum_col4 = st.columns(4)
        
        with sum_col1:
            total_action_time = agent_timings.get('ethics', 0) + agent_timings.get('social', 0)
            st.metric(
                "Action Groups", 
                f"{total_action_time:.2f}s",
                delta="Ethics + Social",
                help="Bedrock Action Groups execution time"
            )
        
        with sum_col2:
            st.metric(
                "LangGraph", 
                f"{agent_timings.get('emotional', 0):.2f}s",
                delta="Emotional Agent",
                help="LangGraph execution time"
            )
        
        with sum_col3:
            st.metric(
                "Strands SDK", 
                f"{agent_timings.get('synthesizer', 0):.2f}s",
                delta="Synthesis",
                help="Strands multi-turn refinement time"
            )
        
        with sum_col4:
            total_time = sum(agent_timings.values())
            st.metric(
                "Total Time", 
                f"{total_time:.2f}s",
                delta="End-to-end",
                help="Complete processing time"
            )
        
        # Detailed sub-agent breakdown
        st.markdown("### 🔍 Sub-Agent Details")
        
        detail_col1, detail_col2 = st.columns(2)
        
        with detail_col1:
            st.markdown("**🧠 Cognitive Sub-Agents** (AgentCore Action Groups)")
            for output in cog_outputs:
                agent_name = output.get('agent', 'Unknown')
                confidence = output.get('confidence', 0)
                analysis = output.get('analysis', '')  # NO TRUNCATION
                
                with st.expander(f"🎯 {agent_name} (Confidence: {confidence:.2%})"):
                    st.write(analysis)  # NO "..."
                    if 'concerns' in output:
                        st.write("**Concerns:**")
                        for concern in output.get('concerns', []):
                            st.write(f"  • {concern}")
                    if 'affected_parties' in output:
                        st.write("**Affected Parties:**")
                        for party in output.get('affected_parties', []):
                            st.write(f"  • {party}")
        
        with detail_col2:
            st.markdown("**❤️ Emotional Sub-Agents** (LangGraph)")
            for output in emo_outputs:
                agent_name = output.get('agent', 'Unknown')
                confidence = output.get('confidence', 0)
                analysis = output.get('analysis', '')  # NO TRUNCATION
                
                with st.expander(f"📊 {agent_name} (Confidence: {confidence:.2%})"):
                    st.write(analysis)  # NO "..."
        
        # Debate visualization
        if debate_time and debate_time > 0:
            st.markdown("### 💬 Cognitive ↔ Emotional Debate")
            st.info(f"**Debate Duration:** {debate_time:.3f}s")
            
            if latest_turn.get('debate_log'):
                with st.expander("View Debate Transcript"):
                    st.code(latest_turn['debate_log'], language=None)
    
    else:
        st.info("💡 Run a query in the Chat tab to see the architecture in action!")
        
        # Show static architecture diagram
        st.markdown("### 📊 Architecture Overview")
        st.markdown("""
        ```
        Controller (Orchestrator)
            ↓
        Synthesizer (Strands SDK)
            ↓
        ┌─────────────────────────────┬─────────────────────────────┐
        │   Cognitive Agent           │   Emotional Agent           │
        │   (AgentCore + Actions)     │   (LangGraph)               │
        ├─────────────────────────────┼─────────────────────────────┤
        │ ├─ Ethics (Action Group)    │ ├─ Sentiment (Sub-Agent)    │
        │ └─ Social (Action Group)    │ └─ PastExperience (Sub)     │
        └─────────────────────────────┴─────────────────────────────┘
                        ↓                           ↓
                Cognitive vs Emotional Debate (Multi-Round)
                        ↓
                Final Synthesis (Strands SDK)
        ```
        """)

# Footer
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #666;'>
    <p>P-CABA v1.0 | Hackathon Demo | Multi-Agent Adaptive Brain System</p>
    <p>Powered by LangGraph + Ollama + Streamlit</p>
</div>
""", unsafe_allow_html=True)

# --- Process agent thinking AFTER rendering (to keep old UI visible) ---
# This runs after all UI is rendered, so old content stays visible during processing
if current_history and current_history[-1].get("thinking"):
    last_prompt = current_history[-1]["query"]
    selected_domain_id = current_history[-1].get("domain_id", "healthcare")  # Get domain from history
    
    # Show processing indicator
    with st.spinner("🧠 Multi-agent reasoning in progress..."):
        start_time = time.time()
        state = {
            "query": last_prompt,
            "domain_id": selected_domain_id,  # Pass domain to agents
            "email_enabled": st.session_state.get("email_enabled", False),  # Pass email settings
            "user_email": st.session_state.get("user_email", "")  # Pass user email
        }
        controller = AdaptiveController()
        result = controller.process(state)
        end_time = time.time()
        processing_time = end_time - start_time
    
    # Add query and domain to result for proper rendering
    result["query"] = last_prompt
    result["domain_id"] = selected_domain_id
    
    # Replace last entry with agent response
    st.session_state.all_sessions[st.session_state.current_session][-1] = result
    st.rerun()