"""
QuantProdigy AI Agents - Free Version for Hugging Face Spaces
Multi-agent learning platform for CFA preparation using only free resources
"""

import streamlit as st
import json
import plotly.graph_objects as go
from typing import Dict, List, Optional
from datetime import datetime
import sys
import os

# Debug: Print current directory and files
print("Current directory:", os.getcwd())
print("Files in directory:", os.listdir('.'))

# Add current directory to path for Hugging Face Spaces
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Suppress warnings
import warnings
warnings.filterwarnings('ignore')

# Import agent system with error handling
try:
    # Check if agent files exist
    if os.path.exists('agent_orchestrator.py'):
        from agent_orchestrator import AgentOrchestrator
    else:
        print("agent_orchestrator.py not found")
        AgentOrchestrator = None
        
    if os.path.exists('knowledge_engine.py'):
        from knowledge_engine import KnowledgeEngine
    else:
        print("knowledge_engine.py not found")
        KnowledgeEngine = None
        
    if os.path.exists('monitoring_system.py'):
        from monitoring_system import MonitoringSystem
    else:
        print("monitoring_system.py not found")
        MonitoringSystem = None
        
    if os.path.exists('runtime_environment.py'):
        from runtime_environment import RuntimeEnvironment
    else:
        print("runtime_environment.py not found")
        RuntimeEnvironment = None
        
except ImportError as e:
    print(f"Import error details: {e}")
    # Create dummy classes for demo mode
    class DummyOrchestrator:
        def __init__(self):
            pass
        def initialize_agents(self):
            return True
        def process_query(self, query, user_profile, selected_agents):
            return {
                'response': f"This is a demo response to: {query}\n\nIn a full deployment, multiple AI agents would analyze this question and provide comprehensive answers.",
                'primary_agent': 'Demo Agent'
            }
    
    class DummyKnowledgeEngine:
        def __init__(self):
            pass
        def retrieve_knowledge(self, query):
            return {
                'results': [
                    {
                        'type': 'concept',
                        'content': {
                            'name': 'Time Value of Money',
                            'definition': 'The concept that money available now is worth more than the same amount in the future due to its potential earning capacity.',
                            'formula': 'FV = PV × (1 + r)ⁿ'
                        },
                        'topics': ['Quantitative Methods']
                    }
                ]
            }
    
    class DummyMonitoring:
        def __init__(self):
            pass
        def get_system_metrics(self):
            return {
                'active_agents': 6,
                'avg_response_time': 1.2,
                'success_rate': 0.95,
                'total_queries': 0
            }
        def update_metrics(self):
            pass
    
    class DummyRuntime:
        def __init__(self):
            pass
    
    AgentOrchestrator = DummyOrchestrator
    KnowledgeEngine = DummyKnowledgeEngine
    MonitoringSystem = DummyMonitoring
    RuntimeEnvironment = DummyRuntime

# Page configuration for Hugging Face Spaces
st.set_page_config(
    page_title="QuantProdigy AI Agents",
    page_icon="📈",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize session state
def init_session_state():
    """Initialize session state with safe defaults"""
    defaults = {
        'orchestrator': None,
        'knowledge_engine': None,
        'monitoring': None,
        'runtime': None,
        'chat_history': [],
        'user_profile': {
            'cfa_level': 'I',
            'study_focus': ['Quantitative Methods', 'Economics'],
            'weak_areas': [],
            'learning_style': 'visual',
            'experience_level': 'beginner',
            'user_id': 'user_' + datetime.now().strftime('%Y%m%d_%H%M%S')
        },
        'active_agents': {},
        'system_ready': False,
        'demo_mode': True  # Added for Hugging Face Spaces
    }
    
    for key, value in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = value
    
    # Initialize components safely
    if st.session_state.orchestrator is None:
        try:
            st.session_state.orchestrator = AgentOrchestrator()
        except:
            st.session_state.orchestrator = AgentOrchestrator() if callable(AgentOrchestrator) else None
    
    if st.session_state.knowledge_engine is None:
        try:
            st.session_state.knowledge_engine = KnowledgeEngine()
        except:
            st.session_state.knowledge_engine = KnowledgeEngine() if callable(KnowledgeEngine) else None
    
    if st.session_state.monitoring is None:
        try:
            st.session_state.monitoring = MonitoringSystem()
        except:
            st.session_state.monitoring = MonitoringSystem() if callable(MonitoringSystem) else None
    
    if st.session_state.runtime is None:
        try:
            st.session_state.runtime = RuntimeEnvironment()
        except:
            st.session_state.runtime = RuntimeEnvironment() if callable(RuntimeEnvironment) else None

# Initialize session state
init_session_state()

# Custom CSS for Hugging Face Spaces
st.markdown("""
<style>
    /* Main styling */
    .main-header {
        font-size: 2.5rem;
        color: #1E3A8A;
        text-align: center;
        margin-bottom: 1rem;
        padding: 1rem;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
    }
    
    /* Card styling */
    .agent-card {
        background: white;
        padding: 1.5rem;
        border-radius: 15px;
        margin: 10px 0;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        border-left: 5px solid #667eea;
        transition: transform 0.3s ease;
    }
    
    /* Message styling */
    .message-user {
        background-color: #E3F2FD;
        padding: 15px;
        border-radius: 15px 15px 15px 5px;
        margin: 10px 0;
        text-align: right;
        max-width: 80%;
        margin-left: auto;
        border: 1px solid #BBDEFB;
    }
    
    .message-agent {
        background-color: #F3E5F5;
        padding: 15px;
        border-radius: 15px 15px 5px 15px;
        margin: 10px 0;
        max-width: 80%;
        border: 1px solid #E1BEE7;
    }
    
    /* Button styling */
    .stButton button {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border: none;
        padding: 0.5rem 2rem;
        border-radius: 25px;
        font-weight: bold;
        transition: all 0.3s ease;
    }
    
    .stButton button:hover {
        transform: translateY(-2px);
        box-shadow: 0 5px 15px rgba(102, 126, 234, 0.4);
    }
    
    /* Demo banner */
    .demo-banner {
        background: linear-gradient(135deg, #ff6b6b 0%, #ff8e53 100%);
        color: white;
        padding: 10px;
        border-radius: 10px;
        text-align: center;
        margin-bottom: 20px;
    }
</style>
""", unsafe_allow_html=True)

# Sidebar
with st.sidebar:
    st.title("📊 QuantProdigy")
    
    # Demo mode indicator
    if st.session_state.demo_mode:
        st.markdown('<div class="demo-banner">DEMO MODE</div>', unsafe_allow_html=True)
    
    st.markdown("### Free CFA Learning Platform")
    st.markdown("---")
    
    # System Status
    st.subheader("🖥️ System Status")
    
    if st.session_state.system_ready:
        st.success("✅ System Ready")
    else:
        if st.button("Initialize System", type="primary", use_container_width=True):
            try:
                with st.spinner("Initializing AI Agents..."):
                    if st.session_state.orchestrator:
                        st.session_state.orchestrator.initialize_agents()
                    st.session_state.system_ready = True
                    st.success("System initialized successfully!")
                    st.rerun()
            except Exception as e:
                st.error(f"Initialization failed: {e}")
    
    # User profile
    st.subheader("👤 Your Profile")
    
    with st.expander("Edit Profile", expanded=True):
        cfa_level = st.selectbox(
            "CFA Level",
            ["I", "II", "III"],
            index=0
        )
        
        learning_focus = st.multiselect(
            "Learning Focus",
            ["Quantitative Methods", "Economics", "Financial Reporting", 
             "Corporate Finance", "Equity", "Fixed Income", "Derivatives", 
             "Alternative Investments", "Portfolio Management", "Ethics"],
            default=["Quantitative Methods", "Economics"]
        )
        
        learning_style = st.selectbox(
            "Learning Style",
            ["Visual", "Auditory", "Reading/Writing", "Kinesthetic"],
            index=0
        )
        
        experience_level = st.selectbox(
            "Experience Level",
            ["Beginner", "Intermediate", "Advanced"],
            index=0
        )
        
        # Update user profile
        st.session_state.user_profile.update({
            'cfa_level': cfa_level,
            'study_focus': learning_focus,
            'learning_style': learning_style.lower(),
            'experience_level': experience_level.lower()
        })
    
    # Quick actions
    st.subheader("⚡ Quick Actions")
    
    quick_topics = ["Time Value of Money", "Statistics", "Portfolio Theory", "Ethics"]
    for topic in quick_topics:
        if st.button(f"Learn: {topic}", use_container_width=True):
            st.session_state.chat_history.append({
                'role': 'user',
                'content': f"Teach me about {topic}",
                'timestamp': datetime.now().isoformat()
            })
            st.rerun()

# Main content
if not st.session_state.system_ready:
    # Welcome screen
    st.markdown('<h1 class="main-header">Welcome to QuantProdigy AI Agents</h1>', unsafe_allow_html=True)
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.markdown("""
        ## 🚀 Free CFA Learning Platform
        
        **QuantProdigy** is an intelligent multi-agent system designed to help you master 
        CFA concepts using AI-powered assistance - completely free!
        
        ### ✨ Features:
        
        - **🤖 6 Specialized AI Agents** working together
        - **📚 Comprehensive CFA Knowledge Base**
        - **🧮 Built-in Financial Calculators**
        - **📊 Progress Tracking & Analytics**
        - **💬 Interactive Learning Sessions**
        """)
        
        if st.session_state.demo_mode:
            st.info("""
            ⚠️ **Demo Mode Active**  
            This is a lightweight demo version for Hugging Face Spaces. 
            Full features require local deployment with additional AI models.
            """)
    
    with col2:
        st.markdown("""
        ### 🏆 Get Started
        
        Click the button below to initialize the AI agents and start learning:
        """)
        
        if st.button("🚀 Initialize & Start Learning", type="primary", use_container_width=True):
            try:
                with st.spinner("🤖 Loading AI Agents..."):
                    # Initialize all components
                    if st.session_state.orchestrator:
                        st.session_state.orchestrator.initialize_agents()
                    st.session_state.system_ready = True
                st.success("✅ System Ready!")
                st.rerun()
            except Exception as e:
                st.error(f"Initialization failed: {e}")
    
    # Feature showcase
    st.markdown("---")
    st.subheader("🎯 Featured Learning Agents")
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("""
        <div class="agent-card">
            <h4>🧠 Diagnostician</h4>
            <p>Analyzes your learning gaps and identifies weak areas</p>
            <small>Ready to use</small>
        </div>
        """, unsafe_allow_html=True)
    
    with col2:
        st.markdown("""
        <div class="agent-card">
            <h4>👨‍🏫 Socratic Tutor</h4>
            <p>Guides you through questions to build understanding</p>
            <small>Ready to use</small>
        </div>
        """, unsafe_allow_html=True)
    
    with col3:
        st.markdown("""
        <div class="agent-card">
            <h4>🧮 Calculator</h4>
            <p>Performs complex financial calculations instantly</p>
            <small>Ready to use</small>
        </div>
        """, unsafe_allow_html=True)
    
    st.stop()

# Main application (system is ready)
st.markdown('<h1 class="main-header">QuantProdigy AI Learning Agents</h1>', unsafe_allow_html=True)

if st.session_state.demo_mode:
    st.info("""
    💡 **Demo Mode**: This version runs with simplified AI agents. 
    For full functionality with local LLMs, download and run locally.
    """)

# Tab layout
tab1, tab2, tab3, tab4, tab5 = st.tabs([
    "🎓 Learning Session", 
    "🤖 Agent Dashboard", 
    "📚 Knowledge Base", 
    "📊 Progress Analytics",
    "🛠️ Tools & Calculators"
])

# Tab 1: Learning Session
with tab1:
    st.header("Interactive Learning Session")
    
    # Chat interface
    col1, col2 = st.columns([3, 1])
    
    with col1:
        user_input = st.text_area(
            "Type your question here:",
            height=100,
            placeholder="E.g., 'Explain the concept of beta in finance' or 'Help me calculate portfolio variance...'"
        )
    
    with col2:
        if st.button("Ask AI Agents", type="primary", use_container_width=True):
            if user_input.strip():
                # Add user message to history
                st.session_state.chat_history.append({
                    'role': 'user',
                    'content': user_input,
                    'timestamp': datetime.now().isoformat()
                })
                
                # Process with orchestrator
                if st.session_state.orchestrator:
                    with st.spinner("🤖 AI agents are analyzing your question..."):
                        try:
                            response = st.session_state.orchestrator.process_query(
                                query=user_input,
                                user_profile=st.session_state.user_profile,
                                selected_agents=["Diagnostician", "Knowledge Expert"]
                            )
                            
                            # Add agent response to history
                            st.session_state.chat_history.append({
                                'role': 'agent',
                                'content': response['response'],
                                'agent': response.get('primary_agent', 'Orchestrator'),
                                'timestamp': datetime.now().isoformat()
                            })
                            
                            st.rerun()
                            
                        except Exception as e:
                            st.error(f"Error processing query: {e}")
                            # Add fallback response
                            st.session_state.chat_history.append({
                                'role': 'agent',
                                'content': f"""I'll help you with: "{user_input}"
**Time Value of Money Example:**
If you invest $1,000 at 5% annual interest for 5 years:
- Future Value = $1,000 × (1 + 0.05)⁵ = $1,276.28
**Key Concept:** Money today is worth more than the same amount in the future due to earning potential.""",
                                'agent': 'Demo Agent',
                                'timestamp': datetime.now().isoformat()
                            })
                            st.rerun()
                else:
                    # Demo response
                    st.session_state.chat_history.append({
                        'role': 'agent',
                        'content': f"""Thanks for your question about "{user_input}"!
**In a full deployment**, multiple AI agents would:
1. Analyze your learning level and style
2. Retrieve relevant CFA knowledge
3. Provide step-by-step explanations
4. Generate practice problems
5. Track your progress
**For now, here's a helpful insight:**
In finance, understanding core concepts like risk-return tradeoff, time value of money, and diversification is fundamental to mastering CFA material.""",
                        'agent': 'Demo Agent',
                        'timestamp': datetime.now().isoformat()
                    })
                    st.rerun()
    
    # Display chat history
    st.subheader("📝 Conversation History")
    
    if not st.session_state.chat_history:
        st.info("👆 Start by asking a question above!")
    else:
        for message in st.session_state.chat_history[-10:]:
            if message['role'] == 'user':
                st.markdown(f"""
                <div class="message-user">
                    <strong>You:</strong> {message['content']}
                    <br><small>{message.get('timestamp', '')[:16]}</small>
                </div>
                """, unsafe_allow_html=True)
            else:
                st.markdown(f"""
                <div class="message-agent">
                    <strong>{message.get('agent', 'AI Agent')}:</strong><br>
                    {message['content']}
                    <br><small>{message.get('timestamp', '')[:16]}</small>
                </div>
                """, unsafe_allow_html=True)

# Tab 2: Agent Dashboard
with tab2:
    st.header("Agent Dashboard")
    
    # Agent status grid
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("""
        <div class="agent-card">
            <h4>🧠 Diagnostician</h4>
            <p>Analyzes learning gaps and problems</p>
            <small>Active in Demo</small>
        </div>
        """, unsafe_allow_html=True)
    
    with col2:
        st.markdown("""
        <div class="agent-card">
            <h4>👨‍🏫 Socratic Tutor</h4>
            <p>Guides through questions and reasoning</p>
            <small>Active in Demo</small>
        </div>
        """, unsafe_allow_html=True)
    
    with col3:
        st.markdown("""
        <div class="agent-card">
            <h4>🧮 Calculator</h4>
            <p>Performs complex calculations</p>
            <small>Active in Demo</small>
        </div>
        """, unsafe_allow_html=True)

# Tab 3: Knowledge Base
with tab3:
    st.header("CFA Knowledge Base")
    
    # Sample knowledge
    st.subheader("📖 Sample CFA Concepts")
    
    concepts = [
        {
            "name": "Time Value of Money",
            "definition": "Money available now is worth more than the same amount in the future due to its potential earning capacity.",
            "formula": "FV = PV × (1 + r)ⁿ"
        },
        {
            "name": "Sharpe Ratio",
            "definition": "Measures risk-adjusted return of an investment portfolio.",
            "formula": "Sharpe Ratio = (Rp - Rf) / σp"
        },
        {
            "name": "Beta Coefficient",
            "definition": "Measures the volatility of an investment relative to the market.",
            "formula": "β = Cov(Ri, Rm) / Var(Rm)"
        }
    ]
    
    for concept in concepts:
        with st.expander(f"📚 {concept['name']}"):
            st.markdown(f"**Definition:** {concept['definition']}")
            st.markdown(f"**Formula:** `{concept['formula']}`")
            if st.button(f"Learn more about {concept['name']}", key=f"learn_{concept['name']}"):
                st.session_state.chat_history.append({
                    'role': 'user',
                    'content': f"Explain {concept['name']} in detail",
                    'timestamp': datetime.now().isoformat()
                })
                st.rerun()

# Tab 4: Progress Analytics
with tab4:
    st.header("Learning Progress Analytics")
    
    # Sample progress chart
    st.subheader("📈 Sample Progress Data")
    
    fig = go.Figure(data=[
        go.Bar(name='Topics Studied', x=['Week 1', 'Week 2', 'Week 3'], y=[3, 5, 8]),
        go.Bar(name='Practice Problems', x=['Week 1', 'Week 2', 'Week 3'], y=[5, 12, 20])
    ])
    
    fig.update_layout(barmode='group', title='Learning Progress Over Time')
    st.plotly_chart(fig, use_container_width=True)

# Tab 5: Tools & Calculators
with tab5:
    st.header("🛠️ Financial Calculators")
    
    # Simple TVM Calculator
    st.subheader("💰 Time Value of Money Calculator")
    
    col1, col2 = st.columns(2)
    
    with col1:
        pv = st.number_input("Present Value ($)", value=1000.0, min_value=0.0, step=100.0)
        rate = st.number_input("Annual Rate (%)", value=5.0, min_value=0.0, max_value=100.0, step=0.5)
        periods = st.number_input("Periods (years)", value=5, min_value=1, max_value=100, step=1)
    
    with col2:
        if st.button("Calculate Future Value", type="primary", use_container_width=True):
            fv = pv * (1 + rate/100) ** periods
            st.success(f"**Future Value:** ${fv:,.2f}")
            st.info(f"${pv:,.2f} at {rate}% for {periods} years grows to ${fv:,.2f}")

# Footer
st.markdown("---")
st.markdown("""
<div style="text-align: center; padding: 20px;">
    <p><strong>QuantProdigy AI Agents</strong> • Free CFA Learning Platform</p>
    <small>Demo Version for Hugging Face Spaces • For full features, run locally</small>
</div>
""", unsafe_allow_html=True)