import streamlit as st
from pymongo import MongoClient
from datetime import datetime, timedelta
from dotenv import load_dotenv
import os
import google.generativeai as genai 
from file_upload_vectorize import resources_collection, vectors_collection
# Load environment variables
load_dotenv()
MONGO_URI = os.getenv("MONGO_URI")
client = MongoClient(MONGO_URI)
db = client["novascholar_db"]
live_chat_sessions_collection = db["live_chat_sessions"]

# Initialize AI model
genai.configure(api_key=os.getenv("GEMINI_KEY"))
model = genai.GenerativeModel("gemini-1.5-flash")


def display_live_chat_interface(session, user_id, course_id):
    """Main interface for live chat sessions - handles both faculty and student views"""
    st.markdown("<div style='margin-top: 20px;'></div>", unsafe_allow_html=True)
    st.markdown("#### Live Class Chat Session")
    
    # Initialize session states
    if 'chat_active' not in st.session_state:
        st.session_state.chat_active = False
    if 'chat_end_time' not in st.session_state:
        st.session_state.chat_end_time = None
    if 'messages' not in st.session_state:
        st.session_state.messages = []

    # Faculty View
    if st.session_state.user_type == "faculty":
        display_faculty_controls(session, user_id, course_id)
    # Student View
    else:
        display_student_view(session, user_id, course_id)

def create_timer_html(end_time):
    """Create a simplified but reliable timer component"""
    end_timestamp = int(end_time.timestamp() * 1000)  # Convert to milliseconds
    current_timestamp = int(datetime.utcnow().timestamp() * 1000)
    
    return f"""
        <div id="timer-container">
            <style>
                .timer-box {{
                    background: #f0f2f6;
                    border-radius: 8px;
                    padding: 15px;
                    text-align: center;
                    margin: 10px 0;
                }}
                .timer-display {{
                    font-size: 24px;
                    font-weight: bold;
                    color: #0066cc;
                    margin-bottom: 10px;
                }}
                .progress-bar {{
                    width: 100%;
                    height: 8px;
                    background: #e0e0e0;
                    border-radius: 4px;
                    overflow: hidden;
                }}
                .progress {{
                    height: 100%;
                    background: #00aa00;
                    transition: width 1s linear;
                }}
            </style>
            
            <div class="timer-box">
                <div id="timer" class="timer-display">Calculating...</div>
                <div class="progress-bar">
                    <div id="progress" class="progress"></div>
                </div>
            </div>

            <script>
                function updateTimer() {{
                    const endTime = {end_timestamp};
                    const startTime = {current_timestamp};
                    const now = new Date().getTime();
                    const totalDuration = endTime - startTime;
                    const timeLeft = endTime - now;
                    
                    const timer = document.getElementById('timer');
                    const progress = document.getElementById('progress');
                    
                    if (timeLeft <= 0) {{
                        timer.innerHTML = 'Session Ended';
                        timer.style.color = '#ff4444';
                        progress.style.width = '100%';
                        progress.style.background = '#ff4444';
                        return;
                    }}
                    
                    const minutes = Math.floor(timeLeft / (1000 * 60));
                    const seconds = Math.floor((timeLeft % (1000 * 60)) / 1000);
                    const progressWidth = ((totalDuration - timeLeft) / totalDuration) * 100;
                    
                    timer.innerHTML = `${{minutes}}:${{seconds < 10 ? '0' : ''}}${{seconds}}`;
                    progress.style.width = `${{progressWidth}}%`;
                    
                    if (minutes < 5) {{
                        progress.style.background = '#ffa500';
                    }}
                }}
                
                const timerInterval = setInterval(updateTimer, 1000);
                updateTimer();
            </script>
        </div>
    """

# def display_timer(end_time):
#     """Display a simple countdown timer"""
#     if not isinstance(end_time, datetime):
#         return
        
#     # Calculate remaining time
#     remaining_time = end_time - datetime.utcnow()
    
#     if remaining_time.total_seconds() <= 0:
#         st.session_state.chat_active = False
#         st.session_state.chat_end_time = None
#         return
    
#     # Convert to minutes and seconds
#     total_seconds = int(remaining_time.total_seconds())
#     minutes = total_seconds // 60
#     seconds = total_seconds % 60
    
#     # Create a progress value between 0 and 100
#     original_duration = st.session_state.get('original_duration', 15) * 60  # default 15 minutes in seconds
#     progress = ((original_duration - total_seconds) / original_duration)
    
#     # Display timer using columns for better layout
#     col1, col2 = st.columns([1, 28])
    
#     with col1:
#         st.markdown("### ⏱️")
#     with col2:
#         # Display time remaining
#         st.markdown(f"### {minutes:02d}:{seconds:02d}")
    
#     # Display progress bar
#     if minutes < 5:
#         color = "orange"
#     else:
#         color = "blue"
    
#     st.progress(progress, text="Session Progress")

def display_timer(end_time):
    """Display a simple countdown timer"""
    if not isinstance(end_time, datetime):
        return
        
    # Calculate remaining time
    remaining_time = end_time - datetime.utcnow()
    
    if remaining_time.total_seconds() <= 0:
        st.session_state.chat_active = False
        st.session_state.chat_end_time = None
        return
    
    # Convert to minutes and seconds
    total_seconds = int(remaining_time.total_seconds())
    minutes = total_seconds // 60
    seconds = total_seconds % 60
    
    # Create a progress value between 0 and 100
    original_duration = st.session_state.get('original_duration', 15) * 60  # default 15 minutes in seconds
    progress = ((original_duration - total_seconds) / original_duration)
    
    # Custom CSS for timer layout
    st.markdown("""
        <style>
        .timer-container {{
            margin: 1rem 0;
        }}
        .timer-row {{
            display: flex;
            justify-content: space-between;
            align-items: center;
            gap: 1rem;
            margin-bottom: 0.15rem;
        }}
        .timer-icon {{
            font-size: 1.5rem;
        }}
        .timer-text {{
            font-size: 1.2rem;
            font-weight: bold;
        }}
        .timer-progress-label {{
            font-size: 1.1rem;
            font-weight: semi-bold;
        }}
        .stProgress {{
            margin-bottom: 0.25rem;
        }}
        .stChatInputContainer {{
            margin-top: 1rem;
        }}
        </style>
        
        <div class="timer-container">
            <div class="timer-row">
                <div class="progress-text"><span class="timer-progress-label">Session Progress</span></div>
                <div class="timer-progress">
                <span class="timer-icon">⏱️</span>
                <span class="timer-text">{:02d}:{:02d}</span>
                </div>
            </div>
        </div>
    """.format(minutes, seconds), unsafe_allow_html=True)
    
    # Display progress bar
    color = "orange" if minutes < 5 else "blue"
    st.progress(progress)

def display_faculty_controls(session, faculty_id, course_id):
    """Display faculty controls for managing chat sessions"""
    
    # Show scheduled sessions
    st.markdown("<div style='margin-top: 20px;'></div>", unsafe_allow_html=True)
    st.markdown("##### 📅 Scheduled Chat Sessions")
    scheduled_sessions = list(live_chat_sessions_collection.find({
        "session_id": session['session_id'],
        "status": "scheduled"
    }))

    # Schedule new session
    with st.expander("➕ Schedule New Chat Session"):
        col1, col2, col3 = st.columns([2, 2, 1])
        with col1:
            session_date = st.date_input("📅 Date", min_value=datetime.now().date())
        with col2:
            session_time = st.time_input("🕒 Time", value=datetime.now().time())
        with col3:
            duration = st.selectbox(
                "⏱️ Duration (mins)",
                options=[5, 10, 15, 20, 30, 45, 60],
                index=2
            )
        
        if st.button("📋 Schedule Session"):
            session_datetime = datetime.combine(session_date, session_time)
            if session_datetime < datetime.now():
                st.error("❌ Cannot schedule sessions in the past!")
            else:
                live_chat_sessions_collection.insert_one({
                    "session_id": session['session_id'],
                    "course_id": course_id,
                    "faculty_id": faculty_id,
                    "start_time": session_datetime,
                    "duration": duration,
                    "status": "scheduled",
                    "chats": []
                })
                st.success("✅ Chat session scheduled successfully!")
                st.rerun()


    
    if scheduled_sessions:
        for scheduled in scheduled_sessions:
            with st.expander(f"📌 Scheduled: {scheduled['start_time'].strftime('%I:%M %p')}"):
                st.write(f"⏱️ Duration: {scheduled['duration']} minutes")
                if st.button("❌ Cancel Session", key=f"cancel_{scheduled['_id']}", type="secondary"):
                    live_chat_sessions_collection.delete_one({"_id": scheduled['_id']})
                    st.rerun()
    
    # Start immediate session
    if not st.session_state.get('chat_active', False):
        st.markdown("<div style='margin-top: 20px;'></div>", unsafe_allow_html=True)
        st.markdown("##### 🎯 Start Immediate Session")
        col1, col2 = st.columns([3, 1])
        with col1:
            immediate_duration = st.selectbox(
                "⏱️ Select duration (minutes)",
                options=[5, 10, 15, 20, 30, 45, 60],
                index=2
            )
        with col2:
            st.markdown("<div style='margin-top: 26px;'>", unsafe_allow_html=True)
            if st.button("▶️ Start", use_container_width=True):
                st.session_state.chat_active = True
                st.session_state.chat_end_time = datetime.utcnow() + timedelta(minutes=immediate_duration)
                st.session_state.original_duration = immediate_duration  # Store original duration
                live_chat_sessions_collection.insert_one({
                    "session_id": session['session_id'],
                    "course_id": course_id,
                    "faculty_id": faculty_id,
                    "start_time": datetime.utcnow(),
                    "duration": immediate_duration,
                    "status": "active",
                    "chats": []
                })
                st.rerun()
            st.markdown("</div>", unsafe_allow_html=True)
    else:
        # Show timer for active session
        if hasattr(st.session_state, 'chat_end_time'):
            display_timer(st.session_state.chat_end_time)
            
            if st.button("⏹️ End Session", type="secondary"):
                st.session_state.chat_active = False
                st.session_state.chat_end_time = None
                live_chat_sessions_collection.update_one(
                    {"session_id": session['session_id'], "status": "active"},
                    {"$set": {"status": "completed", "end_time": datetime.utcnow()}}
                )
                st.rerun()

def get_session_context(session, session_id, course_id):
    """Retrieve session context for AI model"""
    # Get session data
    session = live_chat_sessions_collection.find_one({"session_id": session_id})
    if not session:
        st.error("Session not found")
        return
    
    # Get pre-class materials
    context = ""
    materials = resources_collection.find({"session_id": session_id})
    for material in materials:
        resource_id = material['_id']
        vector_data = vectors_collection.find_one({"resource_id": resource_id})
        if vector_data and 'text' in vector_data:
            context += vector_data['text'] + "\n"
    
    courses_collection = db["courses"]
    course = courses_collection.find_one({"_id": course_id})
    session = courses_collection.find_one({"sessions.session_id": session_id})
    if course:
        context += f"Course: {course['course_name']}\n"
    if session:
        context += f"Session: {session['title']}\n"
        if 'session_learning_outcomes' in session:
            context += f"Session Learning Outcomes: {', '.join(session['session_learning_outcomes'])}\n"
    return context

def display_student_view(session, student_id, course_id):
    """Display student interface for chat sessions"""
    
    # Show upcoming scheduled sessions
    st.markdown("<div style='margin-top: 20px;'></div>", unsafe_allow_html=True)
    st.markdown("##### 📅 Upcoming Chat Sessions")
    upcoming_sessions = list(live_chat_sessions_collection.find({
        "session_id": session['session_id'],
        "status": "scheduled",
        "start_time": {"$gt": datetime.utcnow()} 
    }).sort("start_time", 1))
    
    if upcoming_sessions:
        for upcoming in upcoming_sessions:
            with st.expander(f"📌 {upcoming['start_time'].strftime('%I:%M %p')}"):
                st.write(f"⏱️ Duration: {upcoming['duration']} minutes")
                time_until = upcoming['start_time'] - datetime.utcnow()
                st.write(f"🕒 Starts in: {time_until.seconds // 3600}h {(time_until.seconds % 3600) // 60}m")
    else:
        st.info("📝 No upcoming chat sessions scheduled")
    
    # Check for active session
    active_session = live_chat_sessions_collection.find_one({
        "session_id": session['session_id'],
        "status": "active"
    })
    
    # if active_session:
    #     st.session_state.chat_active = True
    #     st.session_state.chat_end_time = active_session['start_time'] + timedelta(minutes=active_session['duration'])
    if active_session and not st.session_state.get('chat_active', False):
        st.session_state.chat_active = True
        st.session_state.chat_end_time = active_session['start_time'] + timedelta(minutes=active_session['duration'])
        st.session_state.original_duration = active_session['duration']
    
    if st.session_state.get('chat_active', False):
        st.markdown("<div style='margin-top: 20px;'></div>", unsafe_allow_html=True)
        st.markdown("##### 🔴 Live Chat Section")

        if hasattr(st.session_state, 'chat_end_time'):
            reamining_time = st.session_state.chat_end_time - datetime.utcnow()
            if reamining_time.total_seconds() > 0:
                display_timer(st.session_state.chat_end_time)
                active_session = live_chat_sessions_collection.find_one({
                    "session_id": session['session_id'],
                    "status": "active"
                })
                if active_session:
                    faculty_id = active_session['faculty_id']
                    display_chat_interface(session, student_id, course_id, faculty_id)
                else:
                    st.error("❌ No active session found.")
                    st.session_state.chat_active = False
                    st.session_state.chat_end_time = None
            else:
                st.session_state.chat_active = False
                st.session_state.chat_end_time = None
                live_chat_sessions_collection.update_one(
                    {"session_id": session['session_id'], "status": "active"},
                    {"$set": {"status": "completed", "end_time": datetime.utcnow()}}
                )
                st.rerun()

def display_chat_interface(session, user_id, course_id, faculty_id):
    """Display the actual chat interface with messages"""
    # Display chat messages
    # Initialize 'messages' in session_state if it doesn't exist
    if 'messages' not in st.session_state:
        st.session_state.messages = []
    
    # Chat input
    if prompt := st.chat_input("Ask Questions about the Session"):
        # Get session context
        context = ""
        context = get_session_context(session, session['session_id'], course_id)
        
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message
        with st.chat_message("user"):
            st.markdown(prompt)
        
        try:
            # Generate AI response
            context_prompt = f"""
            You are an intelligent teaching assistant participating in a live class discussion.
            
            Session Context:
            {context}

            Current Question by the Student: {prompt}

            Instructions:
            1. Provide clear, concise responses that relate to the session's key concepts and learning outcomes
            2. Encourage critical thinking and discussion
            3. Keep responses focused and relevant to the current topic
            4. If students seem confused, provide clarifying examples
            
            Please provide an appropriate response for this live classroom discussion.
            """
            
            response = model.generate_content(context_prompt)
            assistant_response = response.text
            
            # Display assistant response
            with st.chat_message("assistant"):
                st.markdown(assistant_response)
            
            # Save chat message
            chat_message = {
                "session_id": session['session_id'],
                "timestamp": datetime.utcnow(),
                "user_id": user_id,
                "user_type": st.session_state.user_type,
                "message": prompt,
                "response": assistant_response
            }
            st.session_state.messages.append(chat_message)
            # chat_messages_collection.insert_one(chat_message)
            # Update database
            try:
                current_session = live_chat_sessions_collection.find_one({"session_id": session['session_id'], "status": "active"})
                live_chat_sessions_collection.update_one(
                    {
                        "session_id": session['session_id'],
                        "course_id": course_id,
                        "faculty_id": faculty_id,
                        "start_time": current_session['start_time'],
                        "duration": current_session['duration'],
                        "status": current_session['status'],
                        "chats.user_id": user_id
                    },
                    {
                        "$push": {
                            "chats.$.messages": {
                                "prompt": prompt,
                                "response": assistant_response,
                                "timestamp": datetime.utcnow()
                            }
                        }
                    }
                )

                # If no existing chat object for the user, create a new one
                if live_chat_sessions_collection.find_one({
                    "session_id": session['session_id'],
                    "course_id": course_id,
                    "faculty_id": faculty_id,
                    "start_time": current_session['start_time'],
                    "duration": current_session['duration'],
                    "status": current_session['status'],
                    "chats.user_id": user_id
                }) is None:
                    live_chat_sessions_collection.update_one(
                        {
                            "session_id": session['session_id'],
                            "course_id": course_id,
                            "faculty_id": faculty_id,
                            "start_time": current_session['start_time'],
                            "duration": current_session['duration'],
                            "status": current_session['status']
                        },
                        {
                            "$push": {
                                "chats": {
                                    "user_id": user_id,
                                    "messages": [
                                        {
                                            "prompt": prompt,
                                            "response": assistant_response,
                                            "timestamp": datetime.utcnow()
                                        }
                                    ]
                                }
                            }
                        },
                        upsert=True
                    )
            except Exception as db_error:
                st.error(f"Error saving chat history: {str(db_error)}")
        except Exception as e:
            st.error(f"Error processing message: {str(e)}")