import gradio as gr
from transformers import pipeline
import PyPDF2
import re
import os
import io
import random
import time
import hashlib
import secrets
import json
import sqlite3
from datetime import datetime, timedelta
from pathlib import Path
from groq import Groq

# ==================== CONFIGURATION ====================
try:
    from google import genai
    from google.genai import types
    GENAI_AVAILABLE = True
except ImportError:
    GENAI_AVAILABLE = False
    genai = None
    types = None

# Environment variables
hf_token = os.environ.get("HF_TOKEN")
gemini_key = os.environ.get("GEMINI_API_KEY")
groq_key = os.environ.get("GROQ_API_KEY")
dev_password = os.environ.get("DEV_PASSWORD", "dev123")

# Database path (persistent storage)
DATA_DIR = Path("/tmp/data") if os.path.exists("/tmp") else Path("./data")
DATA_DIR.mkdir(exist_ok=True)
DB_PATH = DATA_DIR / "student_ai_suite.db"

# API Clients
gemini_client = None
groq_client = None

if gemini_key and GENAI_AVAILABLE:
    try:
        gemini_client = genai.Client(api_key=gemini_key)
    except:
        try:
            import google.generativeai as old_genai
            old_genai.configure(api_key=gemini_key)
            gemini_client = old_genai
        except:
            pass

if groq_key:
    try:
        groq_client = Groq(api_key=groq_key)
    except:
        pass

# Lazy load summarizer
summarizer = None

def load_summarizer():
    global summarizer
    if summarizer is None:
        try:
            summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=-1)
        except:
            pass
    return summarizer

# ==================== DATABASE MANAGER ====================
class DatabaseManager:
    def __init__(self, db_path):
        self.db_path = db_path
        self.init_database()
    
    def get_connection(self):
        """Get database connection with row factory"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn
    
    def init_database(self):
        """Initialize database tables"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        # Users table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                name TEXT NOT NULL,
                email TEXT,
                role TEXT DEFAULT 'user',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                last_login TIMESTAMP,
                is_active BOOLEAN DEFAULT 1,
                is_dev BOOLEAN DEFAULT 0
            )
        """)
        
        # Sessions table for tracking active sessions
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                session_token TEXT UNIQUE NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                expires_at TIMESTAMP NOT NULL,
                ip_address TEXT,
                user_agent TEXT,
                is_active BOOLEAN DEFAULT 1,
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
            )
        """)
        
        # Activity logs (optional - for analytics)
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS activity_logs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                action TEXT NOT NULL,
                tool_used TEXT,
                input_summary TEXT,
                output_summary TEXT,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                execution_time_ms INTEGER,
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE SET NULL
            )
        """)
        
        # User preferences
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS user_preferences (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER UNIQUE NOT NULL,
                default_essay_type TEXT DEFAULT 'Argumentative',
                default_essay_tone TEXT DEFAULT 'Academic',
                default_word_count INTEGER DEFAULT 500,
                default_quiz_questions INTEGER DEFAULT 5,
                default_quiz_time INTEGER DEFAULT 2,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
            )
        """)
        
        conn.commit()
        conn.close()
        print(f"✅ Database initialized at {self.db_path}")
    
    # ==================== USER OPERATIONS ====================
    
    def create_user(self, username, password_hash, name, email, role="user", is_dev=False):
        """Create new user"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            cursor.execute("""
                INSERT INTO users (username, password_hash, name, email, role, is_dev)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (username, password_hash, name, email, role, is_dev))
            
            user_id = cursor.lastrowid
            
            # Create default preferences
            cursor.execute("""
                INSERT INTO user_preferences (user_id) VALUES (?)
            """, (user_id,))
            
            conn.commit()
            conn.close()
            return True, user_id
            
        except sqlite3.IntegrityError:
            return False, "Username already exists"
        except Exception as e:
            return False, str(e)
    
    def get_user_by_username(self, username):
        """Get user by username"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT * FROM users WHERE username = ? AND is_active = 1
        """, (username,))
        
        user = cursor.fetchone()
        conn.close()
        
        if user:
            return dict(user)
        return None
    
    def get_user_by_id(self, user_id):
        """Get user by ID"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT * FROM users WHERE id = ? AND is_active = 1
        """, (user_id,))
        
        user = cursor.fetchone()
        conn.close()
        
        if user:
            return dict(user)
        return None
    
    def update_last_login(self, user_id):
        """Update last login timestamp"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            UPDATE users SET last_login = CURRENT_TIMESTAMP WHERE id = ?
        """, (user_id,))
        
        conn.commit()
        conn.close()
    
    # ==================== SESSION OPERATIONS ====================
    
    def create_session(self, user_id, session_token, expires_at, ip_address=None, user_agent=None):
        """Create new session"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        # Invalidate old sessions for this user
        cursor.execute("""
            UPDATE sessions SET is_active = 0 WHERE user_id = ?
        """, (user_id,))
        
        cursor.execute("""
            INSERT INTO sessions (user_id, session_token, expires_at, ip_address, user_agent)
            VALUES (?, ?, ?, ?, ?)
        """, (user_id, session_token, expires_at, ip_address, user_agent))
        
        conn.commit()
        conn.close()
        return True
    
    def get_session(self, session_token):
        """Validate and get session"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT s.*, u.username, u.name, u.role, u.is_dev
            FROM sessions s
            JOIN users u ON s.user_id = u.id
            WHERE s.session_token = ? 
            AND s.is_active = 1 
            AND s.expires_at > CURRENT_TIMESTAMP
            AND u.is_active = 1
        """, (session_token,))
        
        session = cursor.fetchone()
        conn.close()
        
        if session:
            return dict(session)
        return None
    
    def invalidate_session(self, session_token):
        """Logout - invalidate session"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            UPDATE sessions SET is_active = 0 WHERE session_token = ?
        """, (session_token,))
        
        conn.commit()
        conn.close()
        return True
    
    def cleanup_expired_sessions(self):
        """Remove expired sessions"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            UPDATE sessions SET is_active = 0 WHERE expires_at < CURRENT_TIMESTAMP
        """)
        
        conn.commit()
        conn.close()
    
    # ==================== PREFERENCES ====================
    
    def get_user_preferences(self, user_id):
        """Get user preferences"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT * FROM user_preferences WHERE user_id = ?
        """, (user_id,))
        
        prefs = cursor.fetchone()
        conn.close()
        
        if prefs:
            return dict(prefs)
        return None
    
    def update_preferences(self, user_id, **kwargs):
        """Update user preferences"""
        allowed_fields = ['default_essay_type', 'default_essay_tone', 'default_word_count',
                         'default_quiz_questions', 'default_quiz_time']
        
        updates = {k: v for k, v in kwargs.items() if k in allowed_fields}
        
        if not updates:
            return False
        
        set_clause = ", ".join([f"{k} = ?" for k in updates.keys()])
        values = list(updates.values()) + [user_id]
        
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute(f"""
            UPDATE user_preferences 
            SET {set_clause}, updated_at = CURRENT_TIMESTAMP 
            WHERE user_id = ?
        """, values)
        
        conn.commit()
        conn.close()
        return True
    
    # ==================== ACTIVITY LOGGING ====================
    
    def log_activity(self, user_id, action, tool_used=None, input_summary=None, 
                    output_summary=None, execution_time_ms=None):
        """Log user activity (optional analytics)"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            INSERT INTO activity_logs 
            (user_id, action, tool_used, input_summary, output_summary, execution_time_ms)
            VALUES (?, ?, ?, ?, ?, ?)
        """, (user_id, action, tool_used, input_summary, output_summary, execution_time_ms))
        
        conn.commit()
        conn.close()
    
    def get_user_stats(self, user_id):
        """Get user activity statistics"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT 
                COUNT(*) as total_activities,
                COUNT(DISTINCT tool_used) as tools_used,
                MAX(timestamp) as last_activity
            FROM activity_logs 
            WHERE user_id = ?
        """, (user_id,))
        
        stats = cursor.fetchone()
        conn.close()
        
        return dict(stats) if stats else None
    
    # ==================== ADMIN OPERATIONS ====================
    
    def get_all_users(self):
        """Get all users (admin only)"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT id, username, name, email, role, created_at, last_login, is_active, is_dev
            FROM users
            ORDER BY created_at DESC
        """)
        
        users = [dict(row) for row in cursor.fetchall()]
        conn.close()
        return users
    
    def get_database_stats(self):
        """Get database statistics"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        stats = {}
        
        cursor.execute("SELECT COUNT(*) FROM users WHERE is_dev = 0")
        stats['total_users'] = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM sessions WHERE is_active = 1")
        stats['active_sessions'] = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM activity_logs")
        stats['total_activities'] = cursor.fetchone()[0]
        
        cursor.execute("""
            SELECT tool_used, COUNT(*) as count 
            FROM activity_logs 
            GROUP BY tool_used 
            ORDER BY count DESC
        """)
        stats['tool_usage'] = [dict(row) for row in cursor.fetchall()]
        
        conn.close()
        return stats

# Initialize database
db = DatabaseManager(DB_PATH)

# ==================== USER AUTHENTICATION ====================
class UserAuth:
    def __init__(self, database):
        self.db = database
        self.failed_attempts = {}  # Still in-memory for brute force protection
        self._init_dev_account()
    
    def _hash(self, password):
        return hashlib.sha256(password.encode()).hexdigest()
    
    def _init_dev_account(self):
        """Create hidden developer account"""
        if dev_password:
            # Check if dev exists
            existing = self.db.get_user_by_username("admin")
            if not existing:
                self.db.create_user(
                    username="admin",
                    password_hash=self._hash(dev_password),
                    name="Developer",
                    email="dev@local",
                    role="developer",
                    is_dev=True
                )
                print("✅ Developer account created")
    
    def register(self, username, password, name, email):
        """Public user registration with database"""
        # Validation
        if not all([username, password, name]):
            return False, "All fields are required"
        
        if len(username) < 3:
            return False, "Username must be at least 3 characters"
        
        if len(password) < 6:
            return False, "Password must be at least 6 characters"
        
        # Create user
        password_hash = self._hash(password)
        success, result = self.db.create_user(username, password_hash, name, email)
        
        if success:
            return True, "Account created successfully!"
        else:
            return False, result  # Error message
    
    def login(self, username, password):
        """Login with database validation"""
        if not username or not password:
            return None, "Please enter both username and password"
        
        # Brute force check
        if username in self.failed_attempts:
            if self.failed_attempts[username]["count"] >= 5:
                last = self.failed_attempts[username]["last_attempt"]
                if datetime.now() - last < timedelta(minutes=15):
                    return None, "Too many failed attempts. Try again in 15 minutes."
                else:
                    del self.failed_attempts[username]
        
        # Get user from database
        user = self.db.get_user_by_username(username)
        
        if not user:
            self._record_failed_attempt(username)
            return None, "Invalid username or password"
        
        # Verify password
        password_hash = self._hash(password)
        if user['password_hash'] != password_hash:
            self._record_failed_attempt(username)
            return None, "Invalid username or password"
        
        # Success - update last login
        self.db.update_last_login(user['id'])
        
        if username in self.failed_attempts:
            del self.failed_attempts[username]
        
        # Create session
        session_token = secrets.token_urlsafe(32)
        expires_at = datetime.now() + timedelta(hours=24)
        
        self.db.create_session(
            user_id=user['id'],
            session_token=session_token,
            expires_at=expires_at
        )
        
        return session_token, {
            'user_id': user['id'],
            'username': user['username'],
            'name': user['name'],
            'role': user['role'],
            'is_dev': user['is_dev']
        }
    
    def _record_failed_attempt(self, username):
        if username not in self.failed_attempts:
            self.failed_attempts[username] = {"count": 0, "last_attempt": datetime.now()}
        self.failed_attempts[username]["count"] += 1
        self.failed_attempts[username]["last_attempt"] = datetime.now()
    
    def validate_session(self, session_token):
        """Validate session from database"""
        if not session_token:
            return None
        
        session = self.db.get_session(session_token)
        return session
    
    def logout(self, session_token):
        """Logout - invalidate in database"""
        if session_token:
            self.db.invalidate_session(session_token)
        return True
    
    def get_user_count(self):
        """Get count of non-dev users"""
        stats = self.db.get_database_stats()
        return stats.get('total_users', 0)

# Initialize auth with database
auth_system = UserAuth(db)

# ==================== [ALL YOUR EXISTING FUNCTIONS REMAIN THE SAME] ====================
# ... (Include all your previous functions: extract_text_from_pdf, summarize_with_gemini, etc.)

def extract_text_from_pdf(pdf_file):
    if pdf_file is None:
        return None, "Please upload a PDF file."
    
    try:
        if isinstance(pdf_file, str):
            with open(pdf_file, 'rb') as f:
                pdf_reader = PyPDF2.PdfReader(f)
                text = ""
                for page in pdf_reader.pages:
                    page_text = page.extract_text()
                    if page_text:
                        text += page_text + "\n"
        else:
            if hasattr(pdf_file, 'read'):
                pdf_bytes = pdf_file.read()
                if hasattr(pdf_file, 'seek'):
                    pdf_file.seek(0)
            else:
                pdf_bytes = pdf_file
            
            if isinstance(pdf_bytes, bytes):
                pdf_stream = io.BytesIO(pdf_bytes)
            else:
                pdf_stream = io.BytesIO(pdf_bytes.encode() if isinstance(pdf_bytes, str) else pdf_bytes)
            
            pdf_reader = PyPDF2.PdfReader(pdf_stream)
            text = ""
            for page in pdf_reader.pages:
                page_text = page.extract_text()
                if page_text:
                    text += page_text + "\n"
        
        text = re.sub(r'\s+', ' ', text).strip()
        
        if len(text) < 50:
            return None, "Could not extract text. PDF may be image-based or scanned."
        
        return text, None
        
    except Exception as e:
        return None, f"Error reading PDF: {str(e)}"

def summarize_with_gemini(text, max_length, min_length):
    if not gemini_client:
        return None
    
    try:
        if hasattr(gemini_client, 'models'):
            prompt = f"Summarize the following text in {min_length}-{max_length} words:\n\n{text[:15000]}"
            try:
                response = gemini_client.models.generate_content(
                    model="gemini-2.5-flash",
                    contents=prompt
                )
                return response.text
            except:
                if hasattr(gemini_client, 'GenerativeModel'):
                    model = gemini_client.GenerativeModel('gemini-2.5-flash')
                    response = model.generate_content(prompt)
                    return response.text
        elif hasattr(gemini_client, 'GenerativeModel'):
            model = gemini_client.GenerativeModel('gemini-2.5-flash')
            prompt = f"Summarize the following text in {min_length}-{max_length} words:\n\n{text[:15000]}"
            response = model.generate_content(prompt)
            return response.text
    except Exception as e:
        print(f"Gemini error: {e}")
    
    return None

def summarize_pdf(pdf_file, max_length, min_length):
    text, error = extract_text_from_pdf(pdf_file)
    if error:
        return error
    
    gemini_result = summarize_with_gemini(text, max_length, min_length)
    if gemini_result:
        return gemini_result
    
    summ = load_summarizer()
    if summ:
        try:
            result = summ(text[:3500], max_length=max_length, min_length=min_length, do_sample=False)
            return result[0]['summary_text']
        except Exception as e:
            return f"Error: {str(e)}"
    
    return "Error: No summarization available"

def generate_essay_with_gemini(prompt, essay_type, word_count, tone):
    if not gemini_client:
        return None
    
    try:
        full_prompt = f"""Write a {essay_type} essay in {tone} tone (~{word_count} words).
Topic: {prompt}
Requirements: Engaging intro, structured body, strong conclusion."""
        
        if hasattr(gemini_client, 'models'):
            response = gemini_client.models.generate_content(
                model="gemini-2.5-flash",
                contents=full_prompt
            )
            essay = response.text.strip()
        else:
            model = gemini_client.GenerativeModel('gemini-2.5-flash')
            response = model.generate_content(full_prompt)
            essay = response.text.strip()
        
        word_count_actual = len(essay.split())
        return f"""# {essay_type} Essay: {prompt[:50]}{'...' if len(prompt) > 50 else ''}

{essay}

---
*~{word_count_actual} words | {tone} tone*"""
        
    except Exception as e:
        print(f"Essay error: {e}")
        return None

def generate_essay(prompt, essay_type, word_count, tone):
    if not prompt or len(prompt.strip()) < 10:
        return "Please provide a detailed prompt (at least 10 characters)."
    
    result = generate_essay_with_gemini(prompt, essay_type, word_count, tone)
    if result:
        return result
    
    return "❌ Essay generation failed. Please check Gemini API configuration."

def summarize_text(text, max_length, min_length):
    if len(text.strip()) < 100:
        return "Please provide at least 100 characters."
    
    gemini_result = summarize_with_gemini(text, max_length, min_length)
    if gemini_result:
        return gemini_result
    
    summ = load_summarizer()
    if summ:
        try:
            result = summ(text[:3500], max_length=max_length, min_length=min_length, do_sample=False)
            return result[0]['summary_text']
        except Exception as e:
            return f"Error: {str(e)}"
    
    return "Error: No summarization available"

# Quiz Functions
def extract_sentences(text):
    sentences = re.split(r'[.!?]', text)
    return [s.strip() for s in sentences if len(s.split()) > 6]

def create_quiz(text, num_questions):
    sentences = extract_sentences(text)
    if len(sentences) < num_questions:
        num_questions = len(sentences)
    
    selected = random.sample(sentences, num_questions)
    quiz_data = []
    
    for sentence in selected:
        words = sentence.split()
        if len(words) < 5:
            continue
        keyword = random.choice(words[2:-2])
        question = sentence.replace(keyword, "_____")
        all_words = list(set(text.split()))
        wrong = random.sample([w for w in all_words if w != keyword], min(3, len(all_words)))
        options = wrong + [keyword]
        random.shuffle(options)
        
        quiz_data.append({
            "question": question,
            "options": options,
            "answer": keyword
        })
    
    return quiz_data

def start_quiz(text, num_questions, timer_minutes):
    if not text.strip():
        return "⚠️ Please enter study material.", gr.update(choices=[], visible=False), "", None, 0, None, "", gr.update(visible=False)
    
    quiz = create_quiz(text, num_questions)
    if not quiz:
        return "⚠️ Could not generate quiz. Please provide more text.", gr.update(choices=[], visible=False), "", None, 0, None, "", gr.update(visible=False)
    
    end_time = time.time() + (timer_minutes * 60)
    return show_question(quiz, 0, 0, end_time)

def show_question(quiz, index, score, end_time):
    if time.time() > end_time:
        return finish_quiz(quiz, score, len(quiz))
    
    if index >= len(quiz):
        return finish_quiz(quiz, score, len(quiz))
    
    q = quiz[index]
    remaining = int(end_time - time.time())
    timer_text = f"⏳ {remaining // 60}:{remaining % 60:02d}"
    
    return (
        f"**Question {index+1} of {len(quiz)}**\n\n{q['question']}",
        gr.update(choices=q["options"], value=None, visible=True),
        f"Score: {score}/{len(quiz)}",
        quiz,
        index,
        score,
        end_time,
        timer_text,
        gr.update(visible=True)
    )

def submit_answer(selected, quiz, index, score, end_time):
    if not selected:
        return show_question(quiz, index, score, end_time)
    
    if selected == quiz[index]["answer"]:
        score += 1
    
    return show_question(quiz, index + 1, score, end_time)

def finish_quiz(quiz, score, total):
    percentage = (score / total * 100) if total > 0 else 0
    emoji = "🎉" if percentage >= 80 else "👍" if percentage >= 60 else "📚"
    
    return (
        f"""## {emoji} Quiz Complete!

**Final Score: {score}/{total}** ({percentage:.1f}%)

{'Excellent work!' if percentage >= 80 else 'Good job!' if percentage >= 60 else 'Keep practicing!'}""",
        gr.update(choices=[], visible=False),
        "",
        None,
        0,
        score,
        None,
        "⏰ Time's up!" if total > 0 else "",
        gr.update(visible=False)
    )

def translate_to_urdu(text):
    if not text or not text.strip():
        return "Please enter some text to translate."
    
    if not groq_client:
        return "❌ Groq API not configured."
    
    try:
        chat_completion = groq_client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "You are a professional English to Urdu translator. Translate accurately to Urdu (اردو) using natural language. Respond ONLY with the translation."
                },
                {
                    "role": "user",
                    "content": f"Translate to Urdu:\n\n{text}"
                }
            ],
            model="llama-3.3-70b-versatile",
            temperature=0.3,
            max_completion_tokens=2048,
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        return f"Error: {str(e)}"

# ==================== CUSTOM CSS ====================
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=Noto+Nastaliq+Urdu&display=swap');

:root {
    --primary: #6366f1;
    --primary-dark: #4f46e5;
    --secondary: #8b5cf6;
    --success: #10b981;
    --warning: #f59e0b;
    --danger: #ef4444;
    --dark: #1f2937;
    --light: #f3f4f6;
    --gradient: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
}

body {
    font-family: 'Inter', sans-serif !important;
    background: #f8fafc !important;
}

.auth-container {
    max-width: 450px;
    margin: 3rem auto;
    padding: 2.5rem;
    background: white;
    border-radius: 24px;
    box-shadow: 0 25px 80px rgba(0,0,0,0.15);
    text-align: center;
}

.auth-logo {
    width: 80px;
    height: 80px;
    background: var(--gradient);
    border-radius: 20px;
    display: flex;
    align-items: center;
    justify-content: center;
    margin: 0 auto 1.5rem;
    font-size: 2.5rem;
    color: white;
}

.auth-title {
    font-size: 1.875rem;
    font-weight: 700;
    color: var(--dark);
    margin-bottom: 0.5rem;
}

.auth-subtitle {
    color: #6b7280;
    margin-bottom: 2rem;
    font-size: 1rem;
}

.auth-toggle {
    display: flex;
    background: #f3f4f6;
    border-radius: 12px;
    padding: 4px;
    margin-bottom: 1.5rem;
}

.auth-toggle-btn {
    flex: 1;
    padding: 0.75rem;
    border: none;
    background: transparent;
    border-radius: 8px;
    font-weight: 500;
    cursor: pointer;
    transition: all 0.3s;
}

.auth-toggle-btn.active {
    background: white;
    box-shadow: 0 2px 8px rgba(0,0,0,0.1);
    color: var(--primary);
}

.input-group {
    margin-bottom: 1rem;
    text-align: left;
}

.input-label {
    display: block;
    font-size: 0.875rem;
    font-weight: 500;
    color: #374151;
    margin-bottom: 0.5rem;
}

.input-field {
    width: 100%;
    padding: 0.875rem 1rem;
    border: 2px solid #e5e7eb;
    border-radius: 12px;
    font-size: 1rem;
    transition: all 0.3s;
    box-sizing: border-box;
}

.input-field:focus {
    outline: none;
    border-color: var(--primary);
    box-shadow: 0 0 0 4px rgba(99, 102, 241, 0.1);
}

.btn-primary {
    width: 100%;
    padding: 1rem;
    background: var(--gradient) !important;
    color: white !important;
    border: none !important;
    border-radius: 12px !important;
    font-size: 1rem;
    font-weight: 600 !important;
    cursor: pointer;
    transition: all 0.3s !important;
    margin-top: 0.5rem;
}

.btn-primary:hover {
    transform: translateY(-2px);
    box-shadow: 0 10px 25px rgba(102, 126, 234, 0.4) !important;
}

.btn-secondary {
    width: 100%;
    padding: 0.875rem;
    background: white !important;
    color: var(--dark) !important;
    border: 2px solid #e5e7eb !important;
    border-radius: 12px !important;
    font-weight: 500 !important;
    margin-top: 0.75rem;
}

.auth-message {
    padding: 1rem;
    border-radius: 12px;
    margin-top: 1rem;
    font-size: 0.875rem;
}

.auth-message.success {
    background: #d1fae5;
    color: #065f46;
    border: 1px solid #a7f3d0;
}

.auth-message.error {
    background: #fee2e2;
    color: #991b1b;
    border: 1px solid #fecaca;
}

.app-header {
    background: var(--gradient);
    color: white;
    padding: 1.5rem 2rem;
    border-radius: 0 0 30px 30px;
    margin: -20px -20px 2rem -20px;
    box-shadow: 0 10px 40px rgba(102, 126, 234, 0.3);
    position: relative;
}

.app-header h1 {
    font-size: 2rem;
    font-weight: 700;
    margin: 0;
    display: flex;
    align-items: center;
    gap: 0.75rem;
}

.user-widget {
    position: absolute;
    top: 1.5rem;
    right: 2rem;
    background: rgba(255,255,255,0.2);
    backdrop-filter: blur(10px);
    padding: 0.5rem 1rem;
    border-radius: 50px;
    display: flex;
    align-items: center;
    gap: 0.75rem;
    font-size: 0.875rem;
}

.status-bar {
    display: flex;
    gap: 0.75rem;
    margin-bottom: 1.5rem;
    flex-wrap: wrap;
}

.status-pill {
    padding: 0.5rem 1rem;
    border-radius: 50px;
    font-size: 0.8rem;
    font-weight: 500;
    display: flex;
    align-items: center;
    gap: 0.5rem;
}

.status-ok { background: #d1fae5; color: #065f46; }
.status-warn { background: #fef3c7; color: #92400e; }
.status-error { background: #fee2e2; color: #991b1b; }

.tool-card {
    background: white;
    border-radius: 20px;
    padding: 1.5rem;
    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05);
    border: 1px solid #f3f4f6;
    transition: all 0.3s;
}

.tool-card:hover {
    transform: translateY(-4px);
    box-shadow: 0 20px 40px rgba(0, 0, 0, 0.08);
}

.tool-icon {
    width: 56px;
    height: 56px;
    border-radius: 16px;
    display: flex;
    align-items: center;
    justify-content: center;
    font-size: 1.75rem;
    margin-bottom: 1rem;
    background: var(--gradient);
    color: white;
}

.quiz-card {
    background: var(--gradient);
    color: white;
    padding: 2rem;
    border-radius: 20px;
    text-align: center;
}

.quiz-timer {
    font-size: 2rem;
    font-weight: 700;
    font-family: monospace;
    background: rgba(255,255,255,0.2);
    padding: 1rem;
    border-radius: 16px;
    margin-bottom: 1.5rem;
}

.urdu-text { 
    font-family: 'Noto Nastaliq Urdu', serif !important; 
    font-size: 1.5em !important; 
    line-height: 2 !important; 
    direction: rtl !important; 
    text-align: right !important;
    background: #f8fafc;
    padding: 1.5rem;
    border-radius: 16px;
    border: 2px solid #e5e7eb;
}

.app-footer {
    text-align: center;
    padding: 2rem;
    color: #9ca3af;
    font-size: 0.875rem;
    margin-top: 3rem;
    border-top: 1px solid #f3f4f6;
}
"""

# ==================== AUTH HANDLERS ====================
def handle_login(username, password):
    session_token, user_data = auth_system.login(username, password)
    if session_token:
        return {
            session_token_state: gr.update(value=session_token),
            auth_screen: gr.update(visible=False),
            main_app: gr.update(visible=True),
            user_display: gr.update(value=f"👤 {user_data['name']}"),
            login_message: gr.update(visible=False),
            signup_message: gr.update(visible=False)
        }
    else:
        return {
            session_token_state: gr.update(value=""),
            auth_screen: gr.update(visible=True),
            main_app: gr.update(visible=False),
            user_display: gr.update(value=""),
            login_message: gr.update(value=f"❌ {user_data}", visible=True),
            signup_message: gr.update(visible=False)
        }

def handle_signup(username, password, name, email):
    success, message = auth_system.register(username, password, name, email)
    
    if success:
        # Auto-login after signup
        session_token, user_data = auth_system.login(username, password)
        if session_token:
            return {
                session_token_state: gr.update(value=session_token),
                auth_screen: gr.update(visible=False),
                main_app: gr.update(visible=True),
                user_display: gr.update(value=f"👤 {user_data['name']}"),
                signup_message: gr.update(value=f"✅ {message}", visible=True),
                login_message: gr.update(visible=False)
            }
    
    return {
        session_token_state: gr.update(value=""),
        auth_screen: gr.update(visible=True),
        main_app: gr.update(visible=False),
        user_display: gr.update(value=""),
        signup_message: gr.update(value=f"❌ {message}", visible=True),
        login_message: gr.update(visible=False)
    }

def handle_logout(session_token):
    auth_system.logout(session_token)
    return {
        session_token_state: gr.update(value=""),
        auth_screen: gr.update(visible=True),
        main_app: gr.update(visible=False),
        user_display: gr.update(value=""),
        login_username: gr.update(value=""),
        login_password: gr.update(value=""),
        signup_username: gr.update(value=""),
        signup_password: gr.update(value=""),
        signup_name: gr.update(value=""),
        signup_email: gr.update(value=""),
        login_message: gr.update(visible=False),
        signup_message: gr.update(visible=False)
    }

def toggle_auth_mode(mode):
    if mode == "login":
        return {
            login_form: gr.update(visible=True),
            signup_form: gr.update(visible=False),
            login_tab: gr.update(elem_classes="auth-toggle-btn active"),
            signup_tab: gr.update(elem_classes="auth-toggle-btn")
        }
    else:
        return {
            login_form: gr.update(visible=False),
            signup_form: gr.update(visible=True),
            login_tab: gr.update(elem_classes="auth-toggle-btn"),
            signup_tab: gr.update(elem_classes="auth-toggle-btn active")
        }

# ==================== UI BUILDER ====================
def build_auth_screen():
    with gr.Column(visible=True, elem_classes="auth-container") as auth_screen:
        gr.Markdown("""
        <div class="auth-logo">🎓</div>
        <h1 class="auth-title">Student AI Suite</h1>
        <p class="auth-subtitle">Your personal academic assistant</p>
        """)
        
        with gr.Row(elem_classes="auth-toggle"):
            login_tab = gr.Button("Sign In", elem_classes="auth-toggle-btn active")
            signup_tab = gr.Button("Create Account", elem_classes="auth-toggle-btn")
        
        with gr.Column(visible=True) as login_form:
            login_username = gr.Textbox(label="Username", placeholder="Enter your username", elem_classes="input-field")
            login_password = gr.Textbox(label="Password", type="password", placeholder="Enter your password", elem_classes="input-field")
            login_btn = gr.Button("Sign In", elem_classes="btn-primary")
            login_message = gr.Markdown(visible=False, elem_classes="auth-message")
        
        with gr.Column(visible=False) as signup_form:
            signup_name = gr.Textbox(label="Full Name", placeholder="Your name", elem_classes="input-field")
            signup_email = gr.Textbox(label="Email", placeholder="your@email.com", elem_classes="input-field")
            signup_username = gr.Textbox(label="Username", placeholder="Choose a username", elem_classes="input-field")
            signup_password = gr.Textbox(label="Password", type="password", placeholder="Min 6 characters", elem_classes="input-field")
            signup_btn = gr.Button("Create Account", elem_classes="btn-primary")
            signup_message = gr.Markdown(visible=False, elem_classes="auth-message")
        
        gr.Markdown("""
        <div style="margin-top: 2rem; padding-top: 1.5rem; border-top: 1px solid #f3f4f6; color: #9ca3af; font-size: 0.875rem;">
            🔒 Secure • Private • Free Forever
        </div>
        """)
    
    return (auth_screen, login_tab, signup_tab, login_form, signup_form,
            login_username, login_password, login_btn, login_message,
            signup_name, signup_email, signup_username, signup_password, signup_btn, signup_message)

def build_main_app():
    with gr.Column(visible=False) as main_app:
        with gr.Row(elem_classes="app-header"):
            gr.Markdown("""
            <h1>🎓 Student AI Suite</h1>
            <div style="opacity: 0.9; font-size: 1rem; margin-top: 0.25rem;">
                Essay • PDF • Quiz • Translate
            </div>
            """)
            
            with gr.Row(elem_classes="user-widget"):
                user_display = gr.Markdown()
                logout_btn = gr.Button("Logout", size="sm", variant="secondary")
        
        with gr.Row(elem_classes="status-bar"):
            gemini_status = "✅ Gemini" if gemini_client else "❌ Gemini"
            groq_status = "✅ Groq" if groq_client else "❌ Groq"
            gr.Markdown(f"""
            <span class="status-pill {'status-ok' if gemini_client else 'status-error'}">🤖 {gemini_status}</span>
            <span class="status-pill {'status-ok' if groq_client else 'status-error'}">🌐 {groq_status}</span>
            """)
        
        with gr.Tabs():
            with gr.TabItem("📝 Essay & PDF"):
                with gr.Tabs():
                    with gr.Tab("✍️ Essay Generator"):
                        with gr.Row():
                            with gr.Column():
                                essay_prompt = gr.Textbox(label="Essay Topic", placeholder="e.g., 'Impact of Artificial Intelligence on Modern Education'", lines=3)
                                with gr.Row():
                                    essay_type = gr.Dropdown(["Argumentative", "Expository", "Descriptive", "Persuasive"], value="Argumentative", label="Type")
                                    essay_tone = gr.Dropdown(["Academic", "Formal", "Neutral"], value="Academic", label="Tone")
                                essay_words = gr.Slider(200, 1000, 500, step=50, label="Word Count")
                                essay_btn = gr.Button("✨ Generate Essay", variant="primary")
                            
                            with gr.Column():
                                essay_output = gr.Markdown(elem_classes="output-box")
                    
                    with gr.Tab("📄 PDF Summarizer"):
                        with gr.Row():
                            with gr.Column():
                                pdf_file = gr.File(label="Upload PDF", file_types=[".pdf"], type="binary")
                                with gr.Row():
                                    pdf_max = gr.Slider(50, 500, 200, step=10, label="Max Words")
                                    pdf_min = gr.Slider(20, 200, 50, step=10, label="Min Words")
                                pdf_btn = gr.Button("📝 Summarize PDF", variant="primary")
                                
                                gr.Markdown("---")
                                pdf_text = gr.Textbox(label="Or paste text", lines=4)
                                text_btn = gr.Button("Summarize Text")
                            
                            with gr.Column():
                                pdf_output = gr.Textbox(label="Summary", lines=12, elem_classes="output-box")
            
            with gr.TabItem("🎯 Smart Quiz"):
                with gr.Row():
                    with gr.Column(scale=1):
                        quiz_text = gr.Textbox(label="Study Material", placeholder="Paste your notes, textbook content, or any study material here...", lines=8)
                        with gr.Row():
                            quiz_num = gr.Slider(1, 10, 5, step=1, label="Questions")
                            quiz_time = gr.Slider(1, 10, 2, step=1, label="Minutes")
                        quiz_start = gr.Button("🚀 Start Quiz", variant="primary")
                    
                    with gr.Column(scale=2):
                        with gr.Column(elem_classes="quiz-card"):
                            quiz_timer = gr.Markdown("⏳ 02:00", elem_classes="quiz-timer")
                            quiz_question = gr.Markdown("### Ready to test your knowledge?")
                            quiz_options = gr.Radio(choices=[], label="Select Answer", visible=False)
                            quiz_submit = gr.Button("Submit Answer", visible=False)
                            quiz_score = gr.Markdown()
                        
                        quiz_state = gr.State()
                        quiz_idx = gr.State(0)
                        quiz_scr = gr.State(0)
                        quiz_end = gr.State()
            
            with gr.TabItem("🌍 Urdu Translator"):
                with gr.Row():
                    with gr.Column():
                        trans_input = gr.Textbox(label="English Text", placeholder="Enter text to translate to Urdu...", lines=6)
                        trans_btn = gr.Button("🔄 Translate", variant="primary")
                        gr.Examples(examples=["Hello, how are you?", "Pakistan is beautiful", "I love learning"], inputs=trans_input)
                    
                    with gr.Column():
                        trans_output = gr.Textbox(label="اردو ترجمہ", lines=6, elem_classes="urdu-text", interactive=False)
        
        gr.Markdown("""
        <div class="app-footer">
            <p>Made with ❤️ for students worldwide</p>
        </div>
        """)
    
    return (main_app, user_display, logout_btn, essay_prompt, essay_type, essay_tone,
            essay_words, essay_btn, essay_output, pdf_file, pdf_max, pdf_min, pdf_btn,
            pdf_text, text_btn, pdf_output, quiz_text, quiz_num, quiz_time, quiz_start,
            quiz_timer, quiz_question, quiz_options, quiz_submit, quiz_score,
            quiz_state, quiz_idx, quiz_scr, quiz_end, trans_input, trans_btn, trans_output)

# ==================== MAIN ====================
with gr.Blocks(title="Student AI Suite", css=custom_css) as demo:
    session_token_state = gr.State("")
    
    (auth_screen, login_tab, signup_tab, login_form, signup_form,
     login_username, login_password, login_btn, login_message,
     signup_name, signup_email, signup_username, signup_password, signup_btn, signup_message) = build_auth_screen()
    
    (main_app, user_display, logout_btn, essay_prompt, essay_type, essay_tone,
     essay_words, essay_btn, essay_output, pdf_file, pdf_max, pdf_min, pdf_btn,
     pdf_text, text_btn, pdf_output, quiz_text, quiz_num, quiz_time, quiz_start,
     quiz_timer, quiz_question, quiz_options, quiz_submit, quiz_score,
     quiz_state, quiz_idx, quiz_scr, quiz_end, trans_input, trans_btn, trans_output) = build_main_app()
    
    # Event Handlers
    login_tab.click(lambda: toggle_auth_mode("login"), 
                   outputs=[login_form, signup_form, login_tab, signup_tab])
    signup_tab.click(lambda: toggle_auth_mode("signup"), 
                    outputs=[login_form, signup_form, login_tab, signup_tab])
    
    login_btn.click(handle_login,
                   inputs=[login_username, login_password],
                   outputs=[session_token_state, auth_screen, main_app, user_display, login_message, signup_message])
    
    signup_btn.click(handle_signup,
                    inputs=[signup_username, signup_password, signup_name, signup_email],
                    outputs=[session_token_state, auth_screen, main_app, user_display, signup_message, login_message])
    
    logout_btn.click(handle_logout,
                    inputs=[session_token_state],
                    outputs=[session_token_state, auth_screen, main_app, user_display, 
                            login_username, login_password, signup_username, signup_password,
                            signup_name, signup_email, login_message, signup_message])
    
    # Tools
    essay_btn.click(generate_essay, 
                   inputs=[essay_prompt, essay_type, essay_words, essay_tone], 
                   outputs=essay_output)
    
    pdf_btn.click(summarize_pdf, 
                 inputs=[pdf_file, pdf_max, pdf_min], 
                 outputs=pdf_output)
    
    text_btn.click(summarize_text, 
                  inputs=[pdf_text, pdf_max, pdf_min], 
                  outputs=pdf_output)
    
    quiz_start.click(start_quiz,
                    inputs=[quiz_text, quiz_num, quiz_time],
                    outputs=[quiz_question, quiz_options, quiz_score, quiz_state, 
                            quiz_idx, quiz_scr, quiz_end, quiz_timer, quiz_submit])
    
    quiz_submit.click(submit_answer,
                     inputs=[quiz_options, quiz_state, quiz_idx, quiz_scr, quiz_end],
                     outputs=[quiz_question, quiz_options, quiz_score, quiz_state, 
                             quiz_idx, quiz_scr, quiz_end, quiz_timer, quiz_submit])
    
    trans_btn.click(translate_to_urdu, 
                   inputs=trans_input, 
                   outputs=trans_output)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)