src/streamlit_app.py

"""
THISverse AI - Developer Productivity Agent
Modern UI with enhanced visual design
"""

import streamlit as st
import json
from pathlib import Path
import sys
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime

sys.path.append(str(Path(__file__).parent))

from main import DevProductivityAgent, JiraTicket, Config, cost_tracker

# ============================================================================
# Page Config
# ============================================================================

st.set_page_config(
    page_title="THISverse AI - Code Assistant",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ============================================================================
# Enhanced CSS with Modern Design
# ============================================================================

st.markdown("""
<style>
    /* Import Google Fonts */
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700;900&display=swap');
    
    /* Global Styles */
    * {
        font-family: 'Inter', sans-serif;
    }
    
    /* Hide Streamlit defaults */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    header {visibility: hidden;}
    
    /* Main container */
    .main {
        padding-top: 0rem;
    }
    
    /* Responsive container */
    .block-container {
        padding-top: 1rem;
        padding-bottom: 1rem;
        max-width: 100%;
    }
    
    /* Brand Header - Responsive */
    .brand-header {
        background: linear-gradient(135deg, #0f0f0f 0%, #1a1a2e 50%, #16213e 100%);
        padding: 1.5rem 1rem;
        border-radius: 0 0 16px 16px;
        margin: -1rem -1rem 1.5rem -1rem;
        box-shadow: 0 4px 16px rgba(0, 0, 0, 0.4);
        position: relative;
        overflow: hidden;
    }
    
    .brand-header::before {
        content: '';
        position: absolute;
        top: -50%;
        right: -10%;
        width: 200px;
        height: 200px;
        background: radial-gradient(circle, rgba(64, 224, 208, 0.15) 0%, transparent 70%);
        border-radius: 50%;
    }
    
    .brand-container {
        display: flex;
        align-items: center;
        gap: 1rem;
        position: relative;
        z-index: 1;
    }
    
    .brand-logo {
        width: 50px;
        height: 50px;
        min-width: 50px;
        display: flex;
        align-items: center;
        justify-content: center;
        background: linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 25%, #20B2AA 60%, #40E0D0 100%);
        border-radius: 12px;
        box-shadow: 
            0 0 20px rgba(64, 224, 208, 0.6),
            0 4px 16px rgba(0, 0, 0, 0.8),
            inset 0 1px 0 rgba(255, 255, 255, 0.1);
        flex-shrink: 0;
        border: 2px solid rgba(64, 224, 208, 0.3);
        animation: glow 3s ease-in-out infinite;
    }
    
    @keyframes glow {
        0%, 100% { box-shadow: 0 0 20px rgba(64, 224, 208, 0.6), 0 4px 16px rgba(0, 0, 0, 0.8), inset 0 1px 0 rgba(255, 255, 255, 0.1); }
        50% { box-shadow: 0 0 30px rgba(64, 224, 208, 0.9), 0 4px 16px rgba(0, 0, 0, 0.8), inset 0 1px 0 rgba(255, 255, 255, 0.1); }
    }
    
    .brand-logo svg {
        width: 30px;
        height: 30px;
        filter: drop-shadow(0 2px 4px rgba(255, 255, 255, 0.3));
    }
    
    .brand-info {
        flex: 1;
        min-width: 0;
    }
    
    .brand-title {
        font-size: clamp(1.2rem, 4vw, 2rem);
        font-weight: 900;
        background: linear-gradient(135deg, #ffffff 0%, #40E0D0 50%, #20B2AA 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
        margin: 0 0 0.2rem 0;
        letter-spacing: -0.5px;
        line-height: 1.2;
        word-break: break-word;
    }
    
    .brand-subtitle {
        color: rgba(255, 255, 255, 0.7);
        font-size: clamp(0.7rem, 2vw, 0.9rem);
        font-weight: 500;
        margin: 0;
        letter-spacing: 0.3px;
    }
    
    /* Metric Cards - Responsive */
    .metric-card, .savings-card, .cost-card {
        padding: 1.2rem;
        border-radius: 12px;
        color: white;
        text-align: center;
        margin-bottom: 1rem;
        box-shadow: 0 4px 16px rgba(0, 0, 0, 0.2);
        transition: transform 0.3s ease, box-shadow 0.3s ease;
        border: 1px solid rgba(255, 255, 255, 0.1);
    }
    
    .metric-card {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    }
    
    .savings-card {
        background: linear-gradient(135deg, #11998e 0%, #38ef7d 100%);
    }
    
    .cost-card {
        background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    }
    
    .metric-card:hover, .savings-card:hover, .cost-card:hover {
        transform: translateY(-2px);
        box-shadow: 0 6px 20px rgba(0, 0, 0, 0.3);
    }
    
    .metric-value {
        font-size: clamp(1.5rem, 5vw, 2rem);
        font-weight: 900;
        margin-bottom: 0.3rem;
        text-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
    }
    
    .metric-label {
        font-size: clamp(0.75rem, 2vw, 0.9rem);
        font-weight: 600;
        opacity: 0.95;
        letter-spacing: 0.3px;
    }
    
    /* Chat Messages - Responsive */
    .chat-message {
        padding: 0.8rem 1rem;
        border-radius: 10px;
        margin-bottom: 0.8rem;
        box-shadow: 0 2px 6px rgba(0, 0, 0, 0.05);
        transition: transform 0.2s ease;
        animation: slideIn 0.3s ease;
        word-wrap: break-word;
        overflow-wrap: break-word;
    }
    
    @keyframes slideIn {
        from {
            opacity: 0;
            transform: translateY(10px);
        }
        to {
            opacity: 1;
            transform: translateY(0);
        }
    }
    
    .chat-message:hover {
        transform: translateX(2px);
    }
    
    .user-message {
        background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%);
        border-left: 3px solid #2196f3;
    }
    
    .assistant-message {
        background: linear-gradient(135deg, #f5f5f5 0%, #e0e0e0 100%);
        border-left: 3px solid #4caf50;
    }
    
    /* Info Cards - Responsive */
    .info-card, .warning-card, .success-card {
        padding: 1rem;
        border-radius: 10px;
        margin-bottom: 1rem;
        box-shadow: 0 2px 8px rgba(0, 0, 0, 0.05);
        font-size: clamp(0.85rem, 2vw, 1rem);
    }
    
    .info-card {
        background: linear-gradient(135deg, #f0f4f8 0%, #d9e2ec 100%);
        border-left: 3px solid #3b82f6;
    }
    
    .warning-card {
        background: linear-gradient(135deg, #fff4e6 0%, #ffe0b2 100%);
        border-left: 3px solid #ff9800;
    }
    
    .success-card {
        background: linear-gradient(135deg, #e8f5e9 0%, #c8e6c9 100%);
        border-left: 3px solid #4caf50;
    }
    
    /* Tabs - Responsive */
    .stTabs [data-baseweb="tab-list"] {
        gap: 4px;
        background: linear-gradient(135deg, #f8f9fa 0%, #e9ecef 100%);
        padding: 0.4rem;
        border-radius: 10px;
        overflow-x: auto;
        white-space: nowrap;
    }
    
    .stTabs [data-baseweb="tab"] {
        border-radius: 8px;
        padding: 0.6rem 1rem;
        font-weight: 600;
        transition: all 0.3s ease;
        font-size: clamp(0.8rem, 2vw, 0.95rem);
    }
    
    .stTabs [aria-selected="true"] {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white !important;
        box-shadow: 0 2px 8px rgba(102, 126, 234, 0.3);
    }
    
    /* Buttons - Responsive */
    .stButton > button {
        border-radius: 8px;
        font-weight: 600;
        transition: all 0.3s ease;
        border: none;
        box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
        font-size: clamp(0.85rem, 2vw, 1rem);
        padding: 0.5rem 1rem;
        width: 100%;
    }
    
    .stButton > button:hover {
        transform: translateY(-1px);
        box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);
    }
    
    /* Forms - Responsive */
    .stTextInput input, .stTextArea textarea {
        font-size: clamp(0.9rem, 2vw, 1rem);
        border-radius: 8px;
    }
    
    /* Code blocks - Responsive */
    .stCodeBlock {
        border-radius: 10px;
        box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
        font-size: clamp(0.75rem, 1.5vw, 0.9rem);
    }
    
    /* Expanders - Responsive */
    .streamlit-expanderHeader {
        background: linear-gradient(135deg, #f8f9fa 0%, #e9ecef 100%);
        border-radius: 8px;
        font-weight: 600;
        padding: 0.8rem;
        font-size: clamp(0.9rem, 2vw, 1rem);
    }
    
    /* Status badges - Responsive */
    .status-badge {
        display: inline-block;
        padding: 0.3rem 0.8rem;
        border-radius: 16px;
        font-size: clamp(0.75rem, 1.5vw, 0.85rem);
        font-weight: 600;
        margin: 0.2rem;
        white-space: nowrap;
    }
    
    .status-success {
        background: linear-gradient(135deg, #4caf50 0%, #66bb6a 100%);
        color: white;
        box-shadow: 0 2px 6px rgba(76, 175, 80, 0.3);
    }
    
    .status-info {
        background: linear-gradient(135deg, #2196f3 0%, #42a5f5 100%);
        color: white;
        box-shadow: 0 2px 6px rgba(33, 150, 243, 0.3);
    }
    
    .status-warning {
        background: linear-gradient(135deg, #ff9800 0%, #ffa726 100%);
        color: white;
        box-shadow: 0 2px 6px rgba(255, 152, 0, 0.3);
    }
    
    /* Mobile-specific adjustments */
    @media only screen and (max-width: 768px) {
        .main {
            padding: 0.5rem;
        }
        
        .block-container {
            padding: 0.5rem;
        }
        
        .brand-header {
            padding: 1rem 0.8rem;
            margin: -0.5rem -0.5rem 1rem -0.5rem;
        }
        
        .brand-container {
            gap: 0.8rem;
        }
        
        .brand-logo {
            width: 40px;
            height: 40px;
            min-width: 40px;
        }
        
        .brand-logo svg {
            width: 24px;
            height: 24px;
        }
        
        .metric-card, .savings-card, .cost-card {
            padding: 1rem;
            margin-bottom: 0.8rem;
        }
        
        .chat-message {
            padding: 0.7rem 0.8rem;
            font-size: 0.9rem;
        }
        
        .stTabs [data-baseweb="tab"] {
            padding: 0.5rem 0.8rem;
            font-size: 0.85rem;
        }
        
        /* Stack columns on mobile */
        .row-widget.stHorizontal {
            flex-direction: column;
        }
        
        /* Full width buttons on mobile */
        .stButton > button {
            width: 100%;
            margin-bottom: 0.5rem;
        }
    }
    
    /* Tablet adjustments */
    @media only screen and (min-width: 769px) and (max-width: 1024px) {
        .brand-title {
            font-size: 1.8rem;
        }
        
        .metric-value {
            font-size: 1.8rem;
        }
    }
    
    /* Large desktop */
    @media only screen and (min-width: 1920px) {
        .block-container {
            max-width: 1600px;
            margin: 0 auto;
        }
    }
    
    /* Improve scrolling on mobile */
    @media only screen and (max-width: 768px) {
        .stTabs [data-baseweb="tab-list"] {
            -webkit-overflow-scrolling: touch;
            scrollbar-width: thin;
        }
        
        .stTabs [data-baseweb="tab-list"]::-webkit-scrollbar {
            height: 4px;
        }
        
        .stTabs [data-baseweb="tab-list"]::-webkit-scrollbar-thumb {
            background: #667eea;
            border-radius: 4px;
        }
    }
</style>
""", unsafe_allow_html=True)

# ============================================================================
# Session State
# ============================================================================

if "agent" not in st.session_state:
    st.session_state.agent = DevProductivityAgent()

if "messages" not in st.session_state:
    st.session_state.messages = []

if "current_plan" not in st.session_state:
    st.session_state.current_plan = None

if "indexed" not in st.session_state:
    st.session_state.indexed = False

if "auto_indexed" not in st.session_state:
    st.session_state.auto_indexed = False

# ============================================================================
# Sidebar
# ============================================================================

with st.sidebar:
    # Modern Brand Header
    st.markdown("""
    <div class="brand-header">
        <div class="brand-container">
            <div class="brand-logo">
                <svg viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
                    <path d="M12 2L2 7L12 12L22 7L12 2Z" fill="#FFFFFF" opacity="0.95"/>
                    <path d="M2 17L12 22L22 17L12 12L2 17Z" fill="#FFFFFF" opacity="0.95"/>
                    <path d="M2 12L12 17L22 12L12 7L2 12Z" fill="#FFFFFF" opacity="0.75"/>
                </svg>
            </div>
            <div class="brand-info">
                <div class="brand-title">THISverse AI</div>
                <div class="brand-subtitle">Code Intelligence Platform</div>
            </div>
        </div>
    </div>
    """, unsafe_allow_html=True)
    
    st.markdown("---")
    
    # Configuration Section
    st.markdown("### ⚙️ Configuration")
    
    with st.expander("ℹ️ System Info", expanded=False):
        st.markdown("""
        <div style="font-size: 0.9rem;">
            <strong>🗄️ Vector Database:</strong> Pinecone<br>
            <strong>🤖 AI Model:</strong> GPT-4o-mini<br>
            <strong>📊 Architecture:</strong> Divided LLM<br>
            <strong>💾 Embedding:</strong> text-embedding-3-small
        </div>
        """, unsafe_allow_html=True)
    
    st.markdown("---")
    
    # API Keys Section
    st.markdown("### 🔑 API Configuration")
    
    openai_key = st.text_input(
        "OpenAI API Key",
        type="password",
        value=st.session_state.get("openai_key", ""),
        help="Your OpenAI API key for LLM and embeddings"
    )
    
    pinecone_key = st.text_input(
        "Pinecone API Key",
        type="password",
        value=st.session_state.get("pinecone_key", ""),
        help="Your Pinecone API key for vector storage"
    )
    
    if openai_key and pinecone_key:
        st.session_state.openai_key = openai_key
        st.session_state.pinecone_key = pinecone_key
        st.session_state.agent.set_api_keys(openai_key, pinecone_key)
        
        # Auto-index on first load
        if not st.session_state.get("auto_indexed", False):
            app_dir = Path(__file__).parent
            sample_codebase_path = str(app_dir / "sample_codebase")
            
            if Path(sample_codebase_path).exists():
                try:
                    with st.spinner("🔄 Indexing codebase..."):
                        results = st.session_state.agent.index_codebase(
                            sample_codebase_path,
                            extensions=[".py", ".js", ".ts", ".jsx", ".tsx"]
                        )
                        st.session_state.auto_indexed = True
                        st.session_state.indexed = True
                        st.success(f"✅ Indexed: {results['files_indexed']} files, {results['total_chunks']} chunks")
                except Exception as e:
                    st.warning(f"⚠️ Indexing failed: {str(e)[:100]}")
            else:
                st.markdown('<span class="status-badge status-success">✅ API Keys Configured</span>', unsafe_allow_html=True)
        else:
            st.markdown('<span class="status-badge status-success">✅ API Keys Active</span>', unsafe_allow_html=True)
    else:
        st.markdown('<span class="status-badge status-warning">⚠️ API Keys Required</span>', unsafe_allow_html=True)
    
    st.markdown("---")
    
    # Codebase Status
    st.markdown("### 📂 Codebase Status")
    
    try:
        stats = st.session_state.agent.indexer.get_stats()
        total_chunks = stats.get('total_chunks', 0)
        if total_chunks > 0:
            st.markdown(f"""
            <div class="success-card">
                <strong>✅ Codebase Indexed</strong><br>
                <span style="font-size: 1.5rem; font-weight: 700;">{total_chunks:,}</span> chunks ready
            </div>
            """, unsafe_allow_html=True)
            st.session_state.indexed = True
        else:
            st.markdown("""
            <div class="info-card">
                <strong>ℹ️ Ready to Index</strong><br>
                Set API keys to auto-index
            </div>
            """, unsafe_allow_html=True)
    except:
        st.markdown("""
        <div class="info-card">
            <strong>ℹ️ Awaiting Configuration</strong><br>
            Configure API keys above
        </div>
        """, unsafe_allow_html=True)
    
    st.markdown("---")
    
    # Quick Actions
    st.markdown("### ⚡ Quick Actions")
    
    col1, col2 = st.columns(2)
    with col1:
        if st.button("🔄 Refresh", use_container_width=True):
            st.rerun()
    with col2:
        if st.button("🗑️ Clear Chat", use_container_width=True):
            st.session_state.messages = []
            st.rerun()

# ============================================================================
# Main Content
# ============================================================================

# Hero Section
st.markdown("""
<div style="text-align: center; padding: 2rem 0 1rem 0;">
    <h1 style="font-size: 3rem; font-weight: 900; margin-bottom: 0.5rem; 
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
        🚀 AI-Powered Code Assistant
    </h1>
    <p style="font-size: 1.1rem; color: #6c757d; font-weight: 500;">
        Analyze tickets, generate code, and understand your codebase with AI
    </p>
</div>
""", unsafe_allow_html=True)

# Tabs
tab1, tab2, tab3, tab4 = st.tabs([
    "💬 Chat Assistant",
    "🎫 Ticket Processing",
    "📋 Implementation Plan",
    "📊 Cost Analytics"
])

# ============================================================================
# Tab 1: Chat Assistant
# ============================================================================

with tab1:
    st.markdown("### 💬 Ask Questions About Your Code")
    
    # Check indexing status
    if not st.session_state.get("indexed", False):
        try:
            stats = st.session_state.agent.indexer.get_stats()
            if stats.get('total_chunks', 0) > 0:
                st.session_state.indexed = True
        except:
            pass
    
    if not st.session_state.get("indexed", False):
        st.markdown("""
        <div class="warning-card">
            <strong>⚠️ Codebase Not Indexed</strong><br>
            Please configure your API keys in the sidebar to automatically index the codebase.
        </div>
        """, unsafe_allow_html=True)
    else:
        st.markdown("""
        <div class="success-card">
            <strong>✅ Ready to Answer</strong><br>
            Your codebase is indexed and ready for questions!
        </div>
        """, unsafe_allow_html=True)
    
    # Display chat history
    for msg in st.session_state.messages:
        cls = "user-message" if msg["role"] == "user" else "assistant-message"
        icon = "👤" if msg["role"] == "user" else "🤖"
        st.markdown(
            f'<div class="chat-message {cls}"><strong>{icon}</strong> {msg["content"]}</div>',
            unsafe_allow_html=True
        )
    
    # Chat input
    st.markdown("---")
    with st.form("chat_form", clear_on_submit=True):
        prompt = st.text_input(
            "💭 Your Question",
            placeholder="e.g., How does user authentication work in this codebase?",
            key="chat_input"
        )
        
        col1, col2 = st.columns([3, 1])
        with col2:
            submit = st.form_submit_button("🚀 Send", type="primary", use_container_width=True)
        
        if submit and prompt:
            if not (st.session_state.get("openai_key") and st.session_state.get("pinecone_key")):
                st.error("🔑 Please configure API keys in the sidebar")
            elif not st.session_state.get("indexed", False):
                st.warning("⚠️ Please wait for codebase indexing to complete")
            else:
                st.session_state.messages.append({"role": "user", "content": prompt})
                with st.spinner("🤔 Thinking..."):
                    try:
                        response = st.session_state.agent.ask_about_code(prompt)
                        st.session_state.messages.append({"role": "assistant", "content": response})
                        st.rerun()
                    except Exception as e:
                        st.error(f"❌ Error: {str(e)}")
                        st.session_state.messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
                        st.rerun()

# ============================================================================
# Tab 2: Ticket Processing
# ============================================================================

with tab2:
    st.markdown("### 🎫 JIRA Ticket to Implementation")
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.markdown("#### Ticket Details")
        ticket_id = st.text_input("📋 Ticket ID", placeholder="PROJ-123")
        ticket_title = st.text_input("📝 Title", placeholder="Add user preferences feature")
        ticket_desc = st.text_area(
            "📄 Description",
            height=150,
            placeholder="Detailed description of the feature..."
        )
        acceptance = st.text_area(
            "✅ Acceptance Criteria",
            height=100,
            placeholder="List the acceptance criteria..."
        )
        labels = st.text_input("🏷️ Labels", placeholder="backend, frontend (comma-separated)")
    
    with col2:
        st.markdown("#### Quick Actions")
        
        if st.button("📝 Load Example", use_container_width=True):
            st.session_state.example = {
                "id": "PROJ-456",
                "title": "Add notification preferences",
                "desc": "As a user, I want to customize my notification settings.\n\nFeatures:\n- Email notification toggle\n- Push notification toggle\n- Notification frequency settings\n- Save preferences to database",
                "accept": "- User can toggle email notifications\n- User can toggle push notifications\n- Settings persist across sessions\n- Changes take effect immediately",
                "labels": "backend, frontend, notifications"
            }
            st.rerun()
        
        if "example" in st.session_state:
            ticket_id = st.session_state.example["id"]
            ticket_title = st.session_state.example["title"]
            ticket_desc = st.session_state.example["desc"]
            acceptance = st.session_state.example["accept"]
            labels = st.session_state.example["labels"]
            del st.session_state.example
        
        st.markdown("---")
        
        st.markdown("""
        <div class="info-card">
            <strong>💡 Tip</strong><br>
            Provide detailed descriptions for better code generation
        </div>
        """, unsafe_allow_html=True)
    
    st.markdown("---")
    
    if st.button("🚀 Generate Implementation Plan", type="primary", use_container_width=True):
        if not (st.session_state.get("openai_key") and st.session_state.get("pinecone_key")):
            st.error("🔑 Please configure API keys first")
        elif not ticket_title or not ticket_desc:
            st.error("📝 Please provide at least a title and description")
        else:
            ticket = JiraTicket(
                ticket_id=ticket_id or "UNKNOWN",
                title=ticket_title,
                description=ticket_desc,
                acceptance_criteria=acceptance or None,
                labels=[l.strip() for l in labels.split(",")] if labels else None
            )
            
            with st.spinner("⚙️ Generating implementation plan..."):
                try:
                    plan = st.session_state.agent.process_ticket(ticket)
                    st.session_state.current_plan = plan
                    st.session_state.last_ticket = ticket
                    st.success("✅ Plan generated! Check the 'Implementation Plan' tab")
                    st.balloons()
                except Exception as e:
                    st.error(f"❌ Error: {str(e)}")

# ============================================================================
# Tab 3: Implementation Plan
# ============================================================================

with tab3:
    if st.session_state.current_plan:
        plan = st.session_state.current_plan
        
        # Summary Section
        st.markdown("### 📋 Implementation Summary")
        st.markdown(f"""
        <div class="info-card">
            {plan.ticket_summary}
        </div>
        """, unsafe_allow_html=True)
        
        # Metrics
        col1, col2, col3 = st.columns(3)
        with col1:
            st.markdown(f"""
            <div class="metric-card">
                <div class="metric-value">{plan.estimated_complexity}</div>
                <div class="metric-label">Complexity</div>
            </div>
            """, unsafe_allow_html=True)
        with col2:
            st.markdown(f"""
            <div class="savings-card">
                <div class="metric-value">{len(plan.key_entities)}</div>
                <div class="metric-label">Key Entities</div>
            </div>
            """, unsafe_allow_html=True)
        with col3:
            st.markdown(f"""
            <div class="cost-card">
                <div class="metric-value">{len(plan.relevant_files)}</div>
                <div class="metric-label">Related Files</div>
            </div>
            """, unsafe_allow_html=True)
        
        st.markdown("---")
        
        # Details
        col1, col2 = st.columns(2)
        
        with col1:
            st.markdown("### 🏷️ Key Entities")
            for e in plan.key_entities:
                st.markdown(f'<span class="status-badge status-info">{e}</span>', unsafe_allow_html=True)
            
            st.markdown("### ⚠️ Prerequisites")
            if plan.prerequisites:
                for p in plan.prerequisites:
                    st.markdown(f"- {p}")
            else:
                st.info("No prerequisites identified")
        
        with col2:
            st.markdown("### 📝 Implementation Steps")
            for i, s in enumerate(plan.implementation_steps, 1):
                st.markdown(f"**{i}.** {s}")
        
        st.markdown("---")
        
        # Architecture
        st.markdown("### 🏗️ Architecture Notes")
        st.markdown(f"""
        <div class="info-card">
            {plan.architecture_notes}
        </div>
        """, unsafe_allow_html=True)
        
        st.markdown("---")
        
        # Generated Code
        st.markdown("### 💻 Generated Code")
        
        for path, code in plan.boilerplate_code.items():
            with st.expander(f"📄 {path}", expanded=True):
                lang = {
                    '.py': 'python',
                    '.js': 'javascript',
                    '.ts': 'typescript',
                    '.jsx': 'javascript',
                    '.tsx': 'typescript'
                }.get(Path(path).suffix, 'text')
                
                st.code(code, language=lang)
                
                col1, col2 = st.columns([1, 3])
                with col1:
                    st.download_button(
                        "⬇️ Download",
                        code,
                        Path(path).name,
                        key=f"download_{path}"
                    )
        
        st.markdown("---")
        
        # Modification Section
        st.markdown("### ✏️ Request Modifications")
        st.markdown("""
        <div class="info-card">
            <strong>💡 Need Changes?</strong><br>
            Describe modifications and we'll regenerate the code
        </div>
        """, unsafe_allow_html=True)
        
        modification_request = st.text_area(
            "What would you like to change?",
            placeholder="e.g., Add error handling, use async/await, add unit tests...",
            height=100,
            key="modification_request"
        )
        
        if st.button("🔄 Regenerate with Modifications", type="primary", use_container_width=True):
            if not modification_request.strip():
                st.warning("⚠️ Please describe the modifications")
            else:
                with st.spinner("⚙️ Regenerating code..."):
                    try:
                        original_ticket = st.session_state.get("last_ticket")
                        if original_ticket:
                            modified_description = f"{original_ticket.description}\n\n---\n\n**Modification Request:**\n{modification_request}"
                            modified_ticket = JiraTicket(
                                ticket_id=original_ticket.ticket_id,
                                title=original_ticket.title,
                                description=modified_description,
                                acceptance_criteria=original_ticket.acceptance_criteria,
                                labels=original_ticket.labels
                            )
                            new_plan = st.session_state.agent.process_ticket(modified_ticket)
                            st.session_state.current_plan = new_plan
                            st.success("✅ Code regenerated!")
                            st.rerun()
                        else:
                            st.error("Original ticket not found")
                    except Exception as e:
                        st.error(f"❌ Error: {str(e)}")
    else:
        st.markdown("""
        <div style="text-align: center; padding: 3rem 0;">
            <h3>📋 No Plan Generated Yet</h3>
            <p style="color: #6c757d;">Process a ticket in the "Ticket Processing" tab to see the implementation plan here</p>
        </div>
        """, unsafe_allow_html=True)

# ============================================================================
# Tab 4: Cost Analytics
# ============================================================================

with tab4:
    st.markdown("### 📊 Cost Analytics Dashboard")
    
    stats = st.session_state.agent.get_cost_stats()
    
    # Top Metrics
    col1, col2, col3, col4 = st.columns(4)
    
    with col1:
        st.markdown(f"""
        <div class="savings-card">
            <div class="metric-value">${stats['savings']:.4f}</div>
            <div class="metric-label">💰 Total Savings</div>
        </div>
        """, unsafe_allow_html=True)
    
    with col2:
        st.markdown(f"""
        <div class="cost-card">
            <div class="metric-value">${stats['actual_cost']:.4f}</div>
            <div class="metric-label">📉 Actual Cost</div>
        </div>
        """, unsafe_allow_html=True)
    
    with col3:
        st.markdown(f"""
        <div class="metric-card">
            <div class="metric-value">{stats['savings_percentage']:.1f}%</div>
            <div class="metric-label">📊 Cost Reduction</div>
        </div>
        """, unsafe_allow_html=True)
    
    with col4:
        st.markdown(f"""
        <div class="metric-card">
            <div class="metric-value">{stats['api_calls']}</div>
            <div class="metric-label">🔄 API Calls</div>
        </div>
        """, unsafe_allow_html=True)
    
    st.markdown("---")
    
    # Charts
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("### 💵 Cost Comparison")
        fig = go.Figure(data=[
            go.Bar(
                name='Traditional (GPT-4)',
                x=['Cost'],
                y=[stats['traditional_cost']],
                marker_color='#f5576c',
                text=[f"${stats['traditional_cost']:.4f}"],
                textposition='auto'
            ),
            go.Bar(
                name='Our Approach',
                x=['Cost'],
                y=[stats['actual_cost']],
                marker_color='#38ef7d',
                text=[f"${stats['actual_cost']:.4f}"],
                textposition='auto'
            )
        ])
        fig.update_layout(
            barmode='group',
            height=350,
            margin=dict(l=20, r=20, t=20, b=20),
            showlegend=True
        )
        st.plotly_chart(fig, use_container_width=True)
    
    with col2:
        st.markdown("### 📈 Token Distribution")
        tokens = stats['total_tokens']
        fig = go.Figure(data=[go.Pie(
            labels=['Embeddings', 'Architect In', 'Architect Out', 'Developer In', 'Developer Out'],
            values=[
                tokens['embedding'],
                tokens['architect_input'],
                tokens['architect_output'],
                tokens['developer_input'],
                tokens['developer_output']
            ],
            hole=0.4,
            marker_colors=['#667eea', '#764ba2', '#f093fb', '#11998e', '#38ef7d']
        )])
        fig.update_layout(height=350, margin=dict(l=20, r=20, t=20, b=20))
        st.plotly_chart(fig, use_container_width=True)
    
    st.markdown("---")
    
    # Detailed Stats
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("### 📊 Session Stats")
        st.metric("Tickets Processed", stats['tickets_processed'])
        st.metric("Questions Answered", stats['questions_answered'])
        st.metric("Session Duration", f"{stats['session_duration_minutes']} min")
    
    with col2:
        st.markdown("### 💰 Per-Action Costs")
        st.metric("Cost per Ticket", f"${stats['cost_per_ticket']:.4f}")
        if stats['tickets_processed'] > 0:
            trad = stats['traditional_cost'] / stats['tickets_processed']
            st.metric("Traditional Cost", f"${trad:.4f}")
            st.metric("Savings per Ticket", f"${trad - stats['cost_per_ticket']:.4f}")
    
    with col3:
        st.markdown("### 🔢 Token Breakdown")
        st.write(f"**Embedding:** {tokens['embedding']:,}")
        st.write(f"**Architect:** {tokens['architect_input']:,} / {tokens['architect_output']:,}")
        st.write(f"**Developer:** {tokens['developer_input']:,} / {tokens['developer_output']:,}")
        st.write(f"**Total:** {tokens['total']:,}")
    
    st.markdown("---")
    
    # Architecture Comparison
    st.markdown("### 🏗️ Architecture Comparison")
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("""
        <div class="cost-card">
            <h4>Traditional Approach</h4>
            <ul style="text-align: left; padding-left: 1.5rem;">
                <li>Single GPT-4 model</li>
                <li>$30/1M input tokens</li>
                <li>$60/1M output tokens</li>
                <li>~15,000 tokens/ticket</li>
                <li><strong>~$0.45 per ticket</strong></li>
            </ul>
        </div>
        """, unsafe_allow_html=True)
    
    with col2:
        st.markdown("""
        <div class="savings-card">
            <h4>Our Divided Approach</h4>
            <ul style="text-align: left; padding-left: 1.5rem;">
                <li>Two GPT-4o-mini models</li>
                <li>$0.15/1M input tokens</li>
                <li>$0.60/1M output tokens</li>
                <li>~9,000 tokens/ticket</li>
                <li><strong>~$0.0023 per ticket</strong></li>
            </ul>
        </div>
        """, unsafe_allow_html=True)
    
    st.markdown("---")
    
    # Reset Button
    col1, col2, col3 = st.columns([1, 1, 1])
    with col2:
        if st.button("🔄 Reset Cost Tracking", use_container_width=True):
            st.session_state.agent.reset_cost_tracking()
            st.rerun()

# ============================================================================
# Footer
# ============================================================================

st.markdown("---")
st.markdown("""
<div style="text-align: center; color: #6c757d; padding: 1rem 0;">
    <strong>THISverse AI</strong> - Developer Productivity Agent v2.0<br>
    Powered by GPT-4o-mini, Pinecone, and Divided LLM Architecture
</div>
""", unsafe_allow_html=True)