import streamlit as st
from crewai import Agent, Crew, Task, Process
from typing import List
import os
from dotenv import load_dotenv
from crewai_tools import SerperDevTool
import json
from pydantic import BaseModel, Field
from typing import List, Optional, Dict
from enum import Enum
from langchain.llms import GoogleGenerativeAI

# Page configuration
st.set_page_config(
    page_title="Learning Path Generator",
    page_icon="🎓",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        font-weight: 700;
        color: #1E88E5;
        margin-bottom: 1rem;
    }
    .sub-header {
        font-size: 1.8rem;
        font-weight: 600;
        color: #333;
        margin-top: 2rem;
        margin-bottom: 1rem;
    }
    .card {
        background-color: #f9f9f9;
        border-radius: 10px;
        padding: 20px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        margin-bottom: 20px;
    }
    .material-card {
        background-color: white;
        border-left: 5px solid #4CAF50;
        padding: 15px;
        margin-bottom: 15px;
        border-radius: 5px;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
    }
    .video-card {
        border-left-color: #FF5722;
    }
    .article-card {
        border-left-color: #2196F3;
    }
    .exercise-card {
        border-left-color: #9C27B0;
    }
    .badge {
        display: inline-block;
        padding: 5px 10px;
        border-radius: 15px;
        font-size: 0.8rem;
        font-weight: 600;
        color: white;
        margin-right: 10px;
    }
    .badge-video {
        background-color: #FF5722;
    }
    .badge-article {
        background-color: #2196F3;
    }
    .badge-exercise {
        background-color: #9C27B0;
    }
    .badge-beginner {
        background-color: #4CAF50;
    }
    .badge-intermediate {
        background-color: #FF9800;
    }
    .badge-advanced {
        background-color: #F44336;
    }
    .quiz-question {
        background-color: white;
        border-radius: 8px;
        padding: 20px;
        margin-bottom: 20px;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
    }
    .quiz-option {
        padding: 10px;
        background-color: #f5f5f5;
        border-radius: 5px;
        margin-bottom: 10px;
        cursor: pointer;
    }
    .quiz-option-correct {
        background-color: #e8f5e9;
        border-left: 5px solid #4CAF50;
    }
    .project-card {
        background-color: white;
        border-radius: 8px;
        padding: 20px;
        margin-bottom: 20px;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
    }
    .project-header {
        display: flex;
        justify-content: space-between;
        align-items: center;
        margin-bottom: 15px;
    }
    .footer {
        text-align: center;
        padding: 20px;
        color: #666;
        font-size: 0.9rem;
    }
    .loading-container {
        display: flex;
        flex-direction: column;
        align-items: center;
        justify-content: center;
        padding: 50px;
    }
    .progress-bar {
        width: 100%;
        height: 20px;
        background-color: #f0f0f0;
        border-radius: 10px;
        overflow: hidden;
        margin-bottom: 20px;
    }
    .progress {
        height: 100%;
        background-color: #4CAF50;
        width: 0%;
        animation: progress 2s ease infinite;
    }
    @keyframes progress {
        0% { width: 0%; }
        50% { width: 100%; }
        100% { width: 0%; }
    }
    .gemini-badge {
        background-color: #8E24AA;
        color: white;
        padding: 5px 10px;
        border-radius: 15px;
        font-size: 0.8rem;
        font-weight: 600;
        display: inline-block;
    }
</style>
""", unsafe_allow_html=True)

# Load environment variables
load_dotenv()

class ExpertiseLevel(str, Enum):
    BEGINNER = "beginner"
    INTERMEDIATE = "intermediate"
    ADVANCED = "advanced"

# Model definitions
class LearningMaterial(BaseModel):
    title: str
    url: str
    type: str = Field(..., description="video, article, or exercise")
    description: str

class MaterialCollection(BaseModel):
    materials: List[LearningMaterial]

class QuizQuestion(BaseModel):
    question: str
    options: List[str]
    correct_answer: int
    explanation: str

class Quiz(BaseModel):
    questions: List[QuizQuestion]

class ProjectIdea(BaseModel):
    title: str
    description: str
    difficulty: ExpertiseLevel
    estimated_duration: str = Field(..., description="Duration estimation in days")
    required_skills: List[str]
    learning_outcomes: List[str]

class Projects(BaseModel):
    projects: List[ProjectIdea]

# Initialize LLM and search tool
def initialize_services():
    # Get API keys
    google_api_key = os.getenv("GOOGLE_API_KEY")
    serper_api_key = os.getenv("SERPER_API_KEY")
    
    if not google_api_key:
        st.error("Google API Key not found in environment variables. Please set the GOOGLE_API_KEY.")
        st.stop()
    
    if not serper_api_key:
        st.warning("Serper API Key not found. Search functionality may be limited.")
    
    try:
        # Initialize Gemini model
        from langchain_google_genai import ChatGoogleGenerativeAI
        
        gemini_llm = ChatGoogleGenerativeAI(
            model="gemini-2.0-flash-lite",
            google_api_key=google_api_key,
            temperature=0.7,
            convert_system_message_to_human=True
        )
        
        # Test the model connection
        _ = gemini_llm.invoke("Test connection")
        
        # Initialize search tool if API key is available
        search_tool = None
        if serper_api_key:
            search_tool = SerperDevTool(serper_api_key=serper_api_key)
        
        return gemini_llm, search_tool
    
    except ImportError:
        st.error("Required packages not installed. Please install langchain-google-genai.")
        st.stop()
    except Exception as e:
        st.error(f"Error initializing Gemini model: {str(e)}")
        # Fallback to a default model from CrewAI
        from crewai import LLM
        default_llm = LLM(name="openai", model="gpt-3.5-turbo")
        st.warning("Falling back to default model (OpenAI GPT-3.5).")
        return default_llm, None

def create_agents_and_tasks(topics, expertise_level, llm):
    # Create agents
    learning_material_agent = Agent(
        role='Learning Material Curator',
        goal='Curate high-quality learning materials based on user topics and expertise level',
        backstory="""You are an expert educational content curator with years of experience
        in finding the best learning resources for students at different levels. You know how
        to identify reliable and high-quality educational content from reputable sources.""",
        llm=llm,
        verbose=True
    )

    quiz_creator_agent = Agent(
        role='Quiz Creator',
        goal='Create engaging and educational quizzes to test understanding',
        backstory="""You are an experienced educator who specializes in creating
        effective assessment questions that test understanding while promoting learning.""",
        llm=llm,
        verbose=True
    )

    project_suggestion_agent = Agent(
        role='Project Advisor',
        goal='Suggest practical projects that match user expertise and interests',
        backstory="""You are a project-based learning expert who knows how to create
        engaging hands-on projects that reinforce learning objectives.""",
        llm=llm,
        verbose=True
    )

    # Create tasks
    create_learning_material_task = Task(
        description=f"""{topics}.
        Explain {topics} to a {expertise_level} level.
        Include a mix of videos, articles, and practical exercises.
        Ensure all materials are from reputable sources and are current.
        Include GitHub repos for practical exercises. Verify source credibility before including.
        Format response as: {{
            "materials": [
                {{
                    "title": "...",
                    "url": "...",
                    "type": "...",
                    "description": "..."
                }}
            ]
        }}""",
        agent=learning_material_agent,
        expected_output=MaterialCollection.schema_json()
    )

    create_quiz_task = Task(
        description=f"Create a comprehensive quiz for {topics} at {expertise_level} level.",
        agent=quiz_creator_agent,
        expected_output=Quiz.schema_json(),
        output_pydantic=Quiz
    )

    create_project_suggestion_task = Task(
        description=f"""Suggest ONLY 5 BEST practical project ideas for {topics}.
        Projects should be suitable for {expertise_level} level.
        Include title, description, difficulty, estimated duration, required skills, and learning outcomes.
        Suggest projects that have recent community activity (check GitHub).
        Include links to relevant documentation.
        Projects should be engaging and reinforce key concepts.""",
        agent=project_suggestion_agent,
        expected_output=Projects.schema_json(),
        output_pydantic=Projects
    )
    
    return (
        [learning_material_agent, quiz_creator_agent, project_suggestion_agent],
        [create_learning_material_task, create_quiz_task, create_project_suggestion_task]
    )

def display_learning_materials(materials):
    st.markdown("<div class='sub-header'>📚 Curated Learning Materials</div>", unsafe_allow_html=True)
    
    try:
        # Parse the raw JSON string into a Python dictionary
        materials_json = json.loads(materials)
        
        # Group materials by type
        videos = []
        articles = []
        exercises = []
        
        if 'materials' in materials_json:
            for material in materials_json['materials']:
                if material['type'].lower() == 'video':
                    videos.append(material)
                elif material['type'].lower() == 'article':
                    articles.append(material)
                elif material['type'].lower() == 'exercise':
                    exercises.append(material)
            
            # Display materials by type
            if videos:
                st.markdown("### 🎥 Videos")
                for material in videos:
                    st.markdown(f"""
                    <div class='material-card video-card'>
                        <div><span class='badge badge-video'>Video</span> <strong>{material['title']}</strong></div>
                        <div style='margin: 10px 0;'>{material['description']}</div>
                        <a href='{material['url']}' target='_blank'>Watch Video →</a>
                    </div>
                    """, unsafe_allow_html=True)
            
            if articles:
                st.markdown("### 📄 Articles")
                for material in articles:
                    st.markdown(f"""
                    <div class='material-card article-card'>
                        <div><span class='badge badge-article'>Article</span> <strong>{material['title']}</strong></div>
                        <div style='margin: 10px 0;'>{material['description']}</div>
                        <a href='{material['url']}' target='_blank'>Read Article →</a>
                    </div>
                    """, unsafe_allow_html=True)
            
            if exercises:
                st.markdown("### 💻 Exercises")
                for material in exercises:
                    st.markdown(f"""
                    <div class='material-card exercise-card'>
                        <div><span class='badge badge-exercise'>Exercise</span> <strong>{material['title']}</strong></div>
                        <div style='margin: 10px 0;'>{material['description']}</div>
                        <a href='{material['url']}' target='_blank'>Start Exercise →</a>
                    </div>
                    """, unsafe_allow_html=True)
    except json.JSONDecodeError as e:
        st.error(f"Error parsing learning materials: {e}")
        st.write(materials)  # Fallback to display raw output

def display_quiz(quiz):
    st.markdown("<div class='sub-header'>🧠 Knowledge Quiz</div>", unsafe_allow_html=True)
    
    try:
        quiz_json = json.loads(quiz)
        if 'questions' in quiz_json:
            for i, question in enumerate(quiz_json['questions'], 1):
                st.markdown(f"""
                <div class='quiz-question'>
                    <h3>Question {i}</h3>
                    <p><strong>{question['question']}</strong></p>
                </div>
                """, unsafe_allow_html=True)
                
                # Display options
                for j, option in enumerate(question['options'], 1):
                    correct_index = question['correct_answer']
                    
                    # Check if this is the correct answer (add 1 since our display is 1-indexed)
                    is_correct = (j == correct_index + 1)
                    
                    # Create option class based on correctness
                    option_class = "quiz-option quiz-option-correct" if is_correct else "quiz-option"
                    
                    st.markdown(f"""
                    <div class='{option_class}'>
                        {j}. {option} {' ✓' if is_correct else ''}
                    </div>
                    """, unsafe_allow_html=True)
                
                # Show explanation in an expander
                with st.expander("See Explanation"):
                    st.write(question['explanation'])
                    
    except json.JSONDecodeError as e:
        st.error(f"Error parsing quiz: {e}")
        st.write(quiz)  # Fallback to display raw output

def display_projects(projects):
    st.markdown("<div class='sub-header'>🚀 Suggested Projects</div>", unsafe_allow_html=True)
    
    if projects and hasattr(projects, 'projects'):
        for i, project in enumerate(projects.projects, 1):
            # Set badge class based on difficulty
            badge_class = ""
            if project.difficulty == "beginner":
                badge_class = "badge-beginner"
            elif project.difficulty == "intermediate":
                badge_class = "badge-intermediate"
            elif project.difficulty == "advanced":
                badge_class = "badge-advanced"
            
            st.markdown(f"""
            <div class='project-card'>
                <div class='project-header'>
                    <h3>Project #{i}: {project.title}</h3>
                    <div>
                        <span class='badge {badge_class}'>{project.difficulty.capitalize()}</span>
                        <span>⏱️ {project.estimated_duration}</span>
                    </div>
                </div>
                <p>{project.description}</p>
                <hr style='margin: 15px 0;'>
            </div>
            """, unsafe_allow_html=True)
            
            # Skills and outcomes in expandable sections
            col1, col2 = st.columns(2)
            
            with col1:
                with st.expander("📋 Required Skills"):
                    for skill in project.required_skills:
                        st.markdown(f"• {skill}")
            
            with col2:
                with st.expander("🎯 Learning Outcomes"):
                    for outcome in project.learning_outcomes:
                        st.markdown(f"• {outcome}")

def render_welcome_screen():
    col1, col2, col3 = st.columns([1, 3, 1])
    
    with col2:
        st.markdown("""
        <div style="text-align: center; padding: 2rem;">
            <h1 style="color: #1E88E5;">🎓 Learning Path Generator</h1>
            <p style="font-size: 1.2rem; margin: 20px 0;">
                Generate personalized learning paths, quizzes, and project ideas with AI assistance.
            </p>
            <span class="gemini-badge">Powered by Gemini 2.0</span>
        </div>
        """, unsafe_allow_html=True)
        
        st.markdown("""
        <div class="card">
            <h3>How It Works:</h3>
            <ol>
                <li>Enter your learning topics in the sidebar</li>
                <li>Select your expertise level</li>
                <li>Click "Generate Learning Path" to create personalized content</li>
            </ol>
        </div>
        """, unsafe_allow_html=True)
        
        # Feature highlights
        st.markdown("""
        <div style="display: flex; gap: 20px; margin-top: 20px;">
            <div style="flex: 1; padding: 20px; background-color: #e8f5e9; border-radius: 10px; text-align: center;">
                <h3>📚 Curated Resources</h3>
                <p>Get hand-picked learning materials tailored to your level</p>
            </div>
            <div style="flex: 1; padding: 20px; background-color: #e3f2fd; border-radius: 10px; text-align: center;">
                <h3>🧠 Interactive Quizzes</h3>
                <p>Test your knowledge with custom quizzes</p>
            </div>
            <div style="flex: 1; padding: 20px; background-color: #fff3e0; border-radius: 10px; text-align: center;">
                <h3>🚀 Project Ideas</h3>
                <p>Apply your skills with hands-on projects</p>
            </div>
        </div>
        """, unsafe_allow_html=True)

def main():
    # Initialize session state
    if 'generation_complete' not in st.session_state:
        st.session_state.generation_complete = False
    if 'results' not in st.session_state:
        st.session_state.results = None
    
    # Header
    st.markdown("<div class='main-header'>🎓 Learning Path Generator</div>", unsafe_allow_html=True)
    
    # Sidebar for inputs with enhanced styling
    with st.sidebar:
        st.image("https://www.svgrepo.com/show/374122/learning.svg", width=80)
        st.markdown("<h2>Configure Your Learning Path</h2>", unsafe_allow_html=True)
        
        # Gemini badge
        st.markdown("<div style='display: flex; justify-content: center; margin-bottom: 20px;'><span class='gemini-badge'>Powered by Gemini 2.0</span></div>", unsafe_allow_html=True)
        
        st.markdown("### Topics")
        topics = st.text_area(
            "Enter topics to learn (one per line)",
            placeholder="Example:\nPython Data Science\nMachine Learning\nDeep Learning",
            help="Enter the topics you want to learn about",
            height=150
        )
        
        st.markdown("### Your Level")
        expertise_level = st.selectbox(
            "Select your expertise level",
            options=[level.value for level in ExpertiseLevel],
            format_func=lambda x: x.capitalize(),
            help="Choose your current level of expertise"
        )
        
        # Add model selection dropdown
        model_options = [
            "gemini-2.0-flash-lite",
            "gemini-2.0-pro",
            "gpt-3.5-turbo"  # Fallback option
        ]
        selected_model = st.selectbox(
            "AI Model",
            options=model_options,
            index=0,
            help="Select the AI model to use"
        )
        
        # Store the selected model in session state
        if 'selected_model' not in st.session_state:
            st.session_state.selected_model = selected_model
        elif st.session_state.selected_model != selected_model:
            st.session_state.selected_model = selected_model
        
        generate_btn = st.button("🚀 Generate Learning Path", use_container_width=True, type="primary")
        
        st.markdown("---")
        st.markdown("""
        <div style="font-size: 0.8rem; color: #666;">
            Powered by CrewAI and Google Gemini<br>
            © 2025 Learning Path Generator
        </div>
        """, unsafe_allow_html=True)
    
    # Check for API keys
    if not os.getenv("GOOGLE_API_KEY") and st.session_state.selected_model.startswith("gemini"):
        st.warning("⚠️ Google API Key not found. Please add it to your environment variables.", icon="⚠️")
    
    # Main content area
    if not st.session_state.generation_complete and not generate_btn:
        render_welcome_screen()
    
    if generate_btn:
        if not topics:
            st.error("⚠️ Please enter at least one topic")
            return

        topic_list = [topic.strip() for topic in topics.split('\n') if topic.strip()]
        
        # Show a more visually appealing loading state
        st.markdown("""
        <div class='loading-container'>
            <h2>Generating Your Personalized Learning Path...</h2>
            <div class='progress-bar'>
                <div class='progress'></div>
            </div>
            <p>Our AI experts are crafting the perfect resources for you.<br>This may take a minute or two.</p>
        </div>
        """, unsafe_allow_html=True)
        
        try:
            # Try-except for better error handling
            try:
                # Initialize Gemini LLM and tools
                llm, search_tool = initialize_services()
                
                # Create agents and tasks with Gemini
                agents, tasks = create_agents_and_tasks(topics, expertise_level, llm)
                
                # Create and run crew
                crew = Crew(
                    agents=agents,
                    tasks=tasks,
                    process=Process.sequential
                )
                
                result = crew.kickoff({"topics": topic_list, "expertise_level": ExpertiseLevel(expertise_level)})
                
                # Store results in session state
                st.session_state.results = {
                    "materials": tasks[0].output.raw,
                    "quiz": tasks[1].output.raw,
                    "projects": result.pydantic
                }
                st.session_state.generation_complete = True
                
                # Rerun to display results
                st.rerun()
                
            except ImportError as ie:
                st.error(f"Missing package: {str(ie)}")
                st.info("Try installing required packages with: `pip install langchain-google-genai crewai pydantic`")
                
            except AttributeError as ae:
                st.error(f"Configuration issue: {str(ae)}")
                st.info("This might be a compatibility issue between CrewAI and the LLM integration.")
                
            except ValueError as ve:
                st.error(f"Value error: {str(ve)}")
                if "api_key" in str(ve).lower():
                    st.info("There seems to be an issue with your API key. Please check if it's correctly set in the .env file.")
                    
        except Exception as e:
            st.error(f"🚨 An error occurred: {str(e)}")
            st.info("If the issue persists, try switching to a different AI model in the sidebar.")
    
    # Display results if generation is complete
    if st.session_state.generation_complete and st.session_state.results:
        results = st.session_state.results
        
        # Create tabs with icons
        tab1, tab2, tab3 = st.tabs(["📚 Learning Materials", "🧠 Quiz", "🚀 Project Ideas"])
        
        with tab1:
            display_learning_materials(results["materials"])
        
        with tab2:
            display_quiz(results["quiz"])
        
        with tab3:
            display_projects(results["projects"])
        
        # Add footer
        st.markdown("""
        <div class='footer'>
            <p>Need to regenerate? Update your preferences in the sidebar and click 'Generate Learning Path' again.</p>
            <p>© 2025 Learning Path Generator • Powered by Google Gemini 2.0</p>
        </div>
        """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()