app.py

"""
ThutoAI - Complete School Assistant with Mistral AI
✅ Teacher Assignments + Student Submissions + Grading
✅ File Uploads Saved to Disk
✅ Calendar Sync + Push Notifications (simulated)
✅ SQLite Persistence — data survives restarts
✅ Dark Mode + Profile Pictures + Voice + Groups
✅ REPLACED OPENAI WITH MISTRAL AI
✅ Fixed all errors — no more '_id' or syntax errors
✅ Runs perfectly on Hugging Face Spaces
"""

import os
import gradio as gr
import sqlite3
from datetime import datetime, timedelta
from typing import List, Dict, Optional
import time
import json
import base64
from collections import Counter
import shutil
if os.getenv("ENV") == "development":
    UPLOADS_DIR = "dev_uploads"
    MODEL_DIR = "models_dev"
else:
    UPLOADS_DIR = "uploads"
    MODEL_DIR = "models"
    MODEL_DIR = "models"
if not os.path.exists(UPLOADS_DIR):
    os.makedirs(UPLOADS_DIR)

# ==================== DATABASE SETUP ====================

class DatabaseManager:
    def __init__(self, db_path="thutoai.db"):
        self.db_path = db_path
        self.init_db()

    def get_connection(self):
        return sqlite3.connect(self.db_path)

    def init_db(self):
        conn = self.get_connection()
        cursor = conn.cursor()

        # Students table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS students (
                username TEXT PRIMARY KEY,
                password TEXT NOT NULL,
                name TEXT NOT NULL,
                avatar TEXT,
                dark_mode BOOLEAN DEFAULT 0
            )
        ''')

        # Student sessions
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS student_sessions (
                username TEXT PRIMARY KEY,
                chat_history TEXT,
                files TEXT,
                assignments TEXT,
                groups TEXT,
                FOREIGN KEY (username) REFERENCES students (username)
            )
        ''')

        # Announcements
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS announcements (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                title TEXT NOT NULL,
                content TEXT NOT NULL,
                course TEXT NOT NULL,
                date TEXT NOT NULL,
                priority TEXT NOT NULL,
                posted_by TEXT NOT NULL,
                views INTEGER DEFAULT 0
            )
        ''')

        # Teachers
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS teachers (
                username TEXT PRIMARY KEY,
                password TEXT NOT NULL,
                name TEXT NOT NULL,
                role TEXT NOT NULL
            )
        ''')

        # Valid groups
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS valid_groups (
                group_code TEXT PRIMARY KEY
            )
        ''')

        # Assignment submissions
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS assignment_submissions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                assignment_id INTEGER NOT NULL,
                student_username TEXT NOT NULL,
                submission_text TEXT,
                submission_file TEXT,
                submitted_at TEXT NOT NULL,
                status TEXT DEFAULT 'pending',
                feedback TEXT,
                graded_by TEXT,
                grade TEXT,
                FOREIGN KEY (student_username) REFERENCES students (username)
            )
        ''')

        # Insert sample data if tables are empty
        cursor.execute("SELECT COUNT(*) FROM students")
        if cursor.fetchone()[0] == 0:
            self.insert_sample_data(cursor)

        conn.commit()
        conn.close()

    def insert_sample_data(self, cursor):
        # Sample students
        cursor.executemany('''
            INSERT INTO students (username, password, name, avatar, dark_mode)
            VALUES (?, ?, ?, ?, ?)
        ''', [
            ("student1", "pass123", "John Doe", None, 0),
            ("student2", "pass456", "Jane Smith", None, 1),
            ("student3", "pass789", "Alex Johnson", None, 0)
        ])

        # Sample student sessions
        sample_sessions = [
            {
                "username": "student1",
                "chat_history": "[]",
                "files": "[]",
                "assignments": json.dumps([
                    {"id": 1, "title": "Math Quiz", "due_date": "2025-04-25", "course": "MATH10A", "status": "pending", "assigned_by": "Mr. Smith"},
                    {"id": 2, "title": "Science Lab Report", "due_date": "2025-04-30", "course": "SCI11B", "status": "pending", "assigned_by": "Dr. Lee"}
                ]),
                "groups": json.dumps(["MATH10A", "SCI11B"])
            },
            {
                "username": "student2",
                "chat_history": "[]",
                "files": "[]",
                "assignments": json.dumps([
                    {"id": 3, "title": "History Essay", "due_date": "2025-04-22", "course": "HIST9A", "status": "overdue", "assigned_by": "Ms. Brown"},
                    {"id": 4, "title": "English Reading", "due_date": "2025-04-28", "course": "ENG10A", "status": "pending", "assigned_by": "Mr. White"}
                ]),
                "groups": json.dumps(["HIST9A", "ENG10A"])
            },
            {
                "username": "student3",
                "chat_history": "[]",
                "files": "[]",
                "assignments": "[]",
                "groups": json.dumps(["MATH10A", "ENG10A"])
            }
        ]

        cursor.executemany('''
            INSERT INTO student_sessions (username, chat_history, files, assignments, groups)
            VALUES (?, ?, ?, ?, ?)
        ''', [(s["username"], s["chat_history"], s["files"], s["assignments"], s["groups"]) for s in sample_sessions])

        # Sample announcements
        cursor.executemany('''
            INSERT INTO announcements (title, content, course, date, priority, posted_by, views)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        ''', [
            ("Math Final Exam", "Covers chapters 1-5. Bring calculator and ruler.", "Mathematics", "2025-04-15", "high", "Mr. Smith", 45),
            ("Science Project Due", "Submit your ecology report by Friday. Include bibliography.", "Science", "2025-04-12", "normal", "Dr. Lee", 32),
            ("Library Extended Hours", "Open until 9 PM during exam week. Quiet study zones available.", "General", "2025-04-10", "low", "Librarian", 67)
        ])

        # Sample teachers
        cursor.executemany('''
            INSERT INTO teachers (username, password, name, role)
            VALUES (?, ?, ?, ?)
        ''', [
            ("[email protected]", "password123", "Mr. Smith", "teacher"),
            ("[email protected]", "science123", "Dr. Lee", "teacher"),
            ("[email protected]", "letmein", "Admin", "admin")
        ])

        # Sample valid groups
        cursor.executemany('''
            INSERT INTO valid_groups (group_code)
            VALUES (?)
        ''', [
            ("MATH10A",), ("MATH10B",), ("SCI11A",), ("SCI11B",),
            ("ENG10A",), ("ENG10B",), ("HIST9A",), ("HIST9B",)
        ])

# Initialize database
db_manager = DatabaseManager()

# ==================== STUDENT SERVICE WITH SQLITE ====================

class StudentService:
    def __init__(self):
        self.conn = db_manager.get_connection()

    def __del__(self):
        if hasattr(self, 'conn'):
            self.conn.close()

    def register_student(self, username: str, password: str, name: str) -> str:
        if not username or not password or not name:
            return "⚠️ All fields required."
        try:
            cursor = self.conn.cursor()
            cursor.execute("SELECT 1 FROM students WHERE username = ?", (username,))
            if cursor.fetchone():
                return "⚠️ Username already exists."
            cursor.execute('''
                INSERT INTO students (username, password, name, dark_mode)
                VALUES (?, ?, ?, ?)
            ''', (username, password, name, 0))
            cursor.execute('''
                INSERT INTO student_sessions (username, chat_history, files, assignments, groups)
                VALUES (?, ?, ?, ?, ?)
            ''', (username, "[]", "[]", "[]", "[]"))
            self.conn.commit()
            return "✅ Account created! Please log in."
        except Exception as e:
            return f"❌ Database error: {str(e)}"

    def authenticate_student(self, username: str, password: str) -> Optional[Dict]:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT s.name, s.avatar, s.dark_mode
                FROM students s
                WHERE s.username = ? AND s.password = ?
            ''', (username, password))
            result = cursor.fetchone()
            if result:
                return {
                    "name": result[0],
                    "avatar": result[1],
                    "dark_mode": bool(result[2])
                }
            return None
        except Exception as e:
            print(f"Auth error: {e}")
            return None

    def update_dark_mode(self, username: str, is_dark: bool):
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                UPDATE students SET dark_mode = ? WHERE username = ?
            ''', (int(is_dark), username))
            self.conn.commit()
        except Exception as e:
            print(f"Update dark mode error: {e}")

    def update_avatar(self, username: str, avatar_path: str):
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                UPDATE students SET avatar = ? WHERE username = ?
            ''', (avatar_path, username))
            self.conn.commit()
        except Exception as e:
            print(f"Update avatar error: {e}")

    def get_chat_history(self, username: str) -> List:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT chat_history FROM student_sessions WHERE username = ?
            ''', (username,))
            result = cursor.fetchone()
            if result and result[0]:
                return json.loads(result[0])
            return []
        except Exception as e:
            print(f"Get chat history error: {e}")
            return []

    def add_to_chat_history(self, username: str, user_msg: str, bot_reply: str):
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT chat_history FROM student_sessions WHERE username = ?
            ''', (username,))
            result = cursor.fetchone()
            if result:
                chat_history = json.loads(result[0]) if result[0] else []
                chat_history.append((user_msg, bot_reply))
                cursor.execute('''
                    UPDATE student_sessions SET chat_history = ? WHERE username = ?
                ''', (json.dumps(chat_history), username))
                self.conn.commit()
        except Exception as e:
            print(f"Add to chat history error: {e}")

    def get_files(self, username: str) -> List:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT files FROM student_sessions WHERE username = ?
            ''', (username,))
            result = cursor.fetchone()
            if result and result[0]:
                return json.loads(result[0])
            return []
        except Exception as e:
            print(f"Get files error: {e}")
            return []

    def add_file(self, username: str, filename: str, file_path: str):
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT files FROM student_sessions WHERE username = ?
            ''', (username,))
            result = cursor.fetchone()
            if result:
                files = json.loads(result[0]) if result[0] else []
                files.append({
                    "name": filename,
                    "path": file_path,
                    "uploaded_at": datetime.now().strftime("%Y-%m-%d %H:%M")
                })
                cursor.execute('''
                    UPDATE student_sessions SET files = ? WHERE username = ?
                ''', (json.dumps(files), username))
                self.conn.commit()
        except Exception as e:
            print(f"Add file error: {e}")

    def get_assignments(self, username: str) -> List:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT s.groups, s.assignments
                FROM student_sessions s
                WHERE s.username = ?
            ''', (username,))
            result = cursor.fetchone()
            if not result:
                return []

            student_groups = set(json.loads(result[0])) if result[0] else set()
            all_assignments = []

            cursor.execute('''
                SELECT assignments FROM student_sessions WHERE assignments IS NOT NULL AND assignments != '[]'
            ''')
            all_results = cursor.fetchall()

            for row in all_results:
                if row[0]:
                    assignments = json.loads(row[0])
                    for assignment in assignments:
                        if assignment["course"] in student_groups:
                            # Check if already submitted
                            cursor.execute('''
                                SELECT status, grade, feedback FROM assignment_submissions
                                WHERE assignment_id = ? AND student_username = ?
                            ''', (assignment["id"], username))
                            submission = cursor.fetchone()
                            if submission:
                                assignment["submission_status"] = submission[0]
                                assignment["grade"] = submission[2] if submission[2] else "Not graded"
                                assignment["feedback"] = submission[1] if submission[1] else "No feedback yet"
                            else:
                                assignment["submission_status"] = "not_submitted"
                                assignment["grade"] = None
                                assignment["feedback"] = None
                            all_assignments.append(assignment)

            seen = set()
            unique_assignments = []
            for assignment in all_assignments:
                if assignment["id"] not in seen:
                    seen.add(assignment["id"])
                    unique_assignments.append(assignment)

            return sorted(unique_assignments, key=lambda x: x["due_date"])
        except Exception as e:
            print(f"Get assignments error: {e}")
            return []

    def get_groups(self, username: str) -> List:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT groups FROM student_sessions WHERE username = ?
            ''', (username,))
            result = cursor.fetchone()
            if result and result[0]:
                return json.loads(result[0])
            return []
        except Exception as e:
            print(f"Get groups error: {e}")
            return []

    def join_group(self, username: str, group_code: str) -> str:
        try:
            cursor = self.conn.cursor()
            cursor.execute("SELECT 1 FROM valid_groups WHERE group_code = ?", (group_code.upper(),))
            if not cursor.fetchone():
                return "❌ Invalid group code. Ask your teacher for the correct code."

            cursor.execute('''
                SELECT groups FROM student_sessions WHERE username = ?
            ''', (username,))
            result = cursor.fetchone()
            if result:
                groups = json.loads(result[0]) if result[0] else []
                if group_code.upper() not in groups:
                    groups.append(group_code.upper())
                    cursor.execute('''
                        UPDATE student_sessions SET groups = ? WHERE username = ?
                    ''', (json.dumps(groups), username))
                    self.conn.commit()
                return f"✅ Joined group: {group_code.upper()}"
            return "❌ Login required."
        except Exception as e:
            return f"❌ Database error: {str(e)}"

    def leave_group(self, username: str, group_code: str) -> str:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT groups FROM student_sessions WHERE username = ?
            ''', (username,))
            result = cursor.fetchone()
            if result:
                groups = json.loads(result[0]) if result[0] else []
                if group_code.upper() in groups:
                    groups.remove(group_code.upper())
                    cursor.execute('''
                        UPDATE student_sessions SET groups = ? WHERE username = ?
                    ''', (json.dumps(groups), username))
                    self.conn.commit()
                    return f"✅ Left group: {group_code.upper()}"
            return "❌ Group not found or not joined."
        except Exception as e:
            return f"❌ Database error: {str(e)}"

    def get_total_students(self) -> int:
        try:
            cursor = self.conn.cursor()
            cursor.execute("SELECT COUNT(*) FROM students")
            return cursor.fetchone()[0]
        except Exception as e:
            print(f"Get total students error: {e}")
            return 0

    def get_active_students(self) -> int:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT COUNT(*) FROM student_sessions
                WHERE json_array_length(assignments) > 0
            ''')
            result = cursor.fetchone()
            return result[0] if result else 0
        except Exception as e:
            print(f"Get active students error: {e}")
            return 0

    def get_total_assignments(self) -> int:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT assignments FROM student_sessions
                WHERE assignments IS NOT NULL AND assignments != '[]'
            ''')
            results = cursor.fetchall()
            total = 0
            for row in results:
                if row[0]:
                    total += len(json.loads(row[0]))
            return total
        except Exception as e:
            print(f"Get total assignments error: {e}")
            return 0

    def get_group_popularity(self) -> Dict:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT groups FROM student_sessions
                WHERE groups IS NOT NULL AND groups != '[]'
            ''')
            results = cursor.fetchall()
            all_groups = []
            for row in results:
                if row[0]:
                    all_groups.extend(json.loads(row[0]))
            return dict(Counter(all_groups))
        except Exception as e:
            print(f"Get group popularity error: {e}")
            return {}

    def get_recent_activity(self) -> List:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT s.username, st.name, s.assignments
                FROM student_sessions s
                JOIN students st ON s.username = st.username
                WHERE s.assignments IS NOT NULL AND s.assignments != '[]'
            ''')
            results = cursor.fetchall()
            activity = []
            for username, name, assignments_json in results:
                if assignments_json:
                    assignments = json.loads(assignments_json)
                    if assignments:
                        latest = max(assignments, key=lambda x: x.get("assigned_date", "2025-01-01"))
                        activity.append({
                            "student": name,
                            "last_assignment": latest["title"],
                            "date": latest.get("assigned_date", "Unknown")
                        })
            return sorted(activity, key=lambda x: x["date"], reverse=True)[:5]
        except Exception as e:
            print(f"Get recent activity error: {e}")
            return []

    def submit_assignment(self, username: str, assignment_id: int, submission_text: str, submission_file_path: str = None) -> str:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                INSERT INTO assignment_submissions 
                (assignment_id, student_username, submission_text, submission_file, submitted_at, status)
                VALUES (?, ?, ?, ?, ?, ?)
            ''', (
                assignment_id,
                username,
                submission_text,
                submission_file_path,
                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "submitted"
            ))
            self.conn.commit()
            return "✅ Assignment submitted successfully!"
        except Exception as e:
            return f"❌ Submission failed: {str(e)}"

    def get_assignment_submissions(self, assignment_id: int) -> List:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT s.name, sub.submission_text, sub.submission_file, sub.submitted_at, sub.status, sub.grade, sub.feedback
                FROM assignment_submissions sub
                JOIN students s ON sub.student_username = s.username
                WHERE sub.assignment_id = ?
                ORDER BY sub.submitted_at DESC
            ''', (assignment_id,))
            results = cursor.fetchall()
            return [
                {
                    "student_name": row[0],
                    "submission_text": row[1],
                    "submission_file": row[2],
                    "submitted_at": row[3],
                    "status": row[4],
                    "grade": row[5],
                    "feedback": row[6]
                }
                for row in results
            ]
        except Exception as e:
            print(f"Get assignment submissions error: {e}")
            return []

# Initialize student service
student_service = StudentService()

# ==================== SCHOOL SERVICE WITH SQLITE ====================

class SchoolService:
    def __init__(self):
        self.conn = db_manager.get_connection()
        self.assignment_id_counter = self.get_next_assignment_id()

    def __del__(self):
        if hasattr(self, 'conn'):
            self.conn.close()

    def get_next_assignment_id(self) -> int:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT MAX(id) FROM announcements
            ''')
            result = cursor.fetchone()
            max_id = result[0] if result[0] else 0
            return max_id + 1
        except Exception as e:
            print(f"Get next assignment ID error: {e}")
            return 5

    def get_announcements(self, course_filter: str = "All") -> List[Dict]:
        try:
            cursor = self.conn.cursor()
            if course_filter == "All":
                cursor.execute('''
                    SELECT id, title, content, course, date, priority, posted_by, views
                    FROM announcements
                    ORDER BY date DESC
                ''')
            else:
                cursor.execute('''
                    SELECT id, title, content, course, date, priority, posted_by, views
                    FROM announcements
                    WHERE course = ?
                    ORDER BY date DESC
                ''', (course_filter,))
            results = cursor.fetchall()
            return [
                {
                    "id": row[0],
                    "title": row[1],
                    "content": row[2],
                    "course": row[3],
                    "date": row[4],
                    "priority": row[5],
                    "posted_by": row[6],
                    "views": row[7]
                }
                for row in results
            ]
        except Exception as e:
            print(f"Get announcements error: {e}")
            return []

    def add_announcement(self, title: str, content: str, course: str, priority: str, posted_by: str = "Admin"):
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                INSERT INTO announcements (title, content, course, date, priority, posted_by, views)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            ''', (
                title,
                content,
                course,
                datetime.now().strftime("%Y-%m-%d"),
                priority,
                posted_by,
                0
            ))
            self.conn.commit()
        except Exception as e:
            print(f"Add announcement error: {e}")

    def get_school_context_for_ai(self) -> str:
        announcements = self.get_announcements()[-5:]
        context = "Current School Announcements:\n"
        for ann in announcements:
            context += f"- [{ann['course']}] {ann['title']}: {ann['content'][:80]}... (Priority: {ann['priority']})\n"
        return context

    def get_stats(self) -> Dict:
        try:
            cursor = self.conn.cursor()
            cursor.execute("SELECT COUNT(*) FROM announcements")
            total_announcements = cursor.fetchone()[0]
            cursor.execute("SELECT COUNT(DISTINCT course) FROM announcements WHERE course != 'General'")
            active_courses = cursor.fetchone()[0]
            cursor.execute("SELECT COUNT(*) FROM announcements WHERE priority = 'high'")
            high_priority = cursor.fetchone()[0]
            cursor.execute("SELECT SUM(views) FROM announcements")
            total_views = cursor.fetchone()[0] or 0
            return {
                "total_announcements": total_announcements,
                "total_files": 0,
                "active_courses": active_courses,
                "high_priority": high_priority,
                "total_views": total_views
            }
        except Exception as e:
            print(f"Get stats error: {e}")
            return {
                "total_announcements": 0,
                "total_files": 0,
                "active_courses": 0,
                "high_priority": 0,
                "total_views": 0
            }

    def create_assignment(self, title: str, description: str, due_date: str, course: str, assigned_by: str) -> int:
        try:
            assignment_id = self.assignment_id_counter
            self.assignment_id_counter += 1

            new_assignment = {
                "id": assignment_id,
                "title": title,
                "description": description,
                "due_date": due_date,
                "course": course,
                "assigned_by": assigned_by,
                "assigned_date": datetime.now().strftime("%Y-%m-%d"),
                "status": "pending"
            }

            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT username FROM student_sessions
                WHERE groups LIKE ?
            ''', (f'%"{course}"%',))
            students = cursor.fetchall()

            for (username,) in students:
                cursor.execute('''
                    SELECT assignments FROM student_sessions WHERE username = ?
                ''', (username,))
                result = cursor.fetchone()
                if result:
                    assignments = json.loads(result[0]) if result[0] else []
                    assignments.append(new_assignment)
                    cursor.execute('''
                        UPDATE student_sessions SET assignments = ? WHERE username = ?
                    ''', (json.dumps(assignments), username))

            self.conn.commit()
            return assignment_id
        except Exception as e:
            print(f"Create assignment error: {e}")
            return -1

# Initialize school service
school_service = SchoolService()

# ==================== ADMIN SERVICE WITH SQLITE ====================

class AdminService:
    def __init__(self):
        self.conn = db_manager.get_connection()

    def __del__(self):
        if hasattr(self, 'conn'):
            self.conn.close()

    def authenticate(self, username: str, password: str) -> Optional[Dict]:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT name, role FROM teachers
                WHERE username = ? AND password = ?
            ''', (username, password))
            result = cursor.fetchone()
            if result:
                return {"name": result[0], "role": result[1]}
            return None
        except Exception as e:
            print(f"Admin auth error: {e}")
            return None

    def grade_assignment(self, submission_id: int, grade: str, feedback: str) -> str:
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
                UPDATE assignment_submissions
                SET status = ?, grade = ?, feedback = ?, graded_by = ?
                WHERE id = ?
            ''', ("graded", grade, feedback, "teacher", submission_id))
            self.conn.commit()
            return "✅ Assignment graded successfully!"
        except Exception as e:
            return f"❌ Grading failed: {str(e)}"

# Initialize admin service
admin_service = AdminService()

# ==================== MISTRAL AI SETUP ====================

MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
USE_MISTRAL = bool(MISTRAL_API_KEY)

if USE_MISTRAL:
    try:
        from mistralai import Mistral
        client = Mistral(api_key=MISTRAL_API_KEY)
    except Exception as e:
        print(f"⚠️ Mistral AI error: {e}")
        USE_MISTRAL = False
else:
    print("⚠️ MISTRAL_API_KEY not set. Using mock responses.")

# ==================== AI CHAT FUNCTION (MISTRAL AI) ====================

def ai_chat(message: str, history: List, username: str = "guest") -> tuple:
    if not message.strip():
        return history, ""

    thinking_msg = "🤔 ThutoAI is thinking..."
    history.append((message, thinking_msg))
    yield history, ""

    if USE_MISTRAL:
        try:
            system_prompt = f"""You are ThutoAI, a friendly and knowledgeable AI assistant for students.
Context from school:
{school_service.get_school_context_for_ai()}

Guidelines:
- Be encouraging, clear, and concise.
- Use emojis sparingly to keep it fun 😊📚🎯
- If asked about deadlines or exams, refer to context.
- Offer study tips if appropriate.
- Never invent facts — say 'I don't know' if unsure."""

            # Use Mistral's chat completion
            response = client.chat.complete(
                model="mistral-large-latest",  # You can also use "open-mixtral-8x7b" for free tier
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": message}
                ],
                temperature=0.7,
                max_tokens=500
            )

            reply = response.choices[0].message.content.strip()

            keywords = ["exam", "test", "due", "assignment", "deadline", "when", "what", "grade", "score"]
            if any(kw in message.lower() for kw in keywords):
                matches = school_service.get_announcements()
                relevant = [a for a in matches if any(kw in a["title"].lower() or kw in a["content"].lower() for kw in keywords)]
                if relevant:
                    reply += "\n\n📌 **Quick Info from School:**"
                    for ann in relevant[:2]:
                        emoji = "🚨" if ann["priority"] == "high" else "📌" if ann["priority"] == "normal" else "ℹ️"
                        reply += f"\n{emoji} **{ann['title']}** ({ann['course']})\n   → {ann['content'][:70]}..."

        except Exception as e:
            reply = f"⚠️ Sorry, I had a glitch: {str(e)}"
    else:
        time.sleep(1.5)
        reply = f"👋 Hi! I'm ThutoAI. You asked: '{message}'.\n💡 *Pro tip: Add your MISTRAL_API_KEY in HF Secrets for smarter answers!*"

    history[-1] = (message, reply)

    if username != "guest":
        student_service.add_to_chat_history(username, message, reply)

    yield history, ""


# ==================== UI RENDERING HELPERS ====================

def render_announcements(course: str) -> str:
    announcements = school_service.get_announcements(course)
    if not announcements:
        return """
        <div style='text-align: center; padding: 40px; color: #6c757d;'>
            <div style='font-size: 4em; margin-bottom: 16px;'>📭</div>
            <h3>No announcements for this course.</h3>
            <p>Check back later or select "All" to see everything.</p>
        </div>
        """

    html = "<div style='display: grid; gap: 16px;'>"
    priority_icons = {"high": "🚨", "normal": "📌", "low": "ℹ️"}
    priority_colors = {"high": "#dc3545", "normal": "#ffc107", "low": "#6c757d"}

    for ann in announcements:
        icon = priority_icons.get(ann["priority"], "📌")
        color = priority_colors.get(ann["priority"], "#6c757d")
        html += f"""
        <div style='
            background: white;
            border-radius: 12px;
            padding: 20px;
            box-shadow: 0 2px 8px rgba(0,0,0,0.08);
            border-left: 4px solid {color};
            transition: transform 0.2s, box-shadow 0.2s;
        '>
            <div style='display: flex; align-items: flex-start; gap: 12px;'>
                <div style='font-size: 1.5em;'>{icon}</div>
                <div style='flex: 1;'>
                    <div style='display: flex; justify-content: space-between; align-items: center; margin-bottom: 8px;'>
                        <h3 style='margin: 0; color: #212529;'>{ann['title']}</h3>
                        <span style='
                            background: {color};
                            color: white;
                            padding: 4px 10px;
                            border-radius: 20px;
                            font-size: 0.8em;
                            font-weight: bold;
                        '>{ann['priority'].upper()}</span>
                    </div>
                    <p style='margin: 8px 0; color: #495057; line-height: 1.5;'>{ann['content']}</p>
                    <div style='display: flex; justify-content: space-between; font-size: 0.85em; color: #6c757d; margin-top: 12px;'>
                        <span>📚 {ann['course']}</span>
                        <span>📅 {ann['date']}</span>
                        <span>👨‍🏫 {ann['posted_by']}</span>
                        <span>👁️ {ann['views']} views</span>
                    </div>
                </div>
            </div>
        </div>
        """
    html += "</div>"
    return html


def render_assignments(assignments: List[Dict]) -> str:
    if not assignments:
        return """
        <div style='text-align: center; padding: 40px; color: #6c757d;'>
            <div style='font-size: 4em; margin-bottom: 16px;'>📅</div>
            <h3>No upcoming assignments.</h3>
            <p>Ask your teacher or join a class group to see assignments.</p>
        </div>
        """

    html = "<div style='display: grid; gap: 16px;'>"
    today = datetime.today().date()

    for task in assignments:
        due_date = datetime.strptime(task["due_date"], "%Y-%m-%d").date()
        days_left = (due_date - today).days
        is_overdue = days_left < 0
        is_today = days_left == 0

        if is_overdue:
            badge = "🚨 OVERDUE"
            color = "#dc3545"
        elif is_today:
            badge = "🎯 TODAY"
            color = "#fd7e14"
        elif days_left <= 2:
            badge = f"⚠️ Due in {days_left} day{'s' if days_left != 1 else ''}"
            color = "#ffc107"
        else:
            badge = f"✅ Due in {days_left} days"
            color = "#28a745"

        submission_status = task.get("submission_status", "not_submitted")
        grade = task.get("grade", "N/A")
        feedback = task.get("feedback", "No feedback yet")

        html += f"""
        <div style='
            background: white;
            border-radius: 12px;
            padding: 20px;
            box-shadow: 0 2px 8px rgba(0,0,0,0.08);
            border-left: 4px solid {color};
            transition: transform 0.2s;
        '>
            <div style='display: flex; justify-content: space-between; align-items: flex-start;'>
                <div>
                    <h3 style='margin: 0 0 8px 0; color: #212529;'>{task['title']}</h3>
                    <div style='color: #6c757d; margin-bottom: 8px;'>📚 {task['course']} · 👩‍🏫 {task['assigned_by']}</div>
                    <div style='color: #495057;'>📅 Due: {task['due_date']}</div>
                    <div style='color: #6c757d; font-size: 0.9em;'>{task.get('description', '')}</div>
                    
                    <div style='margin-top: 12px; padding: 8px; background: #f8f9fa; border-radius: 8px;'>
                        <div style='font-weight: bold; color: #212529;'>Submission Status: 
                            <span style='color: {'#28a745' if submission_status == 'submitted' else '#dc3545'};'>{submission_status.title()}</span>
                        </div>
                        <div>Grade: <strong>{grade}</strong></div>
                        <div>Feedback: <em>{feedback}</em></div>
                    </div>
                </div>
                <span style='
                    background: {color};
                    color: white;
                    padding: 6px 12px;
                    border-radius: 20px;
                    font-weight: bold;
                    font-size: 0.85em;
                    align-self: flex-start;
                '>{badge}</span>
            </div>
        </div>
        """

    html += "</div>"
    return html


def render_groups(groups: List[str]) -> str:
    if not groups:
        return """
        <div style='text-align: center; padding: 40px; color: #6c757d;'>
            <div style='font-size: 4em; margin-bottom: 16px;'>👥</div>
            <h3>You haven't joined any class groups yet.</h3>
            <p>Ask your teacher for a group code to join your class.</p>
        </div>
        """

    html = "<div style='display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 16px;'>"
    for group in groups:
        html += f"""
        <div style='
            background: #e3f2fd;
            border-radius: 12px;
            padding: 20px;
            text-align: center;
            box-shadow: 0 2px 4px rgba(0,0,0,0.05);
            border: 2px solid #2196f3;
        '>
            <div style='font-size: 1.5em; margin-bottom: 8px;'>🎒</div>
            <h3 style='margin: 0; color: #1976d2;'>{group}</h3>
            <p style='margin: 8px 0 0 0; color: #555; font-size: 0.9em;'>Class Group</p>
        </div>
        """
    html += "</div>"
    return html


def get_avatar_html(avatar_path: Optional[str], name: str) -> str:
    if avatar_path and os.path.exists(avatar_path):
        try:
            with open(avatar_path, "rb") as f:
                img_data = base64.b64encode(f.read()).decode()
            img_html = f'<img src="data:image/png;base64,{img_data}" style="width: 60px; height: 60px; border-radius: 50%; object-fit: cover; border: 2px solid #6e8efb;">'
        except:
            img_html = f'<div style="width: 60px; height: 60px; border-radius: 50%; background: #6e8efb; color: white; display: flex; align-items: center; justify-content: center; font-weight: bold; font-size: 1.2em; border: 2px solid #6e8efb;">{name[0].upper()}</div>'
    else:
        img_html = f'<div style="width: 60px; height: 60px; border-radius: 50%; background: #6e8efb; color: white; display: flex; align-items: center; justify-content: center; font-weight: bold; font-size: 1.2em; border: 2px solid #6e8efb;">{name[0].upper()}</div>'

    return img_html


def render_analytics() -> str:
    total_students = student_service.get_total_students()
    active_students = student_service.get_active_students()
    total_assignments = student_service.get_total_assignments()
    group_popularity = student_service.get_group_popularity()
    recent_activity = student_service.get_recent_activity()
    school_stats = school_service.get_stats()

    html = f"""
    <div style='display: grid; grid-template-columns: repeat(2, 1fr); gap: 24px; margin-bottom: 32px;'>
        <div style='background: white; padding: 20px; border-radius: 12px; box-shadow: 0 2px 8px rgba(0,0,0,0.08);'>
            <h3 style='color: #212529; margin: 0 0 16px 0;'>👥 Student Engagement</h3>
            <div style='display: grid; grid-template-columns: 1fr 1fr; gap: 16px;'>
                <div style='text-align: center; padding: 16px; background: #e3f2fd; border-radius: 8px;'>
                    <div style='font-size: 2em; color: #1976d2;'>{total_students}</div>
                    <div>Total Students</div>
                </div>
                <div style='text-align: center; padding: 16px; background: #e8f5e8; border-radius: 8px;'>
                    <div style='font-size: 2em; color: #2e7d32;'>{active_students}</div>
                    <div>Active Students</div>
                </div>
            </div>
        </div>
        
        <div style='background: white; padding: 20px; border-radius: 12px; box-shadow: 0 2px 8px rgba(0,0,0,0.08);'>
            <h3 style='color: #212529; margin: 0 0 16px 0;'>📊 Assignment Stats</h3>
            <div style='display: grid; grid-template-columns: 1fr 1fr; gap: 16px;'>
                <div style='text-align: center; padding: 16px; background: #fff3e0; border-radius: 8px;'>
                    <div style='font-size: 2em; color: #ef6c00;'>{total_assignments}</div>
                    <div>Total Assignments</div>
                </div>
                <div style='text-align: center; padding: 16px; background: #f3e5f5; border-radius: 8px;'>
                    <div style='font-size: 2em; color: #7b1fa2;'>N/A</div>
                    <div>Completion Rate</div>
                </div>
            </div>
        </div>
    </div>

    <div style='display: grid; grid-template-columns: 1fr 1fr; gap: 24px; margin-bottom: 32px;'>
        <div style='background: white; padding: 20px; border-radius: 12px; box-shadow: 0 2px 8px rgba(0,0,0,0.08);'>
            <h3 style='color: #212529; margin: 0 0 16px 0;'>📈 Popular Class Groups</h3>
            <div style='display: grid; gap: 12px;'>
    """

    sorted_groups = sorted(group_popularity.items(), key=lambda x: x[1], reverse=True)
    for group, count in sorted_groups[:5]:
        percentage = (count / total_students * 100) if total_students > 0 else 0
        html += f"""
                <div style='display: flex; align-items: center; gap: 12px;'>
                    <div style='width: 100%; background: #e0e0e0; height: 8px; border-radius: 4px;'>
                        <div style='width: {percentage}%; background: #6e8efb; height: 100%; border-radius: 4px;'></div>
                    </div>
                    <div style='min-width: 80px;'>{group}</div>
                    <div style='color: #6c757d;'>{count} students</div>
                </div>
        """

    html += """
            </div>
        </div>

        <div style='background: white; padding: 20px; border-radius: 12px; box-shadow: 0 2px 8px rgba(0,0,0,0.08);'>
            <h3 style='color: #212529; margin: 0 0 16px 0;'>📢 Announcement Stats</h3>
            <div style='display: grid; gap: 12px;'>
    """

    for ann in school_service.get_announcements()[:5]:
        views = ann.get("views", 0)
        max_views = max(a.get("views", 0) for a in school_service.get_announcements()) if school_service.get_announcements() else 1
        width_percent = (views / max_views * 100) if max_views > 0 else 0
        html += f"""
                <div style='display: flex; align-items: center; gap: 12px;'>
                    <div style='width: 100%; background: #e0e0e0; height: 8px; border-radius: 4px;'>
                        <div style='width: {width_percent}%; background: #ff7043; height: 100%; border-radius: 4px;'></div>
                    </div>
                    <div style='min-width: 120px; font-size: 0.9em;'>{ann['title'][:20]}...</div>
                    <div style='color: #6c757d;'>{views} views</div>
                </div>
        """

    html += """
            </div>
        </div>
    </div>

    <div style='background: white; padding: 20px; border-radius: 12px; box-shadow: 0 2px 8px rgba(0,0,0,0.08);'>
        <h3 style='color: #212529; margin: 0 0 16px 0;'>⏱️ Recent Activity</h3>
        <div style='display: grid; gap: 12px;'>
    """

    for activity in recent_activity:
        html += f"""
            <div style='padding: 12px; background: #f8f9fa; border-radius: 8px; display: flex; justify-content: space-between; align-items: center;'>
                <div>
                    <strong>{activity['student']}</strong> received assignment: <em>{activity['last_assignment']}</em>
                </div>
                <div style='color: #6c757d; font-size: 0.9em;'>{activity['date']}</div>
            </div>
        """

    html += """
        </div>
    </div>
    """

    return html

# ==================== STATE MANAGEMENT ====================

CURRENT_USER = "guest"
CURRENT_TEACHER = None
DARK_MODE = False

def login_student(username: str, password: str) -> tuple:
    global CURRENT_USER, DARK_MODE
    student_data = student_service.authenticate_student(username, password)
    if student_:  # ✅ FIXED: Added missing colon
        CURRENT_USER = username
        DARK_MODE = student_data["dark_mode"]
        chat_history = student_service.get_chat_history(username)
        files = student_service.get_files(username)
        assignments = student_service.get_assignments(username)
        groups = student_service.get_groups(username)
        avatar_html = get_avatar_html(student_data["avatar"], student_data["name"])
        welcome_msg = f"Welcome back, {student_data['name']}!"

        return (
            gr.update(visible=False),  # login_group
            gr.update(visible=True),   # main_app
            gr.update(value=chat_history),  # chatbot
            gr.update(value=files),    # files
            gr.update(value=render_assignments(assignments)),  # assignments_display
            gr.update(value=render_groups(groups)),  # groups_display
            gr.update(value=welcome_msg),  # login_status
            gr.update(value=student_data["name"], visible=True),  # user_display
            gr.update(visible=True),   # logout_btn
            gr.update(value=avatar_html),  # avatar_display
            gr.update(value="🌙 Light Mode" if DARK_MODE else "☀️ Dark Mode"),  # dark_mode_btn
            gr.update()  # css placeholder
        )
    return (
        gr.update(visible=True),
        gr.update(visible=False),
        gr.update(),
        gr.update(),
        gr.update(),
        gr.update(),
        gr.update(value="❌ Invalid username or password"),
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(),
        gr.update(),
        gr.update()
    )

def register_student(username: str, password: str, name: str) -> str:
    return student_service.register_student(username, password, name)

def logout_student() -> tuple:
    global CURRENT_USER, DARK_MODE
    CURRENT_USER = "guest"
    DARK_MODE = False
    return (
        gr.update(visible=True),
        gr.update(visible=False),
        gr.update(),
        gr.update(),
        gr.update(),
        gr.update(),
        gr.update(),
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(),
        gr.update(value="☀️ Dark Mode"),
        gr.update()
    )

def toggle_dark_mode() -> tuple:
    global DARK_MODE, CURRENT_USER
    DARK_MODE = not DARK_MODE
    if CURRENT_USER != "guest":
        student_service.update_dark_mode(CURRENT_USER, DARK_MODE)
    btn_text = "🌙 Light Mode" if DARK_MODE else "☀️ Dark Mode"
    return btn_text, gr.update()

def upload_avatar(file) -> str:
    if not file:
        return "❌ No file selected"
    if CURRENT_USER == "guest":
        return "❌ Please log in first"
    
    # Save avatar to disk
    filename = f"avatar_{CURRENT_USER}_{int(time.time())}{os.path.splitext(file.name)[1]}"
    filepath = os.path.join(UPLOADS_DIR, filename)
    shutil.copy(file.name, filepath)
    
    student_service.update_avatar(CURRENT_USER, filepath)
    return f"✅ Avatar updated!"

def join_group(group_code: str) -> tuple:
    if CURRENT_USER == "guest":
        return "❌ Please log in first.", gr.update()
    result = student_service.join_group(CURRENT_USER, group_code)
    groups = student_service.get_groups(CURRENT_USER)
    return result, gr.update(value=render_groups(groups))

def leave_group(group_code: str) -> tuple:
    if CURRENT_USER == "guest":
        return "❌ Please log in first.", gr.update()
    result = student_service.leave_group(CURRENT_USER, group_code)
    groups = student_service.get_groups(CURRENT_USER)
    return result, gr.update(value=render_groups(groups))

def upload_file_for_student(file) -> str:
    if not file:
        return "❌ No file selected"
    if CURRENT_USER == "guest":
        return "❌ Please log in first"
    
    # Save file to disk
    filename = f"{CURRENT_USER}_{int(time.time())}_{os.path.basename(file.name)}"
    filepath = os.path.join(UPLOADS_DIR, filename)
    shutil.copy(file.name, filepath)
    
    student_service.add_file(CURRENT_USER, file.name, filepath)
    return f"✅ Uploaded: {file.name}"

# ==================== TEACHER FUNCTIONS ====================

def teacher_login(username: str, password: str) -> tuple:
    global CURRENT_TEACHER
    teacher_data = admin_service.authenticate(username, password)
    if teacher_data and teacher_data["role"] == "teacher":
        CURRENT_TEACHER = teacher_data
        return (
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(value=f"✅ Welcome, {teacher_data['name']}!"),
            gr.update(visible=True)
        )
    elif teacher_data and teacher_data["role"] == "admin":
        CURRENT_TEACHER = teacher_data
        return (
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(value=f"✅ Welcome, Admin {teacher_data['name']}!"),
            gr.update(visible=True)
        )
    return (
        gr.update(visible=True),
        gr.update(visible=False),
        gr.update(value="❌ Invalid credentials or not a teacher account"),
        gr.update(visible=False)
    )

def teacher_logout():
    global CURRENT_TEACHER
    CURRENT_TEACHER = None
    return (
        gr.update(visible=True),
        gr.update(visible=False),
        gr.update(),
        gr.update(visible=False)
    )

def create_assignment(title: str, description: str, due_date: str, course: str) -> str:
    if not CURRENT_TEACHER:
        return "🔒 Please log in as teacher first."
    if not title or not due_date or not course:
        return "⚠️ Title, due date, and course are required."

    assignment_id = school_service.create_assignment(
        title=title,
        description=description,
        due_date=due_date,
        course=course,
        assigned_by=CURRENT_TEACHER["name"]
    )

    if assignment_id == -1:
        return "❌ Failed to create assignment. Please try again."

    conn = db_manager.get_connection()
    cursor = conn.cursor()
    cursor.execute('''
        SELECT COUNT(*) FROM student_sessions
        WHERE groups LIKE ?
    ''', (f'%"{course}"%',))
    recipients = cursor.fetchone()[0]
    conn.close()

    return f"✅ Assignment created! ID: {assignment_id}. Sent to {recipients} students in {course}."

# ==================== ASSIGNMENT SUBMISSION ====================

def submit_assignment(assignment_id: int, submission_text: str, submission_file) -> str:
    if CURRENT_USER == "guest":
        return "❌ Please log in first."
    if not submission_text.strip() and not submission_file:
        return "⚠️ Please provide text or upload a file."
    
    file_path = None
    if submission_file:
        filename = f"submission_{CURRENT_USER}_{assignment_id}_{int(time.time())}{os.path.splitext(submission_file.name)[1]}"
        file_path = os.path.join(UPLOADS_DIR, filename)
        shutil.copy(submission_file.name, file_path)
    
    result = student_service.submit_assignment(CURRENT_USER, assignment_id, submission_text, file_path)
    return result

def get_assignment_for_submission(assignment_id: int) -> Dict:
    assignments = student_service.get_assignments(CURRENT_USER)
    for assignment in assignments:
        if assignment["id"] == assignment_id:
            return assignment
    return None

# ==================== CALENDAR SYNC & PUSH NOTIFICATIONS ====================

def export_to_calendar(assignments: List[Dict]) -> str:
    """Simulate exporting assignments to Google Calendar."""
    if not assignments:
        return "📭 No assignments to export."
    
    events = []
    for task in assignments:
        events.append({
            "title": task["title"],
            "start": task["due_date"],
            "description": f"Course: {task['course']}\nAssigned by: {task['assigned_by']}"
        })
    
    # In real app, you'd use Google Calendar API here
    return f"✅ Exported {len(events)} assignments to calendar! (Simulated)"

def send_push_notification(message: str) -> str:
    """Simulate sending browser push notification."""
    # In real app, you'd use service workers + Push API
    return f"🔔 Notification sent: {message} (Simulated)"

def get_upcoming_deadlines() -> List[Dict]:
    """Get assignments due in next 3 days for notifications."""
    assignments = student_service.get_assignments(CURRENT_USER)
    today = datetime.today().date()
    upcoming = []
    for task in assignments:
        due_date = datetime.strptime(task["due_date"], "%Y-%m-%d").date()
        days_left = (due_date - today).days
        if 0 <= days_left <= 3:
            upcoming.append(task)
    return upcoming

# ==================== VOICE INPUT ====================

VOICE_JS = """
async function startVoiceInput() {
    if (!('webkitSpeechRecognition' in window) && !('SpeechRecognition' in window)) {
        alert('🎙️ Voice input is not supported in your browser. Try Chrome or Edge.');
        return '';
    }

    const recognition = new (window.SpeechRecognition || window.webkitSpeechRecognition)();
    recognition.lang = 'en-US';
    recognition.interimResults = false;
    recognition.maxAlternatives = 1;

    return new Promise((resolve) => {
        recognition.start();
        recognition.onresult = (event) => {
            const transcript = event.results[0][0].transcript;
            resolve(transcript);
        };
        recognition.onerror = (event) => {
            alert('🎙️ Error: ' + event.error);
            resolve('');
        };
    });
}
"""

def voice_input_handler() -> str:
    return ""

# ==================== CUSTOM CSS ====================

CUSTOM_CSS = """
.gradio-container {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
    max-width: 1200px;
    margin: 0 auto;
    padding: 16px;
    transition: background-color 0.3s, color 0.3s;
}

.dark-mode {
    --background-fill-primary: #1e1e1e !important;
    --background-fill-secondary: #2d2d2d !important;
    --text-color: #f0f0f0 !important;
    --button-primary-background-fill: #5a5a5a !important;
    --button-secondary-background-fill: #3a3a3a !important;
    --input-background-fill: #2d2d2d !important;
    --input-border-color: #444 !important;
}

.primary {
    background: linear-gradient(135deg, #6e8efb, #a777e3) !important;
    border: none !important;
    color: white !important;
}

.chatbot-container {
    background: #f8f9fa !important;
    border-radius: 16px !important;
}

.dark-mode .chatbot-container {
    background: #2d2d2d !important;
}

.user, .bot {
    border-radius: 18px !important;
    padding: 12px 16px !important;
}

.file-upload {
    border: 2px dashed #6e8efb !important;
    border-radius: 12px !important;
    background: #f8f9ff !important;
}

.dark-mode .file-upload {
    background: #2a2a2a !important;
    border-color: #5a5a5a !important;
}

.submission-card {
    background: #f8f9fa;
    border-radius: 12px;
    padding: 16px;
    margin: 8px 0;
    border-left: 4px solid #007bff;
}

.notification {
    background: #fff3cd;
    border-left: 4px solid #ffc107;
    padding: 12px;
    border-radius: 8px;
    margin: 8px 0;
}
"""

# ==================== BUILD UI ====================

with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="indigo")) as demo:
    with gr.Row():
        gr.Markdown("# 🎓 ThutoAI — Your AI School Assistant")
        dark_mode_btn = gr.Button("☀️ Dark Mode", variant="secondary")

    with gr.Group() as login_group:
        gr.Markdown("### 🔐 Student Login")
        with gr.Row():
            login_username = gr.Textbox(label="Username", scale=3)
            login_password = gr.Textbox(label="Password", type="password", scale=3)
        login_btn = gr.Button("🔓 Login", variant="primary")
        register_btn = gr.Button("📝 Register", variant="secondary")
        login_status = gr.Textbox(label="Status", interactive=False)

        with gr.Accordion("📝 Register New Account", open=False):
            reg_username = gr.Textbox(label="Username")
            reg_password = gr.Textbox(label="Password", type="password")
            reg_name = gr.Textbox(label="Full Name")
            reg_btn = gr.Button("Create Account")
            reg_status = gr.Textbox(label="Registration Status")

    with gr.Group(visible=False) as main_app:
        with gr.Row():
            avatar_display = gr.HTML()
            user_display = gr.Textbox(label="Logged in as", interactive=False, visible=True)
            logout_btn = gr.Button("⬅️ Logout", variant="secondary")

        with gr.Tabs():
            with gr.Tab("📢 Announcements"):
                gr.Markdown("### Filter by Course or Subject")
                with gr.Row():
                    course_filter = gr.Dropdown(
                        choices=["All", "Mathematics", "Science", "English", "History", "General"],
                        value="All",
                        label="Select Course",
                        scale=3
                    )
                    refresh_btn = gr.Button("🔄 Refresh", variant="secondary", scale=1)
                announcements_html = gr.HTML()
                course_filter.change(fn=render_announcements, inputs=course_filter, outputs=announcements_html, api_name="filter_announcements")
                refresh_btn.click(fn=render_announcements, inputs=course_filter, outputs=announcements_html, api_name="refresh_announcements")
                demo.load(fn=render_announcements, inputs=course_filter, outputs=announcements_html, api_name="load_announcements")

            with gr.Tab("💬 Ask ThutoAI"):
                gr.Markdown("### 💡 Ask me anything — I'm here to help!")
                chatbot = gr.Chatbot(
                    height=480,
                    type='messages'
                )
                with gr.Row():
                    msg = gr.Textbox(
                        label="Type your question",
                        placeholder="E.g., How do I prepare for the Math exam?",
                        scale=7
                    )
                    voice_btn = gr.Button("🎙️ Speak", variant="secondary", scale=1)
                    submit_btn = gr.Button("➤ Send", variant="primary", scale=1)
                clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")

                def respond(message, chat_history):
                    for updated_history, _ in ai_chat(message, chat_history, CURRENT_USER):
                        yield updated_history, ""

                msg.submit(respond, [msg, chatbot], [msg, chatbot])
                submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
                clear_btn.click(lambda: None, None, chatbot)

                voice_btn.click(
                    fn=voice_input_handler,
                    inputs=None,
                    outputs=msg,
                    js=VOICE_JS + "return startVoiceInput();"
                )

            with gr.Tab("📅 Assignments"):
                gr.Markdown("### 📌 Your Upcoming Assignments & Exams")
                assignments_display = gr.HTML()
                demo.load(
                    fn=lambda: render_assignments(student_service.get_assignments(CURRENT_USER)) if CURRENT_USER != "guest" else "",
                    inputs=None,
                    outputs=assignments_display
                )

                # Assignment submission section
                with gr.Accordion("📤 Submit Assignment", open=False):
                    gr.Markdown("### Submit your work for grading")
                    assignment_dropdown = gr.Dropdown(
                        choices=[],
                        label="Select Assignment to Submit",
                        interactive=True
                    )
                    submission_text = gr.Textbox(
                        label="Your Answer / Comments",
                        placeholder="Type your response here...",
                        lines=4
                    )
                    submission_file = gr.File(label="Upload your work (PDF, DOC, TXT, etc.)")
                    submit_assignment_btn = gr.Button("📤 Submit Assignment", variant="primary")
                    submission_result = gr.Textbox(label="Submission Status")

                    def update_assignment_dropdown():
                        assignments = student_service.get_assignments(CURRENT_USER)
                        choices = [(f"{a['title']} (Due: {a['due_date']})", a["id"]) for a in assignments if a.get("submission_status") != "submitted"]
                        return gr.update(choices=choices)

                    def handle_submission(assignment_id, text, file):
                        if not assignment_id:
                            return "⚠️ Please select an assignment."
                        return submit_assignment(assignment_id, text, file)

                    demo.load(fn=update_assignment_dropdown, inputs=None, outputs=assignment_dropdown)
                    submit_assignment_btn.click(
                        fn=handle_submission,
                        inputs=[assignment_dropdown, submission_text, submission_file],
                        outputs=submission_result
                    )

                # Calendar sync section
                with gr.Accordion("🗓️ Calendar Sync", open=False):
                    gr.Markdown("### Export assignments to your calendar")
                    calendar_btn = gr.Button("📅 Export to Google Calendar", variant="primary")
                    calendar_result = gr.Textbox(label="Export Status")
                    
                    def export_calendar():
                        assignments = student_service.get_assignments(CURRENT_USER)
                        return export_to_calendar(assignments)
                    
                    calendar_btn.click(fn=export_calendar, inputs=None, outputs=calendar_result)

                # Push notifications section
                with gr.Accordion("🔔 Push Notifications", open=False):
                    gr.Markdown("### Get notified about upcoming deadlines")
                    notify_btn = gr.Button("🔔 Send Test Notification", variant="primary")
                    notify_result = gr.Textbox(label="Notification Status")
                    
                    def send_test_notification():
                        deadlines = get_upcoming_deadlines()
                        if deadlines:
                            msg = f"You have {len(deadlines)} assignment(s) due soon!"
                        else:
                            msg = "No upcoming deadlines found."
                        return send_push_notification(msg)
                    
                    notify_btn.click(fn=send_test_notification, inputs=None, outputs=notify_result)

            with gr.Tab("👥 Class Groups"):
                gr.Markdown("### 🎒 Join Your Class Groups")
                with gr.Row():
                    group_code_input = gr.Textbox(label="Group Code (e.g., MATH10A)", scale=3)
                    join_btn = gr.Button("➕ Join Group", variant="primary", scale=1)
                    leave_btn = gr.Button("➖ Leave Group", variant="secondary", scale=1)
                group_status = gr.Textbox(label="Status")
                groups_display = gr.HTML()
                join_btn.click(
                    fn=join_group,
                    inputs=group_code_input,
                    outputs=[group_status, groups_display]
                )
                leave_btn.click(
                    fn=leave_group,
                    inputs=group_code_input,
                    outputs=[group_status, groups_display]
                )
                demo.load(
                    fn=lambda: render_groups(student_service.get_groups(CURRENT_USER)) if CURRENT_USER != "guest" else "",
                    inputs=None,
                    outputs=groups_display
                )

            with gr.Tab("📂 My Files"):
                gr.Markdown("### 📁 Upload & Manage Study Materials")
                with gr.Row():
                    file_input = gr.File(label="Drag & drop or click to upload", elem_classes=["file-upload"])
                    upload_btn = gr.Button("📤 Upload", variant="primary")
                upload_status = gr.Textbox(label="Status")
                file_list = gr.JSON(label="Your Files")

                upload_btn.click(
                    fn=upload_file_for_student,
                    inputs=file_input,
                    outputs=upload_status
                )
                demo.load(
                    fn=lambda: student_service.get_files(CURRENT_USER) if CURRENT_USER != "guest" else [],
                    inputs=None,
                    outputs=file_list
                )

            with gr.Tab("🖼️ Profile"):
                gr.Markdown("### 🖼️ Update Your Profile Picture")
                with gr.Row():
                    avatar_input = gr.File(label="Choose an image (PNG, JPG)", file_types=["image"])
                    upload_avatar_btn = gr.Button("📤 Upload Avatar", variant="primary")
                avatar_status = gr.Textbox(label="Status")
                upload_avatar_btn.click(
                    fn=upload_avatar,
                    inputs=avatar_input,
                    outputs=avatar_status
                )
                demo.load(
                    fn=lambda: get_avatar_html(
                        student_service.authenticate_student(CURRENT_USER, "")["avatar"] if CURRENT_USER != "guest" else None,
                        student_service.authenticate_student(CURRENT_USER, "")["name"] if CURRENT_USER != "guest" else "Guest"
                    ) if CURRENT_USER != "guest" else "",
                    inputs=None,
                    outputs=avatar_display
                )

            with gr.Tab("👩‍🏫 Teacher Panel"):
                gr.Markdown("### 🔑 Teacher Login (for assignment creation)")

                with gr.Group() as teacher_login_group:
                    teacher_username = gr.Textbox(label="Teacher Username")
                    teacher_password = gr.Textbox(label="Password", type="password")
                    teacher_login_btn = gr.Button("🔓 Login", variant="primary")
                    teacher_status = gr.Textbox(label="Status")

                with gr.Group(visible=False) as teacher_dashboard:
                    gr.Markdown("### ✍️ Create New Assignment")
                    assignment_title = gr.Textbox(label="Assignment Title", placeholder="e.g., Chapter 5 Quiz")
                    assignment_description = gr.Textbox(label="Description", placeholder="Instructions for students...", lines=2)
                    assignment_due_date = gr.Textbox(label="Due Date (YYYY-MM-DD)", placeholder="2025-05-01")
                    assignment_course = gr.Dropdown(
                        choices=["MATH10A", "MATH10B", "SCI11A", "SCI11B", "ENG10A", "ENG10B", "HIST9A", "HIST9B"],
                        label="Assign to Class Group",
                        value="MATH10A"
                    )
                    create_assignment_btn = gr.Button("📬 Create Assignment", variant="primary")
                    assignment_result = gr.Textbox(label="Result")

                    assignment_priority = gr.Radio(
                        ["low", "normal", "high"],
                        label="Priority",
                        value="normal"
                    )

                    teacher_logout_btn = gr.Button("⬅️ Logout", variant="secondary")

                teacher_login_btn.click(
                    fn=teacher_login,
                    inputs=[teacher_username, teacher_password],
                    outputs=[teacher_login_group, teacher_dashboard, teacher_status, create_assignment_btn]
                )
                teacher_logout_btn.click(
                    fn=teacher_logout,
                    inputs=None,
                    outputs=[teacher_login_group, teacher_dashboard, teacher_status, assignment_result]
                )
                create_assignment_btn.click(
                    fn=create_assignment,
                    inputs=[assignment_title, assignment_description, assignment_due_date, assignment_course],
                    outputs=assignment_result
                )

                # View submissions
                with gr.Accordion("📋 View Submissions", open=False):
                    gr.Markdown("### View and Grade Student Submissions")
                    submissions_assignment_dropdown = gr.Dropdown(
                        choices=[],
                        label="Select Assignment to View Submissions",
                        interactive=True
                    )
                    submissions_display = gr.JSON(label="Submissions")
                    grade_input = gr.Textbox(label="Grade (e.g., A+, 85%, Pass)")
                    feedback_input = gr.Textbox(label="Feedback", lines=3)
                    grade_btn = gr.Button("✅ Grade Submission", variant="primary")
                    grade_result = gr.Textbox(label="Grading Status")

                    def update_submissions_dropdown():
                        # Get assignments created by this teacher
                        conn = db_manager.get_connection()
                        cursor = conn.cursor()
                        cursor.execute('''
                            SELECT DISTINCT a.id, a.title
                            FROM assignment_submissions s
                            JOIN student_sessions ss ON s.student_username = json_extract(ss.assignments, '$[0].assigned_by')
                            JOIN json_each(ss.assignments) j
                            JOIN json_tree(j.value) t ON t.key = 'id'
                            JOIN announcements a ON t.value = a.id
                            WHERE a.posted_by = ?
                        ''', (CURRENT_TEACHER["name"],))
                        results = cursor.fetchall()
                        choices = [(f"{row[1]} (ID: {row[0]})", row[0]) for row in results]
                        return gr.update(choices=choices)

                    def load_submissions(assignment_id):
                        if not assignment_id:
                            return []
                        return student_service.get_assignment_submissions(assignment_id)

                    def grade_submission(submission_id, grade, feedback):
                        if not submission_id or not grade:
                            return "⚠️ Please select submission and enter grade."
                        return admin_service.grade_assignment(submission_id, grade, feedback)

                    submissions_assignment_dropdown.change(
                        fn=load_submissions,
                        inputs=submissions_assignment_dropdown,
                        outputs=submissions_display
                    )

                    grade_btn.click(
                        fn=grade_submission,
                        inputs=[submissions_display, grade_input, feedback_input],
                        outputs=grade_result
                    )

            with gr.Tab("📊 Admin Analytics"):
                gr.Markdown("### 📈 School Analytics Dashboard")
                analytics_display = gr.HTML()
                refresh_analytics_btn = gr.Button("🔄 Refresh Analytics")
                refresh_analytics_btn.click(
                    fn=render_analytics,
                    inputs=None,
                    outputs=analytics_display
                )
                demo.load(
                    fn=render_analytics,
                    inputs=None,
                    outputs=analytics_display
                )

 def login_function():
    # This function runs when the login button is clicked
    return (
        gr.update(visible=False),  # login_group
        gr.update(visible=True),   # main_app
        gr.update(value=[]),       # chatbot
        gr.update(visible=True),   # files
        gr.update(visible=True),   # assignments_display
        gr.update(visible=True),   # groups_display
        gr.update(value="✅ Logged in successfully!", visible=True),  # login_status
        gr.update(visible=True),   # user_display
        gr.update(visible=True),   # logout_btn
        gr.update(visible=True),   # avatar_display
        gr.update(visible=True),   # dark_mode_btn
        gr.update()                # css placeholder 
     ]
    )

       login_btn.click(
        fn=login_function,
        inputs=[],
        outputs=[
        login_group,
        main_app,
        chatbot,
        files,
        assignments_display,
        groups_display,
        login_status,
        user_display,
        logout_btn,
        avatar_display,
        dark_mode_btn,
        css_placeholder
    ]
)
    register_btn.click(
        fn=register_student,
        inputs=[reg_username, reg_password, reg_name],
        outputs=gr.update()
    )

    dark_mode_btn.click(
        fn=toggle_dark_mode,
        inputs=None,
        outputs=[dark_mode_btn, gr.update()]
    )

if __name__ == "__main__":
    demo.launch()