app.py

import gradio as gr
import pandas as pd
import json
import os
import re
from PyPDF2 import PdfReader
from collections import defaultdict
from typing import Dict, List, Optional, Tuple, Union
import html
from pathlib import Path
import fitz  # PyMuPDF for better PDF text extraction
import pytesseract
from PIL import Image
import io
import secrets
import string
from huggingface_hub import HfApi, HfFolder
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import time
import logging
import asyncio

# ========== CONFIGURATION ==========
PROFILES_DIR = "student_profiles"
ALLOWED_FILE_TYPES = [".pdf", ".png", ".jpg", ".jpeg"]
MAX_FILE_SIZE_MB = 5
MIN_AGE = 5
MAX_AGE = 120
SESSION_TOKEN_LENGTH = 32
HF_TOKEN = os.getenv("HF_TOKEN")

# Initialize logging
logging.basicConfig(filename='app.log', level=logging.INFO)

# Model configuration
MODEL_CHOICES = {
    "TinyLlama (Fastest)": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    "Phi-2 (Balanced)": "microsoft/phi-2",
    "DeepSeek-V3 (Most Powerful)": "deepseek-ai/DeepSeek-V3"
}
DEFAULT_MODEL = "TinyLlama (Fastest)"

# Initialize Hugging Face API
if HF_TOKEN:
    hf_api = HfApi(token=HF_TOKEN)
    HfFolder.save_token(HF_TOKEN)

# ========== OPTIMIZED MODEL LOADING ==========
class ModelLoader:
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.loaded = False
        self.loading = False
        self.error = None
        self.current_model = None
    
    def load_model(self, model_name, progress=gr.Progress()):
        """Lazy load the model with progress feedback"""
        if self.loaded and self.current_model == model_name:
            return self.model, self.tokenizer
            
        self.loading = True
        self.error = None
        try:
            progress(0.1, desc="Initializing...")
            
            # Clear previous model if any
            if self.model:
                del self.model
                del self.tokenizer
                torch.cuda.empty_cache()
                time.sleep(2)  # Allow CUDA cleanup
            
            # Load with optimized settings
            model_kwargs = {
                "trust_remote_code": True,
                "torch_dtype": torch.float16,
                "device_map": "auto",
                "low_cpu_mem_usage": True
            }
            
            if "TinyLlama" in model_name:
                model_kwargs["attn_implementation"] = "flash_attention_2"
            
            progress(0.3, desc="Loading tokenizer...")
            self.tokenizer = AutoTokenizer.from_pretrained(
                MODEL_CHOICES[model_name],
                trust_remote_code=True
            )
            
            progress(0.6, desc="Loading model...")
            self.model = AutoModelForCausalLM.from_pretrained(
                MODEL_CHOICES[model_name],
                **model_kwargs
            )
            
            # Verify model responsiveness
            progress(0.8, desc="Verifying model...")
            test_input = self.tokenizer("Test", return_tensors="pt").to(self.model.device)
            _ = self.model.generate(**test_input, max_new_tokens=1)
            
            self.model.eval()  # Disable dropout
            progress(0.9, desc="Finalizing...")
            self.loaded = True
            self.current_model = model_name
            return self.model, self.tokenizer
            
        except torch.cuda.OutOfMemoryError:
            self.error = "Out of GPU memory. Try a smaller model."
            logging.error(self.error)
            return None, None
        except Exception as e:
            self.error = str(e)
            logging.error(f"Model loading error: {self.error}")
            return None, None
        finally:
            self.loading = False

# Initialize model loader
model_loader = ModelLoader()

# ========== UTILITY FUNCTIONS ==========
def generate_session_token() -> str:
    """Generate a random session token for user identification."""
    alphabet = string.ascii_letters + string.digits
    return ''.join(secrets.choice(alphabet) for _ in range(SESSION_TOKEN_LENGTH))

def sanitize_input(text: str) -> str:
    """Sanitize user input to prevent XSS and injection attacks."""
    return html.escape(text.strip())

def validate_name(name: str) -> str:
    """Validate name input."""
    name = name.strip()
    if not name:
        raise gr.Error("Name cannot be empty")
    if len(name) > 100:
        raise gr.Error("Name is too long (max 100 characters)")
    if any(c.isdigit() for c in name):
        raise gr.Error("Name cannot contain numbers")
    return name

def validate_age(age: Union[int, float, str]) -> int:
    """Validate and convert age input."""
    try:
        age_int = int(age)
        if not MIN_AGE <= age_int <= MAX_AGE:
            raise gr.Error(f"Age must be between {MIN_AGE} and {MAX_AGE}")
        return age_int
    except (ValueError, TypeError):
        raise gr.Error("Please enter a valid age number")

def validate_file(file_obj) -> None:
    """Validate uploaded file."""
    if not file_obj:
        raise ValueError("No file uploaded")
    
    file_ext = os.path.splitext(file_obj.name)[1].lower()
    if file_ext not in ALLOWED_FILE_TYPES:
        raise gr.Error(f"Invalid file type. Allowed: {', '.join(ALLOWED_FILE_TYPES)}")
    
    file_size = os.path.getsize(file_obj.name) / (1024 * 1024)  # MB
    if file_size > MAX_FILE_SIZE_MB:
        raise gr.Error(f"File too large. Max size: {MAX_FILE_SIZE_MB}MB")

# ========== TEXT EXTRACTION FUNCTIONS ==========
def extract_text_from_file(file_path: str, file_ext: str) -> str:
    """Enhanced text extraction with better error handling and fallbacks."""
    text = ""
    
    try:
        if file_ext == '.pdf':
            # First try PyMuPDF for better text extraction
            try:
                doc = fitz.open(file_path)
                for page in doc:
                    text += page.get_text("text") + '\n'
                if not text.strip():
                    raise ValueError("PyMuPDF returned empty text")
            except Exception as e:
                logging.warning(f"PyMuPDF failed: {str(e)}. Trying OCR fallback...")
                text = extract_text_from_pdf_with_ocr(file_path)
        
        elif file_ext in ['.png', '.jpg', '.jpeg']:
            text = extract_text_with_ocr(file_path)
            
        # Clean up the extracted text
        text = clean_extracted_text(text)
        
        if not text.strip():
            raise ValueError("No text could be extracted from the file")
            
        return text
    
    except Exception as e:
        logging.error(f"Text extraction error: {str(e)}")
        raise gr.Error(f"Text extraction error: {str(e)}\nTips: Use high-quality images/PDFs with clear text.")

def extract_text_from_pdf_with_ocr(file_path: str) -> str:
    """Fallback PDF text extraction using OCR."""
    text = ""
    try:
        doc = fitz.open(file_path)
        for page in doc:
            pix = page.get_pixmap()
            img = Image.open(io.BytesIO(pix.tobytes()))
            # Preprocess image for better OCR
            img = img.convert('L')  # Grayscale
            img = img.point(lambda x: 0 if x < 128 else 255)  # Binarize
            text += pytesseract.image_to_string(img, config='--psm 6 --oem 3') + '\n'
    except Exception as e:
        raise ValueError(f"PDF OCR failed: {str(e)}")
    return text

def extract_text_with_ocr(file_path: str) -> str:
    """Extract text from image files using OCR with preprocessing."""
    try:
        image = Image.open(file_path)
        
        # Enhanced preprocessing
        image = image.convert('L')  # Convert to grayscale
        image = image.point(lambda x: 0 if x < 128 else 255, '1')  # Thresholding
        
        # Custom Tesseract configuration
        custom_config = r'--oem 3 --psm 6'
        text = pytesseract.image_to_string(image, config=custom_config)
        return text
    except Exception as e:
        raise ValueError(f"OCR processing failed: {str(e)}")

def clean_extracted_text(text: str) -> str:
    """Clean and normalize the extracted text."""
    # Remove multiple spaces and newlines
    text = re.sub(r'\s+', ' ', text).strip()
    
    # Fix common OCR errors
    replacements = {
        '|': 'I',
        '‘': "'",
        '’': "'",
        '“': '"',
        '”': '"',
        'fi': 'fi',
        'fl': 'fl'
    }
    
    for wrong, right in replacements.items():
        text = text.replace(wrong, right)
    
    return text

def remove_sensitive_info(text: str) -> str:
    """Remove potentially sensitive information from transcript text."""
    # Remove social security numbers
    text = re.sub(r'\b\d{3}-\d{2}-\d{4}\b', '[REDACTED]', text)
    # Remove student IDs (assuming 6-9 digit numbers)
    text = re.sub(r'\b\d{6,9}\b', '[ID]', text)
    # Remove email addresses
    text = re.sub(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', '[EMAIL]', text)
    return text

# ========== TRANSCRIPT PARSING ==========
class TranscriptParser:
    def __init__(self):
        self.student_data = {}
        self.requirements = {}
        self.current_courses = []
        self.course_history = []
        
    def parse_transcript(self, text: str) -> Dict:
        """Main method to parse transcript text"""
        self._extract_student_info(text)
        self._extract_requirements(text)
        self._extract_course_history(text)
        self._extract_current_courses(text)
        
        return {
            "student_info": self.student_data,
            "requirements": self.requirements,
            "current_courses": self.current_courses,
            "course_history": self.course_history,
            "completion_status": self._calculate_completion()
        }
    
    def _extract_student_info(self, text: str):
        """Extract student personal information"""
        header_match = re.search(
            r"(\d{7}) - ([\w\s,]+)\s*\|\s*Cohort \w+\s*\|\s*Un-weighted GPA ([\d.]+)\s*\|\s*Comm Serv Hours (\d+)",
            text
        )
        if header_match:
            self.student_data = {
                "id": header_match.group(1),
                "name": header_match.group(2).strip(),
                "unweighted_gpa": float(header_match.group(3)),
                "community_service_hours": int(header_match.group(4))
            }
        
        # Extract additional info
        grade_match = re.search(
            r"Current Grade: (\d+)\s*\|\s*YOG (\d{4})\s*\|\s*Weighted GPA ([\d.]+)\s*\|\s*Total Credits Earned ([\d.]+)",
            text
        )
        if grade_match:
            self.student_data.update({
                "current_grade": grade_match.group(1),
                "graduation_year": grade_match.group(2),
                "weighted_gpa": float(grade_match.group(3)),
                "total_credits": float(grade_match.group(4))
            })
    
    def _extract_requirements(self, text: str):
        """Parse the graduation requirements section"""
        req_table = re.findall(
            r"\|([A-Z]-[\w\s]+)\s*\|([^\|]+)\|([\d.]+)\s*\|([\d.]+)\s*\|([\d.]+)\s*\|([^\|]+)\|",
            text
        )
        
        for row in req_table:
            req_name = row[0].strip()
            self.requirements[req_name] = {
                "required": float(row[2]),
                "completed": float(row[4]),
                "status": f"{row[5].strip()}%"
            }
    
    def _extract_course_history(self, text: str):
        """Parse the detailed course history"""
        course_lines = re.findall(
            r"\|([A-Z]-[\w\s&\(\)]+)\s*\|(\d{4}-\d{4})\s*\|(\d{2})\s*\|([A-Z0-9]+)\s*\|([^\|]+)\|([^\|]+)\|([^\|]+)\|([A-Z])\s*\|([YRXW]?)\s*\|([^\|]+)\|",
            text
        )
        
        for course in course_lines:
            self.course_history.append({
                "requirement_category": course[0].strip(),
                "school_year": course[1],
                "grade_level": course[2],
                "course_code": course[3],
                "description": course[4].strip(),
                "term": course[5].strip(),
                "district_number": course[6].strip(),
                "grade": course[7],
                "inclusion_status": course[8],
                "credits": course[9].strip()
            })
    
    def _extract_current_courses(self, text: str):
        """Identify courses currently in progress"""
        in_progress = [c for c in self.course_history if "inProgress" in c["credits"]]
        self.current_courses = [
            {
                "course": c["description"],
                "category": c["requirement_category"],
                "term": c["term"],
                "credits": c["credits"]
            }
            for c in in_progress
        ]
    
    def _calculate_completion(self) -> Dict:
        """Calculate overall completion status"""
        total_required = sum(req["required"] for req in self.requirements.values())
        total_completed = sum(req["completed"] for req in self.requirements.values())
        
        return {
            "total_required": total_required,
            "total_completed": total_completed,
            "percent_complete": round((total_completed / total_required) * 100, 1),
            "remaining_credits": total_required - total_completed
        }
    
    def to_json(self) -> str:
        """Export parsed data as JSON"""
        return json.dumps({
            "student_info": self.student_data,
            "requirements": self.requirements,
            "current_courses": self.current_courses,
            "course_history": self.course_history,
            "completion_status": self._calculate_completion()
        }, indent=2)

async def parse_transcript_async(file_obj, progress=gr.Progress()):
    """Async wrapper for transcript parsing"""
    return await asyncio.to_thread(parse_transcript, file_obj, progress)

def parse_transcript_with_ai(text: str, progress=gr.Progress()) -> Dict:
    """Use AI model to parse transcript text with progress feedback"""
    model, tokenizer = model_loader.load_model(model_loader.current_model or DEFAULT_MODEL, progress)
    if model is None or tokenizer is None:
        raise gr.Error(f"Model failed to load. {model_loader.error or 'Please try loading a model first.'}")
    
    # First try the structured parser
    try:
        progress(0.1, desc="Parsing transcript structure...")
        parser = TranscriptParser()
        parsed_data = parser.parse_transcript(text)
        progress(0.9, desc="Formatting results...")
        
        # Convert to expected format
        formatted_data = {
            "grade_level": parsed_data["student_info"].get("current_grade", "Unknown"),
            "gpa": {
                "weighted": parsed_data["student_info"].get("weighted_gpa", "N/A"),
                "unweighted": parsed_data["student_info"].get("unweighted_gpa", "N/A")
            },
            "courses": []
        }
        
        # Add courses
        for course in parsed_data["course_history"]:
            formatted_data["courses"].append({
                "code": course["course_code"],
                "name": course["description"],
                "grade": course["grade"],
                "credits": course["credits"],
                "year": course["school_year"],
                "grade_level": course["grade_level"]
            })
        
        progress(1.0)
        return validate_parsed_data(formatted_data)
    
    except Exception as e:
        logging.warning(f"Structured parsing failed, falling back to AI: {str(e)}")
        # Fall back to AI parsing if structured parsing fails
        return parse_transcript_with_ai_fallback(text, progress)

def parse_transcript_with_ai_fallback(text: str, progress=gr.Progress()) -> Dict:
    """Fallback AI parsing method"""
    # Pre-process the text
    text = remove_sensitive_info(text[:15000])  # Limit input size
    
    prompt = f"""
    Analyze this academic transcript and extract structured information:
    - Current grade level
    - Weighted GPA (if available)
    - Unweighted GPA (if available)
    - List of all courses with:
      * Course code
      * Course name
      * Grade received
      * Credits earned
      * Year/semester taken
      * Grade level when taken
    Return the data in JSON format.
    
    Transcript Text:
    {text}
    """
    
    try:
        progress(0.1, desc="Processing transcript with AI...")
        
        # Tokenize and generate response
        inputs = model_loader.tokenizer(prompt, return_tensors="pt").to(model_loader.model.device)
        progress(0.4)
        
        outputs = model_loader.model.generate(
            **inputs,
            max_new_tokens=1500,
            temperature=0.1,
            do_sample=True
        )
        progress(0.8)
        
        # Decode the response
        response = model_loader.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract JSON from response
        try:
            json_str = response.split('```json')[1].split('```')[0].strip()
            parsed_data = json.loads(json_str)
        except (IndexError, json.JSONDecodeError):
            # Fallback: Extract JSON-like substring
            json_str = re.search(r'\{.*\}', response, re.DOTALL).group()
            parsed_data = json.loads(json_str)
        
        progress(1.0)
        return validate_parsed_data(parsed_data)
    
    except torch.cuda.OutOfMemoryError:
        raise gr.Error("The model ran out of memory. Try with a smaller transcript or use a smaller model.")
    except Exception as e:
        logging.error(f"AI parsing error: {str(e)}")
        raise gr.Error(f"Error processing transcript: {str(e)}")

def validate_parsed_data(data: Dict) -> Dict:
    """Validate and clean the parsed data structure."""
    if not isinstance(data, dict):
        raise ValueError("Invalid data format")
    
    # Set default structure if missing
    if 'grade_level' not in data:
        data['grade_level'] = 'Unknown'
    
    if 'gpa' not in data:
        data['gpa'] = {'weighted': 'N/A', 'unweighted': 'N/A'}
    
    if 'courses' not in data:
        data['courses'] = []
    
    # Clean course data
    for course in data['courses']:
        if 'grade' in course:
            course['grade'] = course['grade'].upper().strip()
        
        # Ensure numeric credits are strings
        if 'credits' in course and isinstance(course['credits'], (int, float)):
            course['credits'] = str(course['credits'])
    
    return data

def format_transcript_output(data: Dict) -> str:
    """Format the parsed data into human-readable text."""
    output = []
    output.append(f"Student Transcript Summary\n{'='*40}")
    output.append(f"Current Grade Level: {data.get('grade_level', 'Unknown')}")
    
    if 'gpa' in data:
        output.append(f"\nGPA:")
        output.append(f"- Weighted: {data['gpa'].get('weighted', 'N/A')}")
        output.append(f"- Unweighted: {data['gpa'].get('unweighted', 'N/A')}")
    
    if 'courses' in data:
        output.append("\nCourse History:\n" + '='*40)
        
        # Group courses by grade level
        courses_by_grade = defaultdict(list)
        for course in data['courses']:
            grade_level = course.get('grade_level', 'Unknown')
            courses_by_grade[grade_level].append(course)
        
        # Sort grades numerically
        for grade in sorted(courses_by_grade.keys(), key=lambda x: int(x) if x.isdigit() else x):
            output.append(f"\nGrade {grade}:\n{'-'*30}")
            for course in courses_by_grade[grade]:
                course_str = f"- {course.get('code', '')} {course.get('name', 'Unnamed course')}"
                if 'grade' in course:
                    course_str += f" (Grade: {course['grade']})"
                if 'credits' in course:
                    course_str += f" | Credits: {course['credits']}"
                if 'year' in course:
                    course_str += f" | Year: {course['year']}"
                output.append(course_str)
    
    return '\n'.join(output)

def parse_transcript(file_obj, progress=gr.Progress()) -> Tuple[str, Optional[Dict]]:
    """Main function to parse transcript files."""
    try:
        if not file_obj:
            raise ValueError("Please upload a file first")
            
        validate_file(file_obj)
        file_ext = os.path.splitext(file_obj.name)[1].lower()
        
        # Extract text from file
        text = extract_text_from_file(file_obj.name, file_ext)
        
        # Use AI for parsing
        parsed_data = parse_transcript_with_ai(text, progress)
        
        # Format output text
        output_text = format_transcript_output(parsed_data)
        
        # Prepare the data structure for saving
        transcript_data = {
            "grade_level": parsed_data.get('grade_level', 'Unknown'),
            "gpa": parsed_data.get('gpa', {}),
            "courses": defaultdict(list)
        }
        
        # Organize courses by grade level
        for course in parsed_data.get('courses', []):
            grade_level = course.get('grade_level', 'Unknown')
            transcript_data["courses"][grade_level].append(course)
        
        return output_text, transcript_data
    
    except Exception as e:
        logging.error(f"Transcript processing error: {str(e)}")
        return f"Error processing transcript: {str(e)}", None

# ========== LEARNING STYLE QUIZ ==========
class LearningStyleQuiz:
    def __init__(self):
        self.questions = [
            "When you study for a test, you prefer to:",
            "When you need directions to a new place, you prefer:",
            "When you learn a new skill, you prefer to:",
            "When you're trying to concentrate, you:",
            "When you meet new people, you remember them by:",
            "When you're assembling furniture or a gadget, you:",
            "When choosing a restaurant, you rely most on:",
            "When you're in a waiting room, you typically:",
            "When giving someone instructions, you tend to:",
            "When you're trying to recall information, you:",
            "When you're at a museum or exhibit, you:",
            "When you're learning a new language, you prefer:",
            "When you're taking notes in class, you:",
            "When you're explaining something complex, you:",
            "When you're at a party, you enjoy:",
            "When you're trying to remember a phone number, you:",
            "When you're relaxing, you prefer to:",
            "When you're learning to use new software, you:",
            "When you're giving a presentation, you rely on:",
            "When you're solving a difficult problem, you:"
        ]
        
        self.options = [
            ["Read the textbook (Reading/Writing)", "Listen to lectures (Auditory)", "Use diagrams/charts (Visual)", "Practice problems (Kinesthetic)"],
            ["Look at a map (Visual)", "Have someone tell you (Auditory)", "Write down directions (Reading/Writing)", "Try walking/driving there (Kinesthetic)"],
            ["Read instructions (Reading/Writing)", "Have someone show you (Visual)", "Listen to explanations (Auditory)", "Try it yourself (Kinesthetic)"],
            ["Need quiet (Reading/Writing)", "Need background noise (Auditory)", "Need to move around (Kinesthetic)", "Need visual stimulation (Visual)"],
            ["Their face (Visual)", "Their name (Auditory)", "What you talked about (Reading/Writing)", "What you did together (Kinesthetic)"],
            ["Read the instructions carefully (Reading/Writing)", "Look at the diagrams (Visual)", "Ask someone to explain (Auditory)", "Start putting pieces together (Kinesthetic)"],
            ["Online photos of the food (Visual)", "Recommendations from friends (Auditory)", "Reading the menu online (Reading/Writing)", "Remembering how it felt to eat there (Kinesthetic)"],
            ["Read magazines (Reading/Writing)", "Listen to music (Auditory)", "Watch TV (Visual)", "Fidget or move around (Kinesthetic)"],
            ["Write them down (Reading/Writing)", "Explain verbally (Auditory)", "Demonstrate (Visual)", "Guide them physically (Kinesthetic)"],
            ["See written words in your mind (Visual)", "Hear the information in your head (Auditory)", "Write it down to remember (Reading/Writing)", "Associate it with physical actions (Kinesthetic)"],
            ["Read all the descriptions (Reading/Writing)", "Listen to audio guides (Auditory)", "Look at the displays (Visual)", "Touch interactive exhibits (Kinesthetic)"],
            ["Study grammar rules (Reading/Writing)", "Listen to native speakers (Auditory)", "Use flashcards with images (Visual)", "Practice conversations (Kinesthetic)"],
            ["Write detailed paragraphs (Reading/Writing)", "Record the lecture (Auditory)", "Draw diagrams and charts (Visual)", "Doodle while listening (Kinesthetic)"],
            ["Write detailed steps (Reading/Writing)", "Explain verbally with examples (Auditory)", "Draw diagrams (Visual)", "Use physical objects to demonstrate (Kinesthetic)"],
            ["Conversations with people (Auditory)", "Watching others or the environment (Visual)", "Writing notes or texting (Reading/Writing)", "Dancing or physical activities (Kinesthetic)"],
            ["See the numbers in your head (Visual)", "Say them aloud (Auditory)", "Write them down (Reading/Writing)", "Dial them on a keypad (Kinesthetic)"],
            ["Read a book (Reading/Writing)", "Listen to music (Auditory)", "Watch TV/movies (Visual)", "Do something physical (Kinesthetic)"],
            ["Read the manual (Reading/Writing)", "Ask someone to show you (Visual)", "Call tech support (Auditory)", "Experiment with the software (Kinesthetic)"],
            ["Detailed notes (Reading/Writing)", "Verbal explanations (Auditory)", "Visual slides (Visual)", "Physical demonstrations (Kinesthetic)"],
            ["Write out possible solutions (Reading/Writing)", "Talk through it with someone (Auditory)", "Draw diagrams (Visual)", "Build a model or prototype (Kinesthetic)"]
        ]
        
        self.learning_styles = {
            "Visual": {
                "description": "Visual learners prefer using images, diagrams, and spatial understanding.",
                "tips": [
                    "Use color coding in your notes",
                    "Create mind maps and diagrams",
                    "Watch educational videos",
                    "Use flashcards with images",
                    "Highlight important information in different colors"
                ],
                "careers": [
                    "Graphic Designer", "Architect", "Photographer", 
                    "Engineer", "Surgeon", "Pilot"
                ]
            },
            "Auditory": {
                "description": "Auditory learners learn best through listening and speaking.",
                "tips": [
                    "Record lectures and listen to them",
                    "Participate in study groups",
                    "Explain concepts out loud to yourself",
                    "Use rhymes or songs to remember information",
                    "Listen to educational podcasts"
                ],
                "careers": [
                    "Musician", "Journalist", "Lawyer",
                    "Psychologist", "Teacher", "Customer Service"
                ]
            },
            "Reading/Writing": {
                "description": "These learners prefer information displayed as words.",
                "tips": [
                    "Write detailed notes",
                    "Create summaries in your own words",
                    "Read textbooks and articles",
                    "Make lists to organize information",
                    "Rewrite your notes to reinforce learning"
                ],
                "careers": [
                    "Writer", "Researcher", "Editor",
                    "Accountant", "Programmer", "Historian"
                ]
            },
            "Kinesthetic": {
                "description": "Kinesthetic learners learn through movement and hands-on activities.",
                "tips": [
                    "Use hands-on activities",
                    "Take frequent movement breaks",
                    "Create physical models",
                    "Associate information with physical actions",
                    "Study while walking or pacing"
                ],
                "careers": [
                    "Athlete", "Chef", "Mechanic",
                    "Dancer", "Physical Therapist", "Carpenter"
                ]
            }
        }
    
    def evaluate_quiz(self, *answers) -> str:
        """Evaluate quiz answers and generate enhanced results."""
        answers = list(answers)  # Convert tuple to list
        if len(answers) != len(self.questions):
            raise gr.Error("Not all questions were answered")
        
        scores = {style: 0 for style in self.learning_styles}
        
        for i, answer in enumerate(answers):
            if not answer:
                continue  # Skip unanswered questions
                
            for j, style in enumerate(self.learning_styles):
                if answer == self.options[i][j]:
                    scores[style] += 1
                    break
        
        total_answered = sum(1 for ans in answers if ans)
        if total_answered == 0:
            raise gr.Error("No answers provided")
        
        percentages = {style: (score/total_answered)*100 for style, score in scores.items()}
        sorted_styles = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        
        # Generate enhanced results report
        result = "## Your Learning Style Results\n\n"
        result += "### Scores:\n"
        for style, score in sorted_styles:
            result += f"- **{style}**: {score}/{total_answered} ({percentages[style]:.1f}%)\n"
        
        max_score = max(scores.values())
        primary_styles = [style for style, score in scores.items() if score == max_score]
        
        result += "\n### Analysis:\n"
        if len(primary_styles) == 1:
            primary_style = primary_styles[0]
            style_info = self.learning_styles[primary_style]
            
            result += f"Your primary learning style is **{primary_style}**\n\n"
            result += f"**{primary_style} Characteristics**:\n"
            result += f"{style_info['description']}\n\n"
            
            result += "**Recommended Study Strategies**:\n"
            for tip in style_info['tips']:
                result += f"- {tip}\n"
            
            result += "\n**Potential Career Paths**:\n"
            for career in style_info['careers'][:6]:
                result += f"- {career}\n"
            
            # Add complementary strategies
            complementary = [s for s in sorted_styles if s[0] != primary_style][0][0]
            result += f"\nYou might also benefit from some **{complementary}** strategies:\n"
            for tip in self.learning_styles[complementary]['tips'][:3]:
                result += f"- {tip}\n"
        else:
            result += "You have multiple strong learning styles:\n"
            for style in primary_styles:
                result += f"- **{style}**\n"
            
            result += "\n**Combined Learning Strategies**:\n"
            result += "You may benefit from combining different learning approaches:\n"
            for style in primary_styles:
                result += f"\n**{style}** techniques:\n"
                for tip in self.learning_styles[style]['tips'][:2]:
                    result += f"- {tip}\n"
                
                result += f"\n**{style}** career suggestions:\n"
                for career in self.learning_styles[style]['careers'][:3]:
                    result += f"- {career}\n"
        
        return result

# Initialize quiz instance
learning_style_quiz = LearningStyleQuiz()

# ========== PROFILE MANAGEMENT ==========
class ProfileManager:
    def __init__(self):
        self.profiles_dir = Path(PROFILES_DIR)
        self.profiles_dir.mkdir(exist_ok=True, parents=True)
        self.current_session = None
    
    def set_session(self, session_token: str) -> None:
        """Set the current session token."""
        self.current_session = session_token
    
    def get_profile_path(self, name: str) -> Path:
        """Get profile path with session token if available."""
        if self.current_session:
            return self.profiles_dir / f"{name.replace(' ', '_')}_{self.current_session}_profile.json"
        return self.profiles_dir / f"{name.replace(' ', '_')}_profile.json"
    
    def save_profile(self, name: str, age: Union[int, str], interests: str, 
                    transcript: Dict, learning_style: str, 
                    movie: str, movie_reason: str, show: str, show_reason: str, 
                    book: str, book_reason: str, character: str, character_reason: str, 
                    blog: str) -> str:
        """Save student profile with validation."""
        try:
            # Validate required fields
            name = validate_name(name)
            age = validate_age(age)
            interests = sanitize_input(interests)
            
            # Prepare favorites data
            favorites = {
                "movie": sanitize_input(movie),
                "movie_reason": sanitize_input(movie_reason),
                "show": sanitize_input(show),
                "show_reason": sanitize_input(show_reason),
                "book": sanitize_input(book),
                "book_reason": sanitize_input(book_reason),
                "character": sanitize_input(character),
                "character_reason": sanitize_input(character_reason)
            }
            
            # Prepare full profile data
            data = {
                "name": name,
                "age": age,
                "interests": interests,
                "transcript": transcript if transcript else {},
                "learning_style": learning_style if learning_style else "Not assessed",
                "favorites": favorites,
                "blog": sanitize_input(blog) if blog else "",
                "session_token": self.current_session
            }
            
            # Save to JSON file
            filepath = self.get_profile_path(name)
            
            with open(filepath, "w", encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            
            # Upload to HF Hub if token is available
            if HF_TOKEN:
                try:
                    hf_api.upload_file(
                        path_or_fileobj=filepath,
                        path_in_repo=f"profiles/{filepath.name}",
                        repo_id="your-username/student-learning-assistant",
                        repo_type="dataset"
                    )
                except Exception as e:
                    logging.error(f"Failed to upload to HF Hub: {str(e)}")
            
            return self._generate_profile_summary(data)
        
        except Exception as e:
            logging.error(f"Error saving profile: {str(e)}")
            raise gr.Error(f"Error saving profile: {str(e)}")
    
    def load_profile(self, name: str = None, session_token: str = None) -> Dict:
        """Load profile by name or return the first one found."""
        try:
            if session_token:
                profile_pattern = f"*{session_token}_profile.json"
            else:
                profile_pattern = "*.json"
            
            profiles = list(self.profiles_dir.glob(profile_pattern))
            if not profiles:
                return {}
            
            if name:
                # Find profile by name
                name = name.replace(" ", "_")
                if session_token:
                    profile_file = self.profiles_dir / f"{name}_{session_token}_profile.json"
                else:
                    profile_file = self.profiles_dir / f"{name}_profile.json"
                
                if not profile_file.exists():
                    # Try loading from HF Hub
                    if HF_TOKEN:
                        try:
                            hf_api.download_file(
                                path_in_repo=f"profiles/{profile_file.name}",
                                repo_id="your-username/student-learning-assistant",
                                repo_type="dataset",
                                local_dir=self.profiles_dir
                            )
                        except:
                            raise gr.Error(f"No profile found for {name}")
                    else:
                        raise gr.Error(f"No profile found for {name}")
            else:
                # Load the first profile found
                profile_file = profiles[0]
            
            with open(profile_file, "r", encoding='utf-8') as f:
                return json.load(f)
        
        except Exception as e:
            logging.error(f"Error loading profile: {str(e)}")
            return {}
    
    def list_profiles(self, session_token: str = None) -> List[str]:
        """List all available profile names for the current session."""
        if session_token:
            profiles = list(self.profiles_dir.glob(f"*{session_token}_profile.json"))
        else:
            profiles = list(self.profiles_dir.glob("*.json"))
        
        # Extract just the name part (without session token)
        profile_names = []
        for p in profiles:
            name_part = p.stem.replace("_profile", "")
            if session_token:
                name_part = name_part.replace(f"_{session_token}", "")
            profile_names.append(name_part.replace("_", " "))
        
        return profile_names
    
    def _generate_profile_summary(self, data: Dict) -> str:
        """Generate markdown summary of the profile."""
        transcript = data.get("transcript", {})
        favorites = data.get("favorites", {})
        learning_style = data.get("learning_style", "Not assessed")
        
        markdown = f"""## Student Profile: {data['name']}
### Basic Information
- **Age:** {data['age']}
- **Interests:** {data.get('interests', 'Not specified')}
- **Learning Style:** {learning_style.split('##')[0].strip()}
### Academic Information
{self._format_transcript(transcript)}
### Favorites
- **Movie:** {favorites.get('movie', 'Not specified')}  
  *Reason:* {favorites.get('movie_reason', 'Not specified')}
- **TV Show:** {favorites.get('show', 'Not specified')}  
  *Reason:* {favorites.get('show_reason', 'Not specified')}
- **Book:** {favorites.get('book', 'Not specified')}  
  *Reason:* {favorites.get('book_reason', 'Not specified')}
- **Character:** {favorites.get('character', 'Not specified')}  
  *Reason:* {favorites.get('character_reason', 'Not specified')}
### Personal Blog
{data.get('blog', '_No blog provided_')}
"""
        return markdown
    
    def _format_transcript(self, transcript: Dict) -> str:
        """Format transcript data for display."""
        if not transcript or "courses" not in transcript:
            return "_No transcript information available_"
        
        display = "#### Course History\n"
        courses_by_grade = transcript["courses"]
        
        if isinstance(courses_by_grade, dict):
            for grade in sorted(courses_by_grade.keys(), key=lambda x: int(x) if x.isdigit() else x):
                display += f"\n**Grade {grade}**\n"
                for course in courses_by_grade[grade]:
                    display += f"- {course.get('code', '')} {course.get('name', 'Unnamed course')}"
                    if 'grade' in course and course['grade']:
                        display += f" (Grade: {course['grade']})"
                    if 'credits' in course:
                        display += f" | Credits: {course['credits']}"
                    display += f" | Year: {course.get('year', 'N/A')}\n"
        
        if 'gpa' in transcript:
            gpa = transcript['gpa']
            display += "\n**GPA**\n"
            display += f"- Unweighted: {gpa.get('unweighted', 'N/A')}\n"
            display += f"- Weighted: {gpa.get('weighted', 'N/A')}\n"
        
        return display

# Initialize profile manager
profile_manager = ProfileManager()

# ========== AI TEACHING ASSISTANT ==========
class TeachingAssistant:
    def __init__(self):
        self.context_history = []
        self.max_context_length = 5  # Keep last 5 exchanges for context
    
    def generate_response(self, message: str, history: List[List[Union[str, None]]], session_token: str) -> str:
        """Generate personalized response based on student profile and context."""
        try:
            # Load profile with session token
            profile = profile_manager.load_profile(session_token=session_token)
            if not profile:
                return "Please complete and save your profile first using the previous tabs."
            
            # Update context history
            self._update_context(message, history)
            
            # Extract profile information
            name = profile.get("name", "there")
            learning_style = profile.get("learning_style", "")
            grade_level = profile.get("transcript", {}).get("grade_level", "unknown")
            gpa = profile.get("transcript", {}).get("gpa", {})
            interests = profile.get("interests", "")
            courses = profile.get("transcript", {}).get("courses", {})
            favorites = profile.get("favorites", {})
            
            # Process message with context
            response = self._process_message(message, profile)
            
            # Add follow-up suggestions
            if "study" in message.lower() or "learn" in message.lower():
                response += "\n\nWould you like me to suggest a study schedule based on your courses?"
            elif "course" in message.lower() or "class" in message.lower():
                response += "\n\nWould you like help finding resources for any of these courses?"
            
            return response
        
        except Exception as e:
            logging.error(f"Error generating response: {str(e)}")
            return "I encountered an error processing your request. Please try again."
    
    def _update_context(self, message: str, history: List[List[Union[str, None]]]) -> None:
        """Maintain conversation context."""
        self.context_history.append({"role": "user", "content": message})
        if history:
            for h in history[-self.max_context_length:]:
                if h[0]:  # User message
                    self.context_history.append({"role": "user", "content": h[0]})
                if h[1]:  # Assistant message
                    self.context_history.append({"role": "assistant", "content": h[1]})
        
        # Trim to maintain max context length
        self.context_history = self.context_history[-(self.max_context_length*2):]
    
    def _process_message(self, message: str, profile: Dict) -> str:
        """Process user message with profile context."""
        message_lower = message.lower()
        
        # Greetings
        if any(greet in message_lower for greet in ["hi", "hello", "hey", "greetings"]):
            return f"Hello {profile.get('name', 'there')}! How can I help you with your learning today?"
        
        # Study help
        study_words = ["study", "learn", "prepare", "exam", "test", "homework"]
        if any(word in message_lower for word in study_words):
            return self._generate_study_advice(profile)
        
        # Grade help
        grade_words = ["grade", "gpa", "score", "marks", "results"]
        if any(word in message_lower for word in grade_words):
            return self._generate_grade_advice(profile)
        
        # Interest help
        interest_words = ["interest", "hobby", "passion", "extracurricular"]
        if any(word in message_lower for word in interest_words):
            return self._generate_interest_advice(profile)
        
        # Course help
        course_words = ["courses", "classes", "transcript", "schedule", "subject"]
        if any(word in message_lower for word in course_words):
            return self._generate_course_advice(profile)
        
        # Favorites
        favorite_words = ["movie", "show", "book", "character", "favorite"]
        if any(word in message_lower for word in favorite_words):
            return self._generate_favorites_response(profile)
        
        # General help
        if "help" in message_lower:
            return self._generate_help_response()
        
        # Default response
        return ("I'm your personalized teaching assistant. I can help with study tips, "
                "grade information, course advice, and more. Try asking about how to "
                "study effectively or about your course history.")
    
    def _generate_study_advice(self, profile: Dict) -> str:
        """Generate study advice based on learning style."""
        learning_style = profile.get("learning_style", "")
        response = ""
        
        if "Visual" in learning_style:
            response = ("Based on your visual learning style, I recommend:\n"
                       "- Creating colorful mind maps or diagrams\n"
                       "- Using highlighters to color-code your notes\n"
                       "- Watching educational videos on the topics\n"
                       "- Creating flashcards with images\n\n")
        elif "Auditory" in learning_style:
            response = ("Based on your auditory learning style, I recommend:\n"
                       "- Recording your notes and listening to them\n"
                       "- Participating in study groups to discuss concepts\n"
                       "- Explaining the material out loud to yourself\n"
                       "- Finding podcasts or audio lectures on the topics\n\n")
        elif "Reading/Writing" in learning_style:
            response = ("Based on your reading/writing learning style, I recommend:\n"
                       "- Writing detailed summaries in your own words\n"
                       "- Creating organized outlines of the material\n"
                       "- Reading additional textbooks or articles\n"
                       "- Rewriting your notes to reinforce learning\n\n")
        elif "Kinesthetic" in learning_style:
            response = ("Based on your kinesthetic learning style, I recommend:\n"
                       "- Creating physical models or demonstrations\n"
                       "- Using hands-on activities to learn concepts\n"
                       "- Taking frequent movement breaks while studying\n"
                       "- Associating information with physical actions\n\n")
        else:
            response = ("Here are some general study tips:\n"
                       "- Use the Pomodoro technique (25 min study, 5 min break)\n"
                       "- Space out your study sessions over time\n"
                       "- Test yourself with practice questions\n"
                       "- Teach the material to someone else\n\n")
        
        # Add time management advice
        response += ("**Time Management Tips**:\n"
                    "- Create a study schedule and stick to it\n"
                    "- Prioritize difficult subjects when you're most alert\n"
                    "- Break large tasks into smaller, manageable chunks\n"
                    "- Set specific goals for each study session")
        
        return response
    
    def _generate_grade_advice(self, profile: Dict) -> str:
        """Generate response about grades and GPA."""
        gpa = profile.get("transcript", {}).get("gpa", {})
        courses = profile.get("transcript", {}).get("courses", {})
        
        response = (f"Your GPA information:\n"
                   f"- Unweighted: {gpa.get('unweighted', 'N/A')}\n"
                   f"- Weighted: {gpa.get('weighted', 'N/A')}\n\n")
        
        # Identify any failing grades
        weak_subjects = []
        for grade_level, course_list in courses.items():
            for course in course_list:
                if course.get('grade', '').upper() in ['D', 'F']:
                    weak_subjects.append(f"{course.get('code', '')} {course.get('name', 'Unknown course')}")
        
        if weak_subjects:
            response += ("**Areas for Improvement**:\n"
                        f"You might want to focus on these subjects: {', '.join(weak_subjects)}\n\n")
        
        response += ("**Grade Improvement Strategies**:\n"
                    "- Meet with your teachers to discuss your performance\n"
                    "- Identify specific areas where you lost points\n"
                    "- Create a targeted study plan for weak areas\n"
                    "- Practice with past exams or sample questions")
        
        return response
    
    def _generate_interest_advice(self, profile: Dict) -> str:
        """Generate response based on student interests."""
        interests = profile.get("interests", "")
        response = f"I see you're interested in: {interests}\n\n"
        
        response += ("**Suggestions**:\n"
                    "- Look for clubs or extracurricular activities related to these interests\n"
                    "- Explore career paths that align with these interests\n"
                    "- Find online communities or forums about these topics\n"
                    "- Consider projects or independent study in these areas")
        
        return response
    
    def _generate_course_advice(self, profile: Dict) -> str:
        """Generate response about courses."""
        courses = profile.get("transcript", {}).get("courses", {})
        grade_level = profile.get("transcript", {}).get("grade_level", "unknown")
        
        response = "Here's a summary of your courses:\n"
        for grade in sorted(courses.keys(), key=lambda x: int(x) if x.isdigit() else x):
            response += f"\n**Grade {grade}**:\n"
            for course in courses[grade]:
                response += f"- {course.get('code', '')} {course.get('name', 'Unnamed course')}"
                if 'grade' in course:
                    response += f" (Grade: {course['grade']})"
                response += "\n"
        
        response += f"\nAs a grade {grade_level} student, you might want to:\n"
        if grade_level in ["9", "10"]:
            response += ("- Focus on building strong foundational skills\n"
                        "- Explore different subjects to find your interests\n"
                        "- Start thinking about college/career requirements")
        elif grade_level in ["11", "12"]:
            response += ("- Focus on courses relevant to your college/career goals\n"
                        "- Consider taking AP or advanced courses if available\n"
                        "- Ensure you're meeting graduation requirements")
        
        return response
    
    def _generate_favorites_response(self, profile: Dict) -> str:
        """Generate response about favorite items."""
        favorites = profile.get("favorites", {})
        response = "I see you enjoy:\n"
        
        if favorites.get('movie'):
            response += f"- Movie: {favorites['movie']} ({favorites.get('movie_reason', 'no reason provided')})\n"
        if favorites.get('show'):
            response += f"- TV Show: {favorites['show']} ({favorites.get('show_reason', 'no reason provided')})\n"
        if favorites.get('book'):
            response += f"- Book: {favorites['book']} ({favorites.get('book_reason', 'no reason provided')})\n"
        if favorites.get('character'):
            response += f"- Character: {favorites['character']} ({favorites.get('character_reason', 'no reason provided')})\n"
        
        response += "\nThese preferences suggest you might enjoy:\n"
        response += "- Similar books/movies in the same genre\n"
        response += "- Creative projects related to these stories\n"
        response += "- Analyzing themes or characters in your schoolwork"
        
        return response
    
    def _generate_help_response(self) -> str:
        """Generate help response with available commands."""
        return ("""I can help with:
- **Study tips**: "How should I study for math?"
- **Grade information**: "What's my GPA?"
- **Course advice**: "Show me my course history"
- **Interest suggestions**: "What clubs match my interests?"
- **General advice**: "How can I improve my grades?"
Try asking about any of these topics!""")

# Initialize teaching assistant
teaching_assistant = TeachingAssistant()

# ========== GRADIO INTERFACE ==========
def create_interface():
    with gr.Blocks(theme=gr.themes.Soft(), title="Student Learning Assistant") as app:
        # Session state
        session_token = gr.State(value=generate_session_token())
        profile_manager.set_session(session_token.value)
        
        # Track completion status for each tab
        tab_completed = gr.State({
            0: False,  # Transcript Upload
            1: False,  # Learning Style Quiz
            2: False,  # Personal Questions
            3: False,  # Save & Review
            4: False   # AI Assistant
        })
        
        # Custom CSS for better styling
        app.css = """
        .gradio-container {
            max-width: 1200px !important;
            margin: 0 auto;
        }
        .tab {
            padding: 20px;
            border-radius: 8px;
            background: white;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
        .progress-bar {
            height: 5px;
            background: linear-gradient(to right, #4CAF50, #8BC34A);
            margin-bottom: 15px;
            border-radius: 3px;
        }
        .quiz-question {
            margin-bottom: 15px;
            padding: 15px;
            background: #f5f5f5;
            border-radius: 5px;
        }
        .profile-card {
            border: 1px solid #e0e0e0;
            border-radius: 8px;
            padding: 15px;
            margin-bottom: 15px;
            background: white;
        }
        .chatbot {
            min-height: 500px;
        }
        .completed-tab {
            background: #2196F3 !important;
            color: white !important;
        }
        .incomplete-tab {
            background: #E0E0E0 !important;
        }
        .alert-box {
            padding: 15px;
            margin-bottom: 20px;
            border: 1px solid transparent;
            border-radius: 4px;
            color: #31708f;
            background-color: #d9edf7;
            border-color: #bce8f1;
        }
        .nav-message {
            padding: 10px;
            margin: 10px 0;
            border-radius: 4px;
            background-color: #ffebee;
            color: #c62828;
        }
        .model-loading {
            padding: 15px;
            margin: 15px 0;
            border-radius: 4px;
            background-color: #fff3e0;
            color: #e65100;
        }
        .model-selection {
            margin-bottom: 20px;
            padding: 15px;
            background: #f8f9fa;
            border-radius: 8px;
        }
        """
        
        gr.Markdown("""
        # Student Learning Assistant
        **Your personalized education companion**  
        Complete each step to get customized learning recommendations.
        """)
        
        # Model selection section
        with gr.Group(elem_classes="model-selection"):
            model_selector = gr.Dropdown(
                choices=list(MODEL_CHOICES.keys()),
                value=DEFAULT_MODEL,
                label="Select AI Model",
                interactive=True
            )
            load_model_btn = gr.Button("Load Selected Model", variant="secondary")
            model_status = gr.HTML(
                value="<div class='model-loading'>Model not loaded yet. Please select and load a model.</div>",
                visible=True
            )
        
        # Progress tracker
        with gr.Row():
            with gr.Column(scale=1):
                step1 = gr.Button("1. Upload Transcript", elem_classes="incomplete-tab")
            with gr.Column(scale=1):
                step2 = gr.Button("2. Learning Style Quiz", elem_classes="incomplete-tab", interactive=False)
            with gr.Column(scale=1):
                step3 = gr.Button("3. Personal Questions", elem_classes="incomplete-tab", interactive=False)
            with gr.Column(scale=1):
                step4 = gr.Button("4. Save & Review", elem_classes="incomplete-tab", interactive=False)
            with gr.Column(scale=1):
                step5 = gr.Button("5. AI Assistant", elem_classes="incomplete-tab", interactive=False)
        
        # Alert box for quiz submission
        quiz_alert = gr.HTML(visible=False)
        
        # Navigation message
        nav_message = gr.HTML(elem_classes="nav-message", visible=False)
        
        # Main tabs
        with gr.Tabs() as tabs:
            # ===== TAB 1: Transcript Upload =====
            with gr.Tab("Transcript Upload", id=0) as tab1:
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Step 1: Upload Your Transcript")
                        gr.Markdown("Upload a PDF or image of your academic transcript to analyze your courses and GPA.")
                        
                        with gr.Group():
                            transcript_file = gr.File(
                                label="Transcript (PDF or Image)",
                                file_types=ALLOWED_FILE_TYPES,
                                type="filepath"
                            )
                            upload_btn = gr.Button("Upload & Analyze", variant="primary")
                        
                        gr.Markdown("""
                        **Supported Formats**: PDF, PNG, JPG  
                        **Note**: Your file is processed locally and not stored permanently.
                        """)
                    
                    with gr.Column(scale=2):
                        transcript_output = gr.Textbox(
                            label="Transcript Analysis",
                            lines=20,
                            interactive=False
                        )
                        transcript_data = gr.State()
                
                def process_transcript_and_update(file_obj, current_tab_status, progress=gr.Progress()):
                    try:
                        output_text, data = parse_transcript(file_obj, progress)
                        if "Error" not in output_text:
                            new_status = current_tab_status.copy()
                            new_status[0] = True
                            return (
                                output_text, 
                                data, 
                                new_status,
                                gr.update(elem_classes="completed-tab"),
                                gr.update(interactive=True),
                                gr.update(visible=False)
                    except Exception as e:
                        logging.error(f"Upload error: {str(e)}")
                    return (
                        "Error processing transcript. Please try again.", 
                        None, 
                        current_tab_status,
                        gr.update(),
                        gr.update(),
                        gr.update(visible=True, value=f"<div class='nav-message'>Error: {str(e)}</div>"))
                
                upload_btn.click(
                    fn=process_transcript_and_update,
                    inputs=[transcript_file, tab_completed],
                    outputs=[transcript_output, transcript_data, tab_completed, step1, step2, nav_message],
                    concurrency_limit=1
                )
            
            # ===== TAB 2: Learning Style Quiz =====
            with gr.Tab("Learning Style Quiz", id=1) as tab2:
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Step 2: Discover Your Learning Style")
                        gr.Markdown("Complete this 20-question quiz to identify whether you're a visual, auditory, reading/writing, or kinesthetic learner.")
                        
                        progress = gr.HTML("<div class='progress-bar' style='width: 0%'></div>")
                        quiz_submit = gr.Button("Submit Quiz", variant="primary")
                    
                    with gr.Column(scale=2):
                        quiz_components = []
                        with gr.Accordion("Quiz Questions", open=True):
                            for i, (question, options) in enumerate(zip(learning_style_quiz.questions, learning_style_quiz.options)):
                                with gr.Group(elem_classes="quiz-question"):
                                    q = gr.Radio(
                                        options,
                                        label=f"{i+1}. {question}",
                                        show_label=True
                                    )
                                    quiz_components.append(q)
                        
                        learning_output = gr.Markdown(
                            label="Your Learning Style Results",
                            visible=False
                        )
                
                # Update progress bar as questions are answered
                for component in quiz_components:
                    component.change(
                        fn=lambda *answers: {
                            progress: gr.HTML(
                                f"<div class='progress-bar' style='width: {sum(1 for a in answers if a)/len(answers)*100}%'></div>"
                            )
                        },
                        inputs=quiz_components,
                        outputs=progress
                    )
                
                def submit_quiz_and_update(*args):
                    # The first argument is the tab_completed state, followed by answers
                    current_tab_status = args[0]
                    answers = args[1:]
                    
                    try:
                        result = learning_style_quiz.evaluate_quiz(*answers)
                        new_status = current_tab_status.copy()
                        new_status[1] = True
                        return (
                            result,
                            gr.update(visible=True),
                            new_status,
                            gr.update(elem_classes="completed-tab"),
                            gr.update(interactive=True),
                            gr.update(value="<div class='alert-box'>Quiz submitted successfully! Scroll down to view your results.</div>", visible=True),
                            gr.update(visible=False))
                    except Exception as e:
                        logging.error(f"Quiz error: {str(e)}")
                        return (
                            f"Error evaluating quiz: {str(e)}",
                            gr.update(visible=True),
                            current_tab_status,
                            gr.update(),
                            gr.update(),
                            gr.update(value=f"<div class='nav-message'>Error: {str(e)}</div>", visible=True),
                            gr.update(visible=False))
                
                quiz_submit.click(
                    fn=submit_quiz_and_update,
                    inputs=[tab_completed] + quiz_components,
                    outputs=[learning_output, learning_output, tab_completed, step2, step3, quiz_alert, nav_message]
                )
            
            # ===== TAB 3: Personal Questions =====
            with gr.Tab("Personal Profile", id=2) as tab3:
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Step 3: Tell Us About Yourself")
                        gr.Markdown("This information helps us provide personalized recommendations.")
                        
                        with gr.Group():
                            name = gr.Textbox(label="Full Name", placeholder="Your name")
                            age = gr.Number(label="Age", minimum=MIN_AGE, maximum=MAX_AGE, precision=0)
                            interests = gr.Textbox(
                                label="Your Interests/Hobbies",
                                placeholder="e.g., Science, Music, Sports, Art..."
                            )
                        
                        save_personal_btn = gr.Button("Save Information", variant="primary")
                        save_confirmation = gr.HTML(visible=False)
                        
                        gr.Markdown("### Favorites")
                        with gr.Group():
                            movie = gr.Textbox(label="Favorite Movie")
                            movie_reason = gr.Textbox(label="Why do you like it?", lines=2)
                            show = gr.Textbox(label="Favorite TV Show")
                            show_reason = gr.Textbox(label="Why do you like it?", lines=2)
                            book = gr.Textbox(label="Favorite Book")
                            book_reason = gr.Textbox(label="Why do you like it?", lines=2)
                            character = gr.Textbox(label="Favorite Character (from any story)")
                            character_reason = gr.Textbox(label="Why do you like them?", lines=2)
                    
                    with gr.Column(scale=1):
                        gr.Markdown("### Additional Information")
                        
                        blog_checkbox = gr.Checkbox(
                            label="Would you like to write a short blog about your learning experiences?",
                            value=False
                        )
                        blog_text = gr.Textbox(
                            label="Your Learning Blog",
                            placeholder="Write about your learning journey, challenges, goals...",
                            lines=8,
                            visible=False
                        )
                        blog_checkbox.change(
                            lambda x: gr.update(visible=x),
                            inputs=blog_checkbox,
                            outputs=blog_text
                        )
                
                def save_personal_info(name, age, interests, current_tab_status):
                    try:
                        name = validate_name(name)
                        age = validate_age(age)
                        interests = sanitize_input(interests)
                        
                        new_status = current_tab_status.copy()
                        new_status[2] = True
                        return (
                            new_status,
                            gr.update(elem_classes="completed-tab"),
                            gr.update(interactive=True),
                            gr.update(value="<div class='alert-box'>Information saved!</div>", visible=True),
                            gr.update(visible=False))
                    except Exception as e:
                        return (
                            current_tab_status,
                            gr.update(),
                            gr.update(),
                            gr.update(visible=False),
                            gr.update(visible=True, value=f"<div class='nav-message'>Error: {str(e)}</div>"))
                
                save_personal_btn.click(
                    fn=save_personal_info,
                    inputs=[name, age, interests, tab_completed],
                    outputs=[tab_completed, step3, step4, save_confirmation, nav_message]
                )
            
            # ===== TAB 4: Save & Review =====
            with gr.Tab("Save Profile", id=3) as tab4:
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Step 4: Review & Save Your Profile")
                        gr.Markdown("Verify your information before saving. You can return to previous steps to make changes.")
                        
                        save_btn = gr.Button("Save Profile", variant="primary")
                        
                        # Profile management section
                        with gr.Group():
                            load_profile_dropdown = gr.Dropdown(
                                label="Load Existing Profile",
                                choices=profile_manager.list_profiles(session_token.value),
                                visible=bool(profile_manager.list_profiles(session_token.value))
                            )
                            with gr.Row():
                                load_btn = gr.Button("Load", visible=bool(profile_manager.list_profiles(session_token.value)))
                                delete_btn = gr.Button("Delete", variant="stop", visible=bool(profile_manager.list_profiles(session_token.value)))
                        
                        clear_btn = gr.Button("Clear Form")
                    
                    with gr.Column(scale=2):
                        output_summary = gr.Markdown(
                            "Your profile summary will appear here after saving.",
                            label="Profile Summary"
                        )
                
                # Save profile
                def save_profile_and_update(*args):
                    # Extract inputs
                    inputs = args[:-1]  # All except the last which is tab_completed
                    current_tab_status = args[-1]
                    
                    try:
                        # Call the original save function
                        summary = profile_manager.save_profile(*inputs)
                        
                        # Update completion status
                        new_status = current_tab_status.copy()
                        new_status[3] = True
                        
                        return (
                            summary,
                            new_status,
                            gr.update(elem_classes="completed-tab"),
                            gr.update(interactive=True),
                            gr.update(visible=False))
                    except Exception as e:
                        logging.error(f"Save profile error: {str(e)}")
                        return (
                            f"Error saving profile: {str(e)}",
                            current_tab_status,
                            gr.update(),
                            gr.update(),
                            gr.update(visible=True, value=f"<div class='nav-message'>Error: {str(e)}</div>"))
                
                save_btn.click(
                    fn=save_profile_and_update,
                    inputs=[
                        name, age, interests, transcript_data, learning_output,
                        movie, movie_reason, show, show_reason,
                        book, book_reason, character, character_reason, blog_text,
                        tab_completed
                    ],
                    outputs=[output_summary, tab_completed, step4, step5, nav_message]
                ).then(
                    fn=lambda: profile_manager.list_profiles(session_token.value),
                    outputs=load_profile_dropdown
                ).then(
                    fn=lambda: gr.update(visible=bool(profile_manager.list_profiles(session_token.value))),
                    outputs=load_btn
                ).then(
                    fn=lambda: gr.update(visible=bool(profile_manager.list_profiles(session_token.value))),
                    outputs=delete_btn
                )
                
                # Load profile
                load_btn.click(
                    fn=lambda name: profile_manager.load_profile(name, session_token.value),
                    inputs=load_profile_dropdown,
                    outputs=output_summary
                )
                
                # Delete profile
                def delete_profile(name, session_token):
                    if not name:
                        raise gr.Error("Please select a profile to delete")
                    try:
                        profile_path = profile_manager.get_profile_path(name)
                        if profile_path.exists():
                            profile_path.unlink()
                        return "Profile deleted successfully", ""
                    except Exception as e:
                        logging.error(f"Delete profile error: {str(e)}")
                        raise gr.Error(f"Error deleting profile: {str(e)}")
                
                delete_btn.click(
                    fn=delete_profile,
                    inputs=[load_profile_dropdown, session_token],
                    outputs=[output_summary, load_profile_dropdown]
                ).then(
                    fn=lambda: gr.update(
                        choices=profile_manager.list_profiles(session_token.value),
                        visible=bool(profile_manager.list_profiles(session_token.value))
                    ),
                    outputs=load_profile_dropdown
                ).then(
                    fn=lambda: gr.update(visible=bool(profile_manager.list_profiles(session_token.value))),
                    outputs=load_btn
                ).then(
                    fn=lambda: gr.update(visible=bool(profile_manager.list_profiles(session_token.value))),
                    outputs=delete_btn
                )
                
                # Clear form
                clear_btn.click(
                    fn=lambda: [gr.update(value="") for _ in range(12)],
                    outputs=[
                        name, age, interests, 
                        movie, movie_reason, show, show_reason,
                        book, book_reason, character, character_reason,
                        blog_text
                    ]
                ).then(
                    fn=lambda: gr.update(value=""),
                    outputs=output_summary
                ).then(
                    fn=lambda: gr.update(value=False),
                    outputs=blog_checkbox
                ).then(
                    fn=lambda: gr.update(visible=False),
                    outputs=blog_text
                )
            
            # ===== TAB 5: AI Teaching Assistant =====
            with gr.Tab("AI Assistant", id=4) as tab5:
                gr.Markdown("## Your Personalized Learning Assistant")
                gr.Markdown("Ask me anything about studying, your courses, grades, or learning strategies.")
                
                # Chat interface with session token
                chatbot = gr.ChatInterface(
                    fn=lambda msg, hist: teaching_assistant.generate_response(msg, hist, session_token.value),
                    examples=[
                        "How should I study for my next math test?",
                        "What's my current GPA?",
                        "Show me my course history",
                        "How can I improve my grades in science?",
                        "What study methods match my learning style?"
                    ],
                    title=""
                )
        
        # Tab navigation logic with completion check
        def navigate_to_tab(tab_index: int, tab_completed_status):
            current_tab = tabs.selected
            
            # Allow backward navigation
            if tab_index <= current_tab:
                return gr.Tabs(selected=tab_index), gr.update(visible=False)
            
            # Check if current tab is completed
            if not tab_completed_status.get(current_tab, False):
                return (
                    gr.Tabs(selected=current_tab),
                    gr.update(value=f"⚠️ Complete Step {current_tab+1} first!", visible=True))
            
            return gr.Tabs(selected=tab_index), gr.update(visible=False)
        
        step1.click(
            fn=lambda idx, status: navigate_to_tab(idx, status),
            inputs=[gr.State(0), tab_completed],
            outputs=[tabs, nav_message]
        )
        step2.click(
            fn=lambda idx, status: navigate_to_tab(idx, status),
            inputs=[gr.State(1), tab_completed],
            outputs=[tabs, nav_message]
        )
        step3.click(
            fn=lambda idx, status: navigate_to_tab(idx, status),
            inputs=[gr.State(2), tab_completed],
            outputs=[tabs, nav_message]
        )
        step4.click(
            fn=lambda idx, status: navigate_to_tab(idx, status),
            inputs=[gr.State(3), tab_completed],
            outputs=[tabs, nav_message]
        )
        step5.click(
            fn=lambda idx, status: navigate_to_tab(idx, status),
            inputs=[gr.State(4), tab_completed],
            outputs=[tabs, nav_message]
        )
        
        # Model loading functions
        def load_selected_model(model_name, progress=gr.Progress()):
            try:
                model_loader.load_model(model_name, progress)
                if model_loader.loaded:
                    return gr.update(value=f"<div class='alert-box'>{model_name} loaded successfully!</div>", visible=True)
                else:
                    return gr.update(value=f"<div class='nav-message'>Failed to load model: {model_loader.error}</div>", visible=True)
            except Exception as e:
                logging.error(f"Model loading error: {str(e)}")
                return gr.update(value=f"<div class='nav-message'>Error: {str(e)}</div>", visible=True)
        
        load_model_btn.click(
            fn=load_selected_model,
            inputs=model_selector,
            outputs=model_status
        )
    
    return app

# Create the interface
app = create_interface()

# For Hugging Face Spaces deployment
if __name__ == "__main__":
    app.launch()