import gradio as gr
import PyPDF2
from io import BytesIO
import tempfile
import os
from pptx import Presentation
import requests
import json
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
from gtts import gTTS
import numpy as np
import re

# Global variables to store models and data
tokenizer = None
model = None
tts_pipeline = None
uploaded_docs = {}
current_doc = None

def load_models():
    """Load free models from Hugging Face"""
    global tokenizer, model, tts_pipeline
    
    try:
        # Text generation model (free and good for conversations)
        model_name = "microsoft/DialoGPT-medium"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name)
        
        # Add padding token if it doesn't exist
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        
        print("✅ Models loaded successfully!")
        return "Models loaded successfully!"
        
    except Exception as e:
        print(f"❌ Error loading models: {e}")
        return f"Error loading models: {e}"

def extract_text_from_pdf(pdf_file):
    """Extract text from PDF file"""
    try:
        if pdf_file is None:
            return ""
        
        pdf_reader = PyPDF2.PdfReader(pdf_file.name)
        text = ""
        for page_num, page in enumerate(pdf_reader.pages, 1):
            page_text = page.extract_text()
            text += f"\n--- Page {page_num} ---\n{page_text}\n"
        
        return text.strip()
    except Exception as e:
        return f"Error reading PDF: {str(e)}"

def extract_text_from_ppt(ppt_file):
    """Extract text from PowerPoint file"""
    try:
        if ppt_file is None:
            return ""
        
        presentation = Presentation(ppt_file.name)
        text = ""
        
        for slide_num, slide in enumerate(presentation.slides, 1):
            text += f"\n--- Slide {slide_num} ---\n"
            for shape in slide.shapes:
                if hasattr(shape, "text") and shape.text.strip():
                    text += shape.text + "\n"
        
        return text.strip()
    except Exception as e:
        return f"Error reading PowerPoint: {str(e)}"

def process_document(file, doc_type):
    """Process uploaded document and extract text"""
    global uploaded_docs, current_doc
    
    if file is None:
        return "❌ No file uploaded", "Please upload a document first."
    
    try:
        # Extract text based on file type
        if doc_type == "PDF":
            text_content = extract_text_from_pdf(file)
        elif doc_type == "PowerPoint":
            text_content = extract_text_from_ppt(file)
        else:
            return "❌ Unsupported file type", "Please upload a PDF or PowerPoint file."
        
        if not text_content or text_content.startswith("Error"):
            return f"❌ Failed to process {file.name}", text_content
        
        # Store document
        doc_name = file.name
        uploaded_docs[doc_name] = {
            'content': text_content,
            'type': doc_type,
            'file_path': file.name
        }
        current_doc = doc_name
        
        # Create preview (first 500 characters)
        preview = text_content[:500] + "..." if len(text_content) > 500 else text_content
        
        return f"✅ Successfully processed: {doc_name}", f"Document Preview:\n\n{preview}"
        
    except Exception as e:
        return f"❌ Error processing {file.name}", f"Error: {str(e)}"

def generate_response(user_input, history, tutor_mode):
    """Generate AI response based on user input and document context"""
    global model, tokenizer, current_doc, uploaded_docs
    
    if not user_input.strip():
        return history, ""
    
    if current_doc is None or current_doc not in uploaded_docs:
        response = "Please upload and process a document first before asking questions."
        history.append([user_input, response])
        return history, ""
    
    try:
        # Get document context
        doc_content = uploaded_docs[current_doc]['content']
        
        # Create context-aware prompt based on tutor mode
        mode_contexts = {
            "Explain Concepts": "As an AI tutor, explain the following concept clearly and simply based on the document content:",
            "Quiz Mode": "Create a quiz question or test the user's understanding of:",
            "Practice Problems": "Provide practice exercises or real-world applications for:"
        }
        
        context_prompt = mode_contexts.get(tutor_mode, "Help me understand:")
        
        # Limit document content to avoid token limits
        limited_content = doc_content[:1000] + "..." if len(doc_content) > 1000 else doc_content
        
        # Create conversation prompt
        prompt = f"{context_prompt} {user_input}\n\nDocument context: {limited_content}\n\nResponse:"
        
        # Generate response using the model
        inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
        
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=150,
                num_return_sequences=1,
                temperature=0.7,
                pad_token_id=tokenizer.eos_token_id,
                do_sample=True,
                top_p=0.9
            )
        
        # Decode response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the new generated part
        response = response[len(prompt):].strip()
        
        # Fallback if response is empty or too short
        if len(response) < 10:
            fallback_responses = {
                "Explain Concepts": f"Based on your document, {user_input} is an important concept. From what I can see in your materials, this topic involves several key aspects that are worth understanding in detail.",
                "Quiz Mode": f"Here's a question about {user_input}: Based on your document, what are the main points or key takeaways regarding this topic?",
                "Practice Problems": f"Let's practice with {user_input}. Try to apply the concepts from your document to solve a real-world scenario involving this topic."
            }
            response = fallback_responses.get(tutor_mode, f"Great question about {user_input}! From your document, I can help you understand this concept better.")
        
        # Clean up response
        response = response.replace(prompt, "").strip()
        if not response:
            response = f"I understand you're asking about {user_input}. Based on your document, this is an important topic that deserves careful explanation."
        
        # Add to history
        history.append([user_input, response])
        
        return history, ""
        
    except Exception as e:
        error_response = f"I apologize, but I'm having trouble processing your question right now. However, I can tell you that {user_input} is mentioned in your document and is worth exploring further."
        history.append([user_input, error_response])
        return history, ""

def text_to_speech(text):
    """Convert text to speech using gTTS"""
    try:
        if not text or len(text.strip()) == 0:
            return None
        
        # Clean text for TTS
        clean_text = re.sub(r'[^\w\s.,!?]', '', text)
        clean_text = clean_text[:500]  # Limit length for TTS
        
        if len(clean_text.strip()) == 0:
            return None
        
        # Generate speech
        tts = gTTS(text=clean_text, lang='en', slow=False)
        
        # Save to temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
            tts.save(tmp_file.name)
            return tmp_file.name
            
    except Exception as e:
        print(f"TTS Error: {e}")
        return None

def chat_with_speech(user_input, history, tutor_mode):
    """Chat function that includes speech output"""
    # Generate text response
    updated_history, _ = generate_response(user_input, history, tutor_mode)
    
    # Get the last AI response for TTS
    if updated_history and len(updated_history) > 0:
        last_response = updated_history[-1][1]
        audio_file = text_to_speech(last_response)
        return updated_history, "", audio_file
    
    return updated_history, "", None

def get_document_info():
    """Get information about currently loaded document"""
    global current_doc, uploaded_docs
    
    if current_doc and current_doc in uploaded_docs:
        doc = uploaded_docs[current_doc]
        word_count = len(doc['content'].split())
        return f"📄 Current Document: {current_doc}\n📊 Type: {doc['type']}\n📝 Word Count: ~{word_count} words\n✅ Ready for tutoring!"
    else:
        return "❌ No document loaded. Please upload a PDF or PowerPoint file."

# Load models on startup
print("Loading AI models...")
load_models()

# Create Gradio interface
with gr.Blocks(title="🧠 AI Tutor - Free Models", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🧠 AI Tutor - Free Models
    
    Upload your PDF or PowerPoint files and start learning with AI-powered tutoring!
    
    **Features:**
    - 📄 Support for PDF and PowerPoint files
    - 🤖 AI-powered tutoring with free Hugging Face models
    - 🔊 Text-to-speech for AI responses
    - 📚 Multiple learning modes (Explain, Quiz, Practice)
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## 📁 Document Upload")
            
            # File upload
            file_input = gr.File(
                label="Upload Document",
                file_types=[".pdf", ".pptx", ".ppt"],
                type="filepath"
            )
            
            doc_type = gr.Radio(
                choices=["PDF", "PowerPoint"],
                label="Document Type",
                value="PDF"
            )
            
            process_btn = gr.Button("📊 Process Document", variant="primary")
            
            # Document status
            doc_status = gr.Textbox(
                label="Status",
                interactive=False,
                placeholder="Upload a document to get started..."
            )
            
            # Document preview
            doc_preview = gr.Textbox(
                label="Document Preview",
                interactive=False,
                lines=8,
                placeholder="Document content will appear here..."
            )
            
            # Current document info
            gr.Markdown("## 📋 Document Info")
            doc_info = gr.Textbox(
                label="Current Document",
                interactive=False,
                lines=4
            )
            
            # Update document info periodically
            doc_info_btn = gr.Button("🔄 Refresh Info")
        
        with gr.Column(scale=2):
            gr.Markdown("## 💬 AI Tutor Chat")
            
            # Tutor mode selection
            tutor_mode = gr.Radio(
                choices=["Explain Concepts", "Quiz Mode", "Practice Problems"],
                label="🎯 Learning Mode",
                value="Explain Concepts"
            )
            
            # Chat interface
            chatbot = gr.Chatbot(
                label="Chat with AI Tutor",
                height=400,
                bubble_full_width=False
            )
            
            with gr.Row():
                msg_input = gr.Textbox(
                    label="Your Question",
                    placeholder="Ask me anything about your document...",
                    scale=4
                )
                send_btn = gr.Button("📤 Send", scale=1, variant="primary")
            
            # Audio output
            audio_output = gr.Audio(
                label="🔊 AI Response (Audio)",
                type="filepath",
                autoplay=True
            )
            
            # Clear chat button
            clear_btn = gr.Button("🗑️ Clear Chat")
    
    # Event handlers
    process_btn.click(
        fn=process_document,
        inputs=[file_input, doc_type],
        outputs=[doc_status, doc_preview]
    )
    
    send_btn.click(
        fn=chat_with_speech,
        inputs=[msg_input, chatbot, tutor_mode],
        outputs=[chatbot, msg_input, audio_output]
    )
    
    msg_input.submit(
        fn=chat_with_speech,
        inputs=[msg_input, chatbot, tutor_mode],
        outputs=[chatbot, msg_input, audio_output]
    )
    
    clear_btn.click(
        fn=lambda: ([], None),
        outputs=[chatbot, audio_output]
    )
    
    doc_info_btn.click(
        fn=get_document_info,
        outputs=[doc_info]
    )
    
    # Load document info on startup
    demo.load(
        fn=get_document_info,
        outputs=[doc_info]
    )

    gr.Markdown("""
    ## 🚀 How to Use:
    1. **Upload** your PDF or PowerPoint file
    2. **Process** the document to extract text
    3. **Choose** your learning mode (Explain, Quiz, or Practice)
    4. **Start chatting** with the AI tutor about your document
    5. **Listen** to AI responses with text-to-speech
    
    ## 🔧 Models Used:
    - **Text Generation**: Microsoft DialoGPT-medium (Free)
    - **Text-to-Speech**: Google TTS (gTTS) - Free
    - **Document Processing**: PyPDF2 & python-pptx (Free)
    """)

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