src/app.py

"""
Main Streamlit application for Data Insights App.
Provides UI for data visualization and AI agent chat interface.
"""
import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from database import DataManager
from tools import AgentTools
from agent import DataInsightsAgent
from support_ticket import SupportTicketManager
from utils import setup_logger, format_price
from config import OPENAI_API_KEY, JIRA_SERVER, JIRA_EMAIL, JIRA_API_TOKEN
import sys

logger = setup_logger(__name__)

# Page configuration
st.set_page_config(
    page_title="Data Insights App",
    page_icon="📱",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .metric-card {
        background-color: #f0f2f6;
        padding: 20px;
        border-radius: 10px;
        margin: 10px 0;
    }
    .stButton>button {
        width: 100%;
    }
    .chat-message {
        padding: 10px;
        border-radius: 5px;
        margin: 5px 0;
    }
    .user-message {
        background-color: #e3f2fd;
    }
    .assistant-message {
        background-color: #f5f5f5;
    }
</style>
""", unsafe_allow_html=True)


@st.cache_resource
def initialize_app():
    """
    Initializes application components (cached for performance).
    Inputs: None
    Outputs: tuple (data_manager, agent_tools, agent, ticket_manager)
    """
    logger.info("Initializing application components...")
    
    try:
        # Initialize data manager
        data_manager = DataManager()
        
        # Initialize tools
        agent_tools = AgentTools(data_manager)
        
        # Initialize agent (will be re-initialized with user's API key if provided)
        agent = None
        
        # Initialize ticket manager
        ticket_manager = SupportTicketManager()
        
        logger.info("Application components initialized successfully")
        
        return data_manager, agent_tools, ticket_manager
        
    except Exception as e:
        logger.error(f"Failed to initialize application: {str(e)}")
        st.error(f"Failed to initialize application: {str(e)}")
        sys.exit(1)


def display_business_metrics(data_manager: DataManager):
    """
    Displays business metrics dashboard.
    Inputs: data_manager (DataManager)
    Outputs: None
    """
    st.header("📊 Business Overview")
    
    # Get dataset stats
    stats = data_manager.get_summary_stats()
    df = data_manager.get_dataframe()
    
    # Display key metrics in columns
    col1, col2, col3, col4 = st.columns(4)
    
    with col1:
        st.metric(
            label="Total Phones",
            value=f"{stats['total_rows']:,}",
            help="Total number of phones in database"
        )
    
    with col2:
        st.metric(
            label="Average Price",
            value=format_price(stats['avg_price']),
            help="Average price across all phones"
        )
    
    with col3:
        st.metric(
            label="Price Range",
            value=f"{format_price(stats['min_price'])} - {format_price(stats['max_price'])}",
            help="Min and max prices"
        )
    
    with col4:
        unique_brands = df['brand'].nunique()
        st.metric(
            label="Brands",
            value=unique_brands,
            help="Number of unique brands"
        )
    
    # Charts
    st.subheader("Data Visualization")
    
    chart_col1, chart_col2 = st.columns(2)
    
    with chart_col1:
        # Price distribution histogram
        fig_price = px.histogram(
            df,
            x='price_usd',
            nbins=50,
            title='Price Distribution',
            labels={'price_usd': 'Price (USD)', 'count': 'Count'},
            color_discrete_sequence=['#1f77b4']
        )
        fig_price.update_layout(showlegend=False)
        st.plotly_chart(fig_price, use_container_width=True)
    
    with chart_col2:
        # Top brands by count
        brand_counts = df['brand'].value_counts().head(10)
        fig_brands = px.bar(
            x=brand_counts.values,
            y=brand_counts.index,
            orientation='h',
            title='Top 10 Brands by Count',
            labels={'x': 'Number of Models', 'y': 'Brand'},
            color_discrete_sequence=['#ff7f0e']
        )
        st.plotly_chart(fig_brands, use_container_width=True)
    
    # Additional charts
    chart_col3, chart_col4 = st.columns(2)
    
    with chart_col3:
        # OS distribution
        os_counts = df['os'].value_counts()
        fig_os = px.pie(
            values=os_counts.values,
            names=os_counts.index,
            title='Operating System Distribution',
            color_discrete_sequence=px.colors.qualitative.Set3
        )
        st.plotly_chart(fig_os, use_container_width=True)
    
    with chart_col4:
        # 5G support
        fiveg_counts = df['5g_support'].value_counts()
        fig_5g = px.pie(
            values=fiveg_counts.values,
            names=fiveg_counts.index,
            title='5G Support Distribution',
            color_discrete_sequence=['#2ecc71', '#e74c3c']
        )
        st.plotly_chart(fig_5g, use_container_width=True)


def display_sample_queries():
    """
    Displays sample query buttons.
    Inputs: None
    Outputs: None
    """
    st.subheader("💡 Sample Queries")
    st.markdown("Try these example queries to get started:")
    
    col1, col2, col3 = st.columns(3)
    
    sample_queries = [
        "What are the top 5 most expensive phones?",
        "Show me Apple phones under $1000",
        "What's the average price by brand?",
        "Find phones with 5G support and 12GB RAM",
        "Compare average prices between Android and iOS",
        "Which brand has the best rated phones?",
        "Show me phones with the largest battery capacity",
        "What are the available brands?",
        "Give me price statistics for Samsung phones"
    ]
    
    cols = [col1, col2, col3]
    for idx, query in enumerate(sample_queries):
        with cols[idx % 3]:
            if st.button(query, key=f"sample_{idx}"):
                st.session_state.sample_query = query


def display_chat_interface(agent: DataInsightsAgent, ticket_manager: SupportTicketManager):
    """
    Displays chat interface with agent.
    Inputs: agent (DataInsightsAgent), ticket_manager (SupportTicketManager)
    Outputs: None
    """
    st.header("💬 Chat with AI Assistant")
    
    # Initialize session state for chat history
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    
    if 'sample_query' not in st.session_state:
        st.session_state.sample_query = None
    
    if 'show_ticket_suggestion' not in st.session_state:
        st.session_state.show_ticket_suggestion = False
    
    # Display chat history
    chat_container = st.container()
    with chat_container:
        for message in st.session_state.chat_history:
            role = message['role']
            content = message['content']
            
            if role == 'user':
                st.markdown(f'<div class="chat-message user-message"><strong>You:</strong> {content}</div>', 
                          unsafe_allow_html=True)
            else:
                st.markdown(f'<div class="chat-message assistant-message"><strong>Assistant:</strong> {content}</div>', 
                          unsafe_allow_html=True)
                
                # Show tool calls if available
                if 'tool_calls' in message and message['tool_calls']:
                    with st.expander("View Function Calls"):
                        for tool_call in message['tool_calls']:
                            st.json(tool_call)
    
    # Chat input
    user_input = st.chat_input("Ask me anything about the mobile phone data...")
    
    # Handle sample query
    if st.session_state.sample_query:
        user_input = st.session_state.sample_query
        st.session_state.sample_query = None
    
    if user_input:
        # Add user message to history
        st.session_state.chat_history.append({
            'role': 'user',
            'content': user_input
        })
        
        # Get agent response
        with st.spinner("Thinking..."):
            try:
                response = agent.chat(user_input)
                
                if response['success']:
                    # Add assistant response to history
                    st.session_state.chat_history.append({
                        'role': 'assistant',
                        'content': response['response'],
                        'tool_calls': response.get('tool_calls_made', [])
                    })
                    
                    # Store support ticket suggestion state
                    if response.get('suggest_support_ticket', False):
                        st.session_state.show_ticket_suggestion = True
                else:
                    st.error(f"Error: {response.get('error', 'Unknown error')}")
                    
            except Exception as e:
                logger.error(f"Error in chat: {str(e)}")
                st.error(f"An error occurred: {str(e)}")
        
        st.rerun()
    
    # Show support ticket suggestion if needed (outside user_input block so it persists)
    if st.session_state.get('show_ticket_suggestion', False):
        st.info("💡 Would you like to create a support ticket for human assistance?")
        col1, col2 = st.columns([1, 4])
        with col1:
            if st.button("Create Support Ticket", key="create_ticket_btn"):
                st.session_state.create_ticket_request = True
                st.session_state.show_ticket_suggestion = False
                st.rerun()
        with col2:
            if st.button("Dismiss", key="dismiss_ticket_btn"):
                st.session_state.show_ticket_suggestion = False
                st.rerun()
    
    # Support ticket creation
    if st.session_state.get('create_ticket_request', False):
        display_ticket_form(agent, ticket_manager)


def display_ticket_form(agent: DataInsightsAgent, ticket_manager: SupportTicketManager):
    """
    Displays support ticket creation form.
    Inputs: agent (DataInsightsAgent), ticket_manager (SupportTicketManager)
    Outputs: None
    """
    st.subheader("🎫 Create Support Ticket")
    
    with st.form("support_ticket_form"):
        ticket_summary = st.text_input(
            "Brief summary of your issue",
            placeholder="e.g., Unable to find specific data"
        )
        
        ticket_description = st.text_area(
            "Detailed description",
            placeholder="Please describe your issue in detail...",
            height=150
        )
        
        user_email = st.text_input(
            "Your email (optional)",
            placeholder="your.email@example.com"
        )
        
        submit_button = st.form_submit_button("Submit Ticket")
        
        if submit_button:
            if not ticket_summary:
                st.error("Please provide a summary for the ticket")
            else:
                # Get conversation context
                conversation_context = agent.get_conversation_summary()
                
                # Create full description
                full_description = f"{ticket_description}\n\n--- Conversation Context ---\n{conversation_context}"
                
                # Create ticket
                with st.spinner("Creating support ticket..."):
                    result = ticket_manager.create_ticket_from_conversation(
                        user_query=ticket_summary,
                        conversation_context=full_description,
                        user_email=user_email if user_email else None
                    )
                
                if result['success']:
                    st.success(f"✅ {result['message']}")
                    st.markdown(f"**Ticket URL:** [{result['ticket_key']}]({result['ticket_url']})")
                    st.session_state.create_ticket_request = False
                else:
                    st.error(f"❌ Failed to create ticket: {result.get('error', 'Unknown error')}")


def main():
    """
    Main application entry point.
    Inputs: None
    Outputs: None
    """
    logger.info("Starting Data Insights App...")
    
    # Sidebar
    with st.sidebar:
        st.title("⚙️ Configuration")
        
        # API Key input
        api_key_input = st.text_input(
            "OpenAI API Key",
            type="password",
            value=OPENAI_API_KEY,
            help="Enter your OpenAI API key"
        )
        
        st.markdown("---")
        
        # Jira configuration status
        st.subheader("Jira Configuration")
        if JIRA_SERVER and JIRA_EMAIL and JIRA_API_TOKEN:
            st.success("✅ Jira configured")
        else:
            st.warning("⚠️ Jira not configured")
            st.info("Set JIRA_SERVER, JIRA_EMAIL, and JIRA_API_TOKEN in .env file")
        
        st.markdown("---")
        
        # Actions
        st.subheader("Actions")
        
        if st.button("🔄 Reset Chat"):
            st.session_state.chat_history = []
            if 'agent' in st.session_state and st.session_state.agent:
                st.session_state.agent.reset_conversation()
            st.rerun()
        
        if st.button("📋 View Logs"):
            st.info("Check your console for detailed logs")
        
        st.markdown("---")
        st.caption("Data Insights App v1.0")
    
    # Main content
    st.title("📱 Mobile Phone Data Insights")
    st.markdown("Get insights from our mobile phone database using AI")
    
    # Initialize components
    data_manager, agent_tools, ticket_manager = initialize_app()
    
    # Initialize agent with API key
    if not api_key_input:
        st.warning("⚠️ Please enter your OpenAI API key in the sidebar to use the chat feature.")
        api_key_valid = False
    else:
        api_key_valid = True
        if 'agent' not in st.session_state or st.session_state.get('api_key') != api_key_input:
            st.session_state.agent = DataInsightsAgent(agent_tools, api_key=api_key_input)
            st.session_state.api_key = api_key_input
            logger.info("Agent initialized with user API key")
    
    # Display sections
    display_business_metrics(data_manager)
    
    st.markdown("---")
    
    display_sample_queries()
    
    st.markdown("---")
    
    if api_key_valid:
        display_chat_interface(st.session_state.agent, ticket_manager)
    else:
        st.info("👆 Enter your OpenAI API key in the sidebar to start chatting with the AI assistant")
    
    # Footer
    st.markdown("---")
    st.markdown("""
    <div style='text-align: center; color: #666;'>
        <p>🛡️ <strong>Safety Features Enabled:</strong> All write operations (INSERT, UPDATE, DELETE, DROP) are blocked</p>
        <p>📞 Need help? Use the chat to create a support ticket</p>
    </div>
    """, unsafe_allow_html=True)


if __name__ == "__main__":
    main()