# -*- coding: utf-8 -*-
"""app
Automatically generated by Colab.
Original file is located at
    https://colab.research.google.com/drive/1pwwcBb5Zlw1DA3u5K8W8mjrwBTBWXc1L
"""

import gradio as gr
import numpy as np
import os
import time
import groq
import uuid
import re
import tempfile

# LangChain imports
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_groq import ChatGroq

# Other imports
import chardet
import fitz  # PyMuPDF for PDFs
import docx  # python-docx for Word files
import gtts  # Google Text-to-Speech library
from pptx import Presentation  # python-pptx for PowerPoint files

# Set API Key
groq.api_key = os.getenv("GROQ_API_KEY")

# Initialize Chat Model
chat_model = ChatGroq(model_name="llama-3.3-70b-versatile", api_key=groq.api_key)

# Initialize Embeddings and chromaDB
os.makedirs("chroma_db", exist_ok=True)
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = Chroma(
    embedding_function=embedding_model,
    persist_directory="chroma_db"
)

# Short-term memory for the LLM
chat_memory = []

# Prompt for quiz generation
quiz_prompt = """
You are an AI assistant specialized in education and assessment creation. Given an uploaded document or text, generate a quiz with a mix of multiple-choice questions (MCQs) and fill-in-the-blank questions. The quiz should be directly based on the key concepts, facts, and details from the provided material.
Generate 20 Questions.
Remove all unnecessary formatting generated by the LLM, including <think> tags, asterisks, markdown formatting, and any bold or italic text, as well as **, ###, ##, and # tags.
For each question:
- Provide 4 answer choices (for MCQs), with only one correct answer.
- Ensure fill-in-the-blank questions focus on key terms, phrases, or concepts from the document.
- Include an answer key for all questions.
- Ensure questions vary in difficulty and encourage comprehension rather than memorization.
- Additionally, implement an instant feedback mechanism:
    - When a user selects an answer, indicate whether it is correct or incorrect.
    - If incorrect, provide a brief explanation from the document to guide learning.
    - Ensure responses are concise and educational to enhance understanding.
Output Example:
1. Fill in the blank: The LLM Agent framework has a central decision-making unit called the _______________________.
Answer: Agent Core
Feedback: The Agent Core is the central component of the LLM Agent framework, responsible for managing goals, tool instructions, planning modules, memory integration, and agent persona.
2. What is the main limitation of LLM-based applications?
a) Limited token capacity
b) Lack of domain expertise
c) Prone to hallucination
d) All of the above
Answer: d) All of the above
Feedback: LLM-based applications have several limitations, including limited token capacity, lack of domain expertise, and being prone to hallucination, among others.
3. Given the following info, what is the value of P(jam|Rain)?
P(no Rain) = 0.8;
P(no Jam) = 0.2;
P(Rain|Jam) = 0.1
a) 0.016
b) 0.025
c) 0.1
d) 0.4
Answer: d) 0.4
Feedback: This question tests understanding of Bayes' Theorem by requiring the calculation of conditional probability using the given values.
"""

# Groq Whisper Transcriber - RELIABLE SOLUTION
class GroqWhisperTranscriber:
    def __init__(self):
        self.client = groq.Client(api_key=groq.api_key)
        print("✅ Groq Whisper transcriber initialized")
    
    def transcribe_audio(self, audio):
        """Transcribe audio using Groq's reliable Whisper API"""
        if audio is None:
            return "Please record audio first"
        
        try:
            sr, y = audio
            
            print(f"Audio received - Sample rate: {sr}, Length: {len(y)}")
            
            # Basic validation
            if len(y) == 0:
                return "Empty audio detected"
            
            # Convert to mono if stereo
            if y.ndim > 1:
                y = np.mean(y, axis=1)
            
            # Convert to proper format
            y = y.astype(np.float32)
            
            # Normalize audio
            max_val = np.max(np.abs(y))
            if max_val > 0:
                y = y / max_val
            
            # Check audio duration
            audio_duration = len(y) / sr
            print(f"Audio duration: {audio_duration:.2f} seconds")
            
            if audio_duration < 0.5:
                return "Audio too short. Speak for at least 1 second."
            
            if audio_duration > 60:
                return "Audio too long. Keep it under 60 seconds."
            
            # Convert to 16-bit PCM for WAV file
            y_int16 = (y * 32767).astype(np.int16)
            
            # Create temporary WAV file
            import scipy.io.wavfile
            
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                temp_path = f.name
            
            # Save as WAV file
            scipy.io.wavfile.write(temp_path, sr, y_int16)
            
            print("Sending to Groq Whisper API...")
            
            # Transcribe with Groq API - USE TURBO VERSION
            with open(temp_path, "rb") as audio_file:
                transcription = self.client.audio.transcriptions.create(
                    file=(temp_path, audio_file.read(), "audio/wav"),
                    model="whisper-large-v3-turbo",  # Use the best model
                    response_format="text",
                    language="en"  # Optional: specify English for better accuracy
                )
            
            # Clean up temporary file
            os.unlink(temp_path)
            
            text = transcription.strip()
            print(f"Groq transcription: '{text}'")
            
            if not text:
                return "No speech detected. Please try again."
            
            return text
            
        except Exception as e:
            print(f"Groq transcription error: {str(e)}")
            # Clean up temp file if it exists
            try:
                if 'temp_path' in locals():
                    os.unlink(temp_path)
            except:
                pass
            return f"Transcription failed: {str(e)}"

# Initialize transcriber
try:
    transcriber = GroqWhisperTranscriber()
    print("✅ Transcriber initialized successfully with Groq API")
except Exception as e:
    print(f"❌ Failed to initialize transcriber: {e}")
    transcriber = None

def transcribe_audio(audio):
    """Main transcription function"""
    if transcriber is None:
        return "Speech recognition not available"
    
    return transcriber.transcribe_audio(audio)

def get_transcription_status(audio):
    """Status updates"""
    if audio is None:
        return "Click record to start"
    
    try:
        sr, y = audio
        duration = len(y) / sr if sr > 0 else 0
        
        if duration < 0.5:
            return "Recording... (keep speaking)"
        elif duration > 10:
            return "Processing longer audio..."
        else:
            return "Processing audio with Groq API..."
    except:
        return "Ready to record"

# Function to clean AI response by removing unwanted formatting
def clean_response(response):
    """Removes <think> tags, asterisks, and markdown formatting."""
    cleaned_text = re.sub(r"<think>.*?</think>", "", response, flags=re.DOTALL)
    cleaned_text = re.sub(r"(\*\*|\*|\[|\])", "", cleaned_text)
    cleaned_text = re.sub(r"^##+\s*", "", cleaned_text, flags=re.MULTILINE)
    cleaned_text = re.sub(r"\\", "", cleaned_text)
    cleaned_text = re.sub(r"---", "", cleaned_text)
    return cleaned_text.strip()

# Function to generate quiz based on content
def generate_quiz(content):
    prompt = f"{quiz_prompt}\n\nDocument content:\n{content}"
    response = chat_model.invoke([HumanMessage(content=prompt)])
    cleaned_response = clean_response(response.content)
    return cleaned_response

# Function to retrieve relevant documents from vectorstore based on user query
def retrieve_documents(query):
    results = vectorstore.similarity_search(query, k=3)
    return [doc.page_content for doc in results]

# Function to convert tuple format to message format
def convert_to_message_format(chat_history):
    message_format = []
    for user_msg, bot_msg in chat_history:
        message_format.append({"role": "user", "content": user_msg})
        message_format.append({"role": "assistant", "content": bot_msg})
    return message_format

# Function to convert message format to tuple format for processing
def convert_to_tuple_format(chat_history):
    tuple_format = []
    for i in range(0, len(chat_history), 2):
        if i+1 < len(chat_history):
            user_msg = chat_history[i]["content"]
            bot_msg = chat_history[i+1]["content"]
            tuple_format.append((user_msg, bot_msg))
    return tuple_format

# Function to handle chatbot interactions with short-term memory
def chat_with_groq(user_input, chat_history):
    try:
        # Convert message format to tuple format for processing
        tuple_history = convert_to_tuple_format(chat_history)
        
        # Retrieve relevant documents for additional context
        relevant_docs = retrieve_documents(user_input)
        context = "\n".join(relevant_docs) if relevant_docs else "No relevant documents found."

        # Construct proper prompting with conversation history
        system_prompt = "You are a helpful AI assistant. Answer questions accurately and concisely."
        conversation_history = "\n".join(chat_memory[-10:])
        prompt = f"{system_prompt}\n\nConversation History:\n{conversation_history}\n\nUser Input: {user_input}\n\nContext:\n{context}"

        # Call the chat model
        response = chat_model.invoke([HumanMessage(content=prompt)])

        # Clean response to remove any unwanted formatting
        cleaned_response_text = clean_response(response.content)

        # Append conversation history
        chat_memory.append(f"User: {user_input}")
        chat_memory.append(f"AI: {cleaned_response_text}")

        # Update chat history
        chat_history.append({"role": "user", "content": user_input})
        chat_history.append({"role": "assistant", "content": cleaned_response_text})

        # Convert response to speech
        audio_file = speech_playback(cleaned_response_text)

        return chat_history, "", audio_file
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        chat_history.append({"role": "user", "content": user_input})
        chat_history.append({"role": "assistant", "content": error_msg})
        return chat_history, "", None

# Function to play response as speech using gTTS
def speech_playback(text):
    try:
        # Generate a unique filename for each audio file
        unique_id = str(uuid.uuid4())
        audio_file = f"output_audio_{unique_id}.mp3"

        # Convert text to speech
        tts = gtts.gTTS(text, lang='en')
        tts.save(audio_file)

        # Return the path to the audio file
        return audio_file
    except Exception as e:
        print(f"Error in speech_playback: {e}")
        return None

# Function to detect encoding safely
def detect_encoding(file_path):
    try:
        with open(file_path, "rb") as f:
            raw_data = f.read(4096)
            detected = chardet.detect(raw_data)
            encoding = detected["encoding"]
        return encoding if encoding else "utf-8"
    except Exception:
        return "utf-8"

# Function to extract text from PDF
def extract_text_from_pdf(pdf_path):
    try:
        doc = fitz.open(pdf_path)
        text = "\n".join([page.get_text("text") for page in doc])
        return text if text.strip() else "No extractable text found."
    except Exception as e:
        return f"Error extracting text from PDF: {str(e)}"

# Function to extract text from Word files (.docx)
def extract_text_from_docx(docx_path):
    try:
        doc = docx.Document(docx_path)
        text = "\n".join([para.text for para in doc.paragraphs])
        return text if text.strip() else "No extractable text found."
    except Exception as e:
        return f"Error extracting text from Word document: {str(e)}"

# Function to extract text from PowerPoint files (.pptx)
def extract_text_from_pptx(pptx_path):
    try:
        presentation = Presentation(pptx_path)
        text = ""
        for slide in presentation.slides:
            for shape in slide.shapes:
                if hasattr(shape, "text"):
                    text += shape.text + "\n"
        return text if text.strip() else "No extractable text found."
    except Exception as e:
        return f"Error extracting text from PowerPoint: {str(e)}"

# Function to process documents safely
def process_document(file):
    try:
        file_extension = os.path.splitext(file.name)[-1].lower()
        if file_extension in [".png", ".jpg", ".jpeg"]:
            return "Error: Images cannot be processed for text extraction."
        if file_extension == ".pdf":
            content = extract_text_from_pdf(file.name)
        elif file_extension == ".docx":
            content = extract_text_from_docx(file.name)
        elif file_extension == ".pptx":
            content = extract_text_from_pptx(file.name)
        else:
            encoding = detect_encoding(file.name)
            with open(file.name, "r", encoding=encoding, errors="replace") as f:
                content = f.read()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
        documents = [Document(page_content=chunk) for chunk in text_splitter.split_text(content)]
        vectorstore.add_documents(documents)

        quiz = generate_quiz(content)
        return f"Document processed successfully (File Type: {file_extension}). Quiz generated:\n{quiz}"
    except Exception as e:
        return f"Error processing document: {str(e)}"

# Clear chat history function
def clear_chat_history():
    chat_memory.clear()
    return [], None

def tutor_ai_chatbot():
    """Main Gradio interface for the Tutor AI Chatbot."""
    with gr.Blocks() as app:
        gr.Markdown("# AI Tutor - We.(POC)")
        gr.Markdown("An interactive Personal AI Tutor chatbot to help with your learning needs.")

        # Chatbot Tab
        with gr.Tab("AI Chatbot"):
            with gr.Row():
                with gr.Column(scale=3):
                    chatbot = gr.Chatbot(height=500, type="messages")
                    
                with gr.Column(scale=1):
                    audio_playback = gr.Audio(label="Audio Response", type="filepath")
            
            # Move the input controls here to span full width
            with gr.Row():
                msg = gr.Textbox(
                    label="Ask a question", 
                    placeholder="Type your question here...",
                    container=False
                )
                submit = gr.Button("Send")
            
            with gr.Row():
                with gr.Column(scale=1):
                    audio_input = gr.Audio(type="numpy", label="Record or Upload Audio")
            
            # Add transcription status indicator
            transcription_status = gr.Textbox(
                label="Transcription Status",
                interactive=False,
                value="Click record to start",
                max_lines=2
            )
            
            # Voice recording tips - ONLY in AI Chatbot tab
            with gr.Accordion("Voice Recording Tips", open=False):
                gr.Markdown("""
                **For perfect transcription:**
                - 🎤 Speak clearly and directly into microphone
                - 🔇 Record in QUIET environment (no background noise)
                - 📏 Keep recording between 2-10 seconds
                - 🗣️ Speak at normal volume and pace
                - 📱 Use a good quality microphone
                
                **Using Distill Whisper API:**
                - ✅ High accuracy transcription
                - ✅ No more "B-B-B" or "oh-oh-oh" errors
                - ✅ Fast and reliable
                - ✅ Professional grade speech recognition
                """)
            
            # Clear chat history button
            clear_btn = gr.Button("Clear Chat")

            # Handle chat interaction
            submit.click(
                chat_with_groq,
                inputs=[msg, chatbot],
                outputs=[chatbot, msg, audio_playback]
            )

            # Clear chat history function
            clear_btn.click(
                lambda: [],
                inputs=None, 
                outputs=[chatbot]
            )

            # Also allow Enter key to submit
            msg.submit(
                chat_with_groq,
                inputs=[msg, chatbot],
                outputs=[chatbot, msg, audio_playback]
            )

            # Add some examples of questions students might ask
            with gr.Accordion("Example Questions", open=False):
                gr.Examples(
                    examples=[
                        "Can you explain the concept of RLHF AI?",
                        "What are AI transformers?",
                        "What is MoE AI?",
                        "What's gate networks AI?",
                        "I am making a switch, please generating baking recipe?"
                    ],
                    inputs=msg
                )

            # Connect audio input to transcription with status updates
            audio_input.change(
                fn=get_transcription_status,
                inputs=audio_input,
                outputs=transcription_status
            ).then(
                fn=transcribe_audio,
                inputs=audio_input,
                outputs=msg
            ).then(
                fn=lambda x: "✅ Transcription completed!" if x and "failed" not in x.lower() and "error" not in x.lower() and "sorry" not in x.lower() else "Ready for new recording",
                inputs=msg,
                outputs=transcription_status
            )

        # Upload Notes & Generate Quiz Tab
        with gr.Tab("Upload Notes & Generate Quiz"):
            with gr.Row():
                with gr.Column(scale=2):
                    file_input = gr.File(label="Upload Lecture Notes (PDF, DOCX, PPTX)")
                with gr.Column(scale=3):
                    quiz_output = gr.Textbox(label="Generated Quiz", lines=10)

            # Connect file input to document processing
            file_input.change(process_document, inputs=file_input, outputs=quiz_output)

        # Introduction Video Tab - Now with the working video
        with gr.Tab("Introduction Video"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Welcome to the Introduction Video")
                    gr.Markdown("Music from Xu Mengyuan - China-O, musician Xu Mengyuan YUAN! | 徐梦圆 - China-O 音乐人徐梦圆YUAN!")
                    # Use the local video file that's stored in your Space
                    gr.Video("We_not_me_video.mp4", label="Introduction Video")

        # Launch the application
        app.launch(share=False)

# Launch the AI chatbot
if __name__ == "__main__":
    tutor_ai_chatbot()