"""Smart Analytics Copilot - Complete Version
With Export, OpenAI, Save/Load, Chart Customization, Power BI Export"""

import streamlit as st
import pandas as pd
import os
from datetime import datetime
from dotenv import load_dotenv

#Load environment variables
load_dotenv()

from data_processor import DataProcessor
from analyzer import Analyzer
from insight_generator import InsightGenerator
from dashboard import DashboardGenerator
from query_engine import QueryEngine
from export_utils import ExportUtils
from session_manager import SessionManager
from chart_customizer import ChartCustomizer

#Page config
st.set_page_config(
    page_title="Smart Analytics Copilot",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ============ DARK THEME CSS ============
st.markdown("""
<style>
    /* Main background */
    .stApp {
        background-color: #0a0e17 !important;
    }
    
    /* All text - light color */
    .stMarkdown, .stMarkdown p, .stMarkdown div, .stMarkdown span, 
    .stText, p, div, span, label {
        color: #e8e8e8 !important;
    }
    
    /* Headers */
    h1, h2, h3, h4, h5, h6 {
        color: #00ff9d !important;
        font-weight: 600 !important;
    }
    
    /* Main header */
    .main-header {
        font-size: 2.8rem;
        font-weight: bold;
        background: linear-gradient(135deg, #00ff9d 0%, #00d4ff 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        margin-bottom: 1rem;
        text-align: center;
    }
    
    /* Sidebar */
    .css-1d391kg, .stSidebar, .sidebar-content {
        background-color: #111827 !important;
    }
    
    /* Metrics */
    div[data-testid="stMetricValue"] {
        color: #00ff9d !important;
        font-size: 2rem !important;
        font-weight: bold !important;
    }
    
    div[data-testid="stMetricLabel"] {
        color: #a0aec0 !important;
        font-size: 0.9rem !important;
    }
    
    /* Tabs */
    .stTabs [data-baseweb="tab-list"] {
        gap: 4px;
        background-color: #111827;
        border-radius: 10px;
        padding: 6px;
    }
    
    .stTabs [data-baseweb="tab"] {
        background-color: #1f2937;
        border-radius: 8px;
        padding: 8px 24px;
        color: #e8e8e8 !important;
    }
    
    .stTabs [aria-selected="true"] {
        background: linear-gradient(135deg, #00ff9d 0%, #00d4ff 100%) !important;
        color: #0a0e17 !important;
        font-weight: bold;
    }
    
    /* Buttons */
    .stButton button {
        background: linear-gradient(135deg, #00ff9d 0%, #00d4ff 100%) !important;
        color: #0a0e17 !important;
        font-weight: bold !important;
        border: none !important;
        border-radius: 8px !important;
    }
    
    /* File uploader */
    .stFileUploader {
        background-color: #1f2937 !important;
        border: 2px dashed #374151 !important;
        border-radius: 12px !important;
    }
    
    /* Expander */
    .streamlit-expanderHeader {
        background-color: #1f2937 !important;
        color: #00ff9d !important;
        border-radius: 8px;
    }
    
    /* Success/Info/Warning boxes */
    .stAlert {
        background-color: #1f2937 !important;
        border: 1px solid #374151 !important;
        border-radius: 10px !important;
    }
    
    .stAlert p, .stAlert div {
        color: #e8e8e8 !important;
    }
    
    /* Dataframe */
    .stDataFrame {
        background-color: #111827 !important;
    }
    
    .stDataFrame thead th {
        background-color: #1f2937 !important;
        color: #00ff9d !important;
    }
    
    /* Text input */
    .stTextInput input {
        background-color: #1f2937 !important;
        color: #e8e8e8 !important;
        border: 1px solid #374151 !important;
        border-radius: 8px !important;
    }
    
    /* Select box */
    .stSelectbox div[data-baseweb="select"] {
        background-color: #1f2937 !important;
        border-color: #374151 !important;
    }
    
    /* Download button */
    .stDownloadButton button {
        background: linear-gradient(135deg, #00ff9d 0%, #00d4ff 100%) !important;
        color: #0a0e17 !important;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'data_loaded' not in st.session_state:
    st.session_state.data_loaded = False
if 'df' not in st.session_state:
    st.session_state.df = None
if 'schema' not in st.session_state:
    st.session_state.schema = None
if 'analysis' not in st.session_state:
    st.session_state.analysis = None
if 'insights' not in st.session_state:
    st.session_state.insights = None
if 'charts' not in st.session_state:
    st.session_state.charts = None
if 'use_openai' not in st.session_state:
    st.session_state.use_openai = False

# Initialize managers
session_mgr = SessionManager()


def main():
    st.markdown('<div class="main-header">🚀 Smart Analytics Copilot</div>', unsafe_allow_html=True)
    st.caption("✨ Upload any CSV/JSON - AI analyzes, visualizes, and answers questions")
    st.markdown("---")
    
    # Sidebar
    with st.sidebar:
        st.markdown("### 📁 Data Source")
        
        # Data source selection
        source = st.radio("Choose data source:", ["📤 Upload File", "💾 Load Saved Session"])
        
        if source == "📤 Upload File":
            uploaded_file = st.file_uploader("Choose CSV or JSON", type=['csv', 'json'])
            if uploaded_file and not st.session_state.data_loaded:
                with st.spinner("🔄 Processing your data..."):
                    process_data(uploaded_file)
        else:
            # Load saved sessions
            sessions = session_mgr.list_sessions()
            if sessions:
                session_names = [s['name'] for s in sessions]
                selected_session = st.selectbox("Select saved session:", session_names)
                if st.button("📂 Load Session"):
                    with st.spinner("Loading..."):
                        load_session(selected_session)
            else:
                st.info("No saved sessions found")
        
        st.markdown("---")
        
        # Settings
        with st.expander("⚙️ Settings"):
            st.session_state.use_openai = st.checkbox("Use OpenAI (better insights)", 
                                                      value=st.session_state.use_openai)
            if st.session_state.use_openai:
                api_key = st.text_input("OpenAI API Key:", type="password")
                if api_key:
                    os.environ['OPENAI_API_KEY'] = api_key
                    st.success("API Key set!")
        
        st.markdown("---")
        
        # Export section (only if data loaded)
        if st.session_state.data_loaded:
            st.markdown("### 💾 Export Options")
            export_utils = ExportUtils(st.session_state.df)
            
            export_format = st.selectbox("Export format:", 
                                         ["CSV", "Excel", "JSON", "Power BI CSV", "Power BI ZIP (Complete)"])
            
            if st.button("📥 Download"):
                if export_format == "CSV":
                    data = export_utils.to_csv()
                    mime = "text/csv"
                    ext = "csv"
                elif export_format == "Excel":
                    data = export_utils.to_excel()
                    mime = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                    ext = "xlsx"
                elif export_format == "JSON":
                    data = export_utils.to_json()
                    mime = "application/json"
                    ext = "json"
                elif export_format == "Power BI CSV":
                    data = export_utils.to_powerbi_ready()
                    mime = "text/csv"
                    ext = "csv"
                else:  # Power BI ZIP (Complete)
                    data = export_utils.to_powerbi_zip()
                    mime = "application/zip"
                    ext = "zip"
                
                st.download_button(
                    label="✅ Click to Download",
                    data=data,
                    file_name=f"export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.{ext}",
                    mime=mime
                )
            
            # Save session button
            st.markdown("---")
            if st.button("💾 Save Current Session"):
                name, path = session_mgr.save_session(st.session_state.df, st.session_state.schema)
                st.success(f"✅ Session saved as: {name}")
    
    # Main content
    if st.session_state.data_loaded:
        tab1, tab2, tab3, tab4, tab5 = st.tabs([
            "📊 Dashboard", "💡 AI Insights", "🎨 Custom Charts", "🔍 Query", "📋 Data"
        ])
        
        with tab1:
            show_dashboard()
        
        with tab2:
            show_insights()
        
        with tab3:
            show_chart_customizer()
        
        with tab4:
            show_query_interface()
        
        with tab5:
            show_data_preview()
    
    else:
        show_welcome()


def process_data(uploaded_file):
    """Process uploaded data"""
    try:
        processor = DataProcessor()
        st.session_state.df = processor.load_from_upload(uploaded_file)
        st.session_state.df = processor.preprocess()
        st.session_state.schema = processor.detect_schema()
        
        analyzer = Analyzer(st.session_state.df, st.session_state.schema)
        st.session_state.analysis = analyzer.run_full_analysis()
        
        # Use OpenAI if enabled
        api_key = os.environ.get('OPENAI_API_KEY')
        insight_gen = InsightGenerator(use_openai=st.session_state.use_openai, api_key=api_key)
        st.session_state.insights = insight_gen.generate_insights(
            st.session_state.df, 
            st.session_state.schema, 
            st.session_state.analysis
        )
        
        dashboard_gen = DashboardGenerator(st.session_state.df, st.session_state.schema)
        st.session_state.charts = dashboard_gen.generate_all_charts()
        
        st.session_state.data_loaded = True
        st.success(f"✅ Successfully loaded {len(st.session_state.df):,} rows with {len(st.session_state.df.columns)} columns")
        st.balloons()
        st.rerun()
    except Exception as e:
        st.error(f"Error: {e}")


def load_session(session_name):
    """Load saved session and regenerate insights"""
    session = session_mgr.load_session(session_name)
    if session:
        st.session_state.df = session['df']
        st.session_state.schema = session['schema']
        
        # Regenerate analysis and insights for loaded session
        with st.spinner("🔄 Regenerating analysis..."):
            analyzer = Analyzer(st.session_state.df, st.session_state.schema)
            st.session_state.analysis = analyzer.run_full_analysis()
            
            # Regenerate insights
            api_key = os.environ.get('OPENAI_API_KEY')
            insight_gen = InsightGenerator(use_openai=st.session_state.use_openai, api_key=api_key)
            st.session_state.insights = insight_gen.generate_insights(
                st.session_state.df, 
                st.session_state.schema, 
                st.session_state.analysis
            )
            
            # Regenerate charts
            dashboard_gen = DashboardGenerator(st.session_state.df, st.session_state.schema)
            st.session_state.charts = dashboard_gen.generate_all_charts()
        
        st.session_state.data_loaded = True
        st.success(f"✅ Loaded session: {session_name}")
        st.rerun()
    else:
        st.error("Failed to load session")


def show_dashboard():
    """Display dashboard"""
    st.markdown("### 📈 Key Metrics")
    st.markdown("---")
    
    # Check if data exists
    if st.session_state.df is None:
        st.warning("No data loaded. Please upload a file first.")
        return
    
    # Display metrics
    if st.session_state.schema['numeric']:
        cols = st.columns(min(4, len(st.session_state.schema['numeric'])))
        for idx, col in enumerate(st.session_state.schema['numeric'][:4]):
            with cols[idx]:
                total = st.session_state.df[col].sum()
                avg = st.session_state.df[col].mean()
                st.metric(
                    label=f"💰 {col.upper()}",
                    value=f"{total:,.0f}",
                    delta=f"Avg: {avg:,.0f}"
                )
    
    st.markdown("---")
    st.markdown("### 📊 Visualizations")
    
    if st.session_state.charts:
        for chart in st.session_state.charts[:4]:
            st.plotly_chart(chart['figure'], use_container_width=True)
    else:
        st.info("No charts available. Try uploading data first.")
    
    st.markdown("---")
    st.markdown("### 📋 Summary Statistics")
    if st.session_state.schema['numeric']:
        summary = st.session_state.df[st.session_state.schema['numeric']].describe()
        st.dataframe(summary, use_container_width=True)


def show_insights():
    """Display AI insights"""
    st.markdown("### 🧠 AI-Powered Insights")
    st.markdown("Here's what we discovered in your data:")
    st.markdown("---")
    
    # Check if insights exist
    if st.session_state.insights is None:
        st.info("💡 Insights will appear after data is analyzed.")
        return
    
    for insight in st.session_state.insights:
        if "Dataset" in insight:
            st.info(f"📊 {insight}")
        elif "correlation" in insight.lower():
            st.success(f"✅ {insight}")
        elif "skewed" in insight.lower():
            st.warning(f"📈 {insight}")
        elif "Recommendation" in insight:
            st.info(f"💡 {insight}")
        else:
            st.markdown(f"• {insight}")
    
    # Power BI template section
    st.markdown("---")
    with st.expander("📊 Power BI Resources"):
        export_utils = ExportUtils(st.session_state.df)
        
        col1, col2 = st.columns(2)
        
        with col1:
            # Show DAX template
            template = export_utils.get_powerbi_template()
            st.code(template, language="dax")
            
            st.download_button(
                label="📥 Download DAX Template",
                data=template,
                file_name="powerbi_measures.dax",
                mime="text/plain"
            )
        
        with col2:
            # Show instructions
            instructions = """
**Power BI Import Steps:**

1. **Export Data**: Use sidebar to export as "Power BI CSV"
2. **Open Power BI Desktop**
3. **Get Data** → **Text/CSV**
4. **Select your exported CSV**
5. **Click Load**
6. **Copy DAX measures** from above
7. **Create visuals** using the measures
            """
            st.info(instructions)


def show_chart_customizer():
    """Show chart customization interface"""
    st.markdown("### 🎨 Custom Chart Builder")
    st.markdown("Create your own custom visualizations")
    st.markdown("---")
    
    customizer = ChartCustomizer(st.session_state.df)
    available_charts = customizer.get_available_charts()
    
    col1, col2, col3 = st.columns([1, 1, 1])
    
    with col1:
        chart_type = st.selectbox("Chart Type:", available_charts)
    
    with col2:
        # Get appropriate columns
        if 'Histogram' in chart_type or 'Box' in chart_type:
            columns = st.session_state.schema['numeric']
            if not columns:
                columns = list(st.session_state.df.select_dtypes(include=['number']).columns)
        elif 'Pie' in chart_type or 'Bar' in chart_type:
            columns = st.session_state.schema['categorical']
            if not columns:
                columns = list(st.session_state.df.select_dtypes(include=['object']).columns)
        else:
            columns = list(st.session_state.df.columns)
        
        if columns:
            x_col = st.selectbox("X-Axis / Category:", columns)
        else:
            x_col = None
            st.warning("No suitable columns found")
    
    with col3:
        # For charts that need Y-axis
        if any(t in chart_type for t in ['Line', 'Scatter', 'Bar']) and 'Histogram' not in chart_type:
            y_cols = ['None'] + st.session_state.schema['numeric']
            y_col = st.selectbox("Y-Axis / Value:", y_cols)
            y_col = None if y_col == 'None' else y_col
        else:
            y_col = None
    
    # Color column (optional)
    color_cols = ['None'] + st.session_state.schema['categorical']
    color_col = st.selectbox("Color By (optional):", color_cols)
    color_col = None if color_col == 'None' else color_col
    
    # Title
    title = st.text_input("Chart Title:", value=f"{chart_type} of {x_col if x_col else 'data'}")
    
    if st.button("🎨 Generate Chart", use_container_width=True):
        if x_col:
            with st.spinner("Creating chart..."):
                fig = customizer.create_chart(chart_type, x_col, y_col, color_col, title)
                if fig:
                    st.plotly_chart(fig, use_container_width=True)
                    
                    # Download chart button
                    try:
                        st.download_button(
                            label="📸 Download as PNG",
                            data=fig.to_image(format="png"),
                            file_name="custom_chart.png",
                            mime="image/png"
                        )
                    except:
                        st.info("💡 Install kaleido for PNG export: `pip install kaleido`")
                else:
                    st.error("Could not create chart. Try different settings.")
        else:
            st.error("Please select a column for X-Axis")


def show_query_interface():
    """Natural language query interface"""
    st.markdown("### 💬 Natural Language Query")
    st.markdown("Ask any question about your data in plain English:")
    st.markdown("---")
    
    query_engine = QueryEngine(st.session_state.df, st.session_state.schema)
    
    # Example questions
    with st.expander("🔍 View Example Questions"):
        if st.session_state.schema['numeric']:
            example_col = st.session_state.schema['numeric'][0]
            st.markdown(f"• 'Statistics {example_col}'")
            st.markdown(f"• 'Total {example_col}'")
            st.markdown(f"• 'Average {example_col}'")
        
        if st.session_state.schema['categorical'] and st.session_state.schema['numeric']:
            st.markdown(f"• 'Top 5 {st.session_state.schema['categorical'][0]} by {st.session_state.schema['numeric'][0]}'")
        
        st.markdown("• 'Summary statistics'")
        st.markdown("• 'Show me the data'")
    
    st.markdown("---")
    
    question = st.text_input("Ask a question:", placeholder="e.g., What is the average of time_in_hospital?")
    
    if question:
        with st.spinner("🤔 Analyzing your question..."):
            answer = query_engine.answer_question(question)
            st.markdown("### ✅ Answer")
            st.success(answer)


def show_data_preview():
    """Show data preview and info with better formatting"""
    st.markdown("### 📋 Data Preview")
    st.markdown("---")
    
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("📊 Total Rows", f"{len(st.session_state.df):,}")
    with col2:
        st.metric("📋 Total Columns", len(st.session_state.df.columns))
    with col3:
        memory = st.session_state.df.memory_usage(deep=True).sum() / 1024**2
        st.metric("💾 Memory Usage", f"{memory:.2f} MB")
    
    st.markdown("---")
    st.markdown("### 📄 Data Sample (First 100 rows)")
    
    # Create a copy for display
    display_df = st.session_state.df.head(100).copy()
    
    # Clean datetime columns for better display
    for col in display_df.columns:
        if 'datetime' in col.lower() or 'date' in col.lower() or 'time' in col.lower():
            try:
                display_df[col] = pd.to_datetime(display_df[col]).dt.strftime('%Y-%m-%d %H:%M:%S')
            except:
                pass
    
    st.dataframe(display_df, use_container_width=True)
    
    st.markdown("---")
    st.markdown("### 📊 Column Information")
    
    col_info = pd.DataFrame({
        'Column': st.session_state.df.columns,
        'Type': st.session_state.df.dtypes.astype(str),
        'Non-Null': st.session_state.df.count().values,
        'Nulls': st.session_state.df.isnull().sum().values,
        'Unique': st.session_state.df.nunique().values
    })
    st.dataframe(col_info, use_container_width=True)


def show_welcome():
    """Welcome screen"""
    st.markdown("""
    <div style="text-align: center; padding: 2rem; background: linear-gradient(135deg, #111827 0%, #0a0e17 100%); border-radius: 20px; margin: 2rem 0;">
        <h2 style="color: #00ff9d;">🚀 Welcome to Smart Analytics Copilot</h2>
        <p style="font-size: 1.1rem;">Upload any CSV or JSON file and let AI analyze it instantly</p>
        <hr>
        <p>👈 <strong>Get Started</strong>: Upload a file or load a saved session from the sidebar</p>
    </div>
    """, unsafe_allow_html=True)
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("""
        <div style="background: linear-gradient(135deg, #1f2937 0%, #111827 100%); padding: 1.5rem; border-radius: 15px; text-align: center;">
            <h3 style="color: #00ff9d;">📊 Auto Dashboard</h3>
            <p>Smart charts based on your data</p>
        </div>
        """, unsafe_allow_html=True)
    
    with col2:
        st.markdown("""
        <div style="background: linear-gradient(135deg, #1f2937 0%, #111827 100%); padding: 1.5rem; border-radius: 15px; text-align: center;">
            <h3 style="color: #00ff9d;">💡 AI Insights</h3>
            <p>Natural language explanations</p>
        </div>
        """, unsafe_allow_html=True)
    
    with col3:
        st.markdown("""
        <div style="background: linear-gradient(135deg, #1f2937 0%, #111827 100%); padding: 1.5rem; border-radius: 15px; text-align: center;">
            <h3 style="color: #00ff9d;">🎨 Custom Charts</h3>
            <p>Build your own visualizations</p>
        </div>
        """, unsafe_allow_html=True)


if __name__ == "__main__":
    main()