import gradio as gr
from transformers import pipeline, AutoTokenizer
from gtts import gTTS
import os
import random
import re

# Initialize NLP pipelines
qa = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
try:
    summarizer = pipeline("summarization", model="t5-small")
except Exception as e:
    print(f"Error loading summarizer: {e}")
    summarizer = None

# Load tokenizer for rule-based question generation
try:
    tokenizer = AutoTokenizer.from_pretrained("t5-small")
except Exception as e:
    print(f"Error loading tokenizer: {e}")
    tokenizer = None

# Initialize user stats and log files
if not os.path.exists("decision_log.txt"):
    with open("decision_log.txt", "w") as f:
        f.write("Decision Log Initialized\n")
if not os.path.exists("feedback.txt"):
    with open("feedback.txt", "w") as f:
        f.write("Feedback Log Initialized\n")
if not os.path.exists("user_score.txt"):
    with open("user_score.txt", "w") as f:
        f.write("0")
if not os.path.exists("questions_answered.txt"):
    with open("questions_answered.txt", "w") as f:
        f.write("0")
if not os.path.exists("avatar.txt"):
    with open("avatar.txt", "w") as f:
        f.write("default")

# Avatar options
avatars = {
    "default": "https://cdn.pixabay.com/photo/2016/08/08/09/17/avatar-1577909_1280.png",
    "robot": "https://cdn.pixabay.com/photo/2017/01/31/21/23/robot-2027195_1280.png",
    "cat": "https://cdn.pixabay.com/photo/2017/01/31/21/23/cat-2027196_1280.png",
    "dog": "https://cdn.pixabay.com/photo/2017/01/31/21/23/dog-2027197_1280.png"
}

def update_score(points):
    with open("user_score.txt", "r") as f:
        score = int(f.read())
    score += points
    with open("user_score.txt", "w") as f:
        f.write(str(score))
    return score

def update_questions_answered():
    with open("questions_answered.txt", "r") as f:
        count = int(f.read())
    count += 1
    with open("questions_answered.txt", "w") as f:
        f.write(str(count))
    return count

def get_progress():
    with open("questions_answered.txt", "r") as f:
        count = int(f.read())
    progress = min(count * 10, 100)  # 10% per question, max 100%
    return f"Progress: {progress}%"

def get_motivational_message():
    messages = [
        "Great job! Keep learning!",
        "You're doing awesome!",
        "Amazing effort! Keep it up!",
        "You're making great progress!"
    ]
    return random.choice(messages)

def get_hint(context):
    try:
        if len(context.split()) < 10:
            return "Hint: The answer is in the context, but it's too short to summarize."
        summary = summarizer(context, max_length=30, min_length=10)[0]["summary_text"]
        return f"Hint: {summary}"
    except Exception as e:
        return f"Hint unavailable: {str(e)}"

def set_avatar(avatar_choice):
    with open("avatar.txt", "w") as f:
        f.write(avatar_choice)
    return avatars[avatar_choice]

def get_avatar():
    with open("avatar.txt", "r") as f:
        avatar_choice = f.read().strip()
    return avatars.get(avatar_choice, avatars["default"])

def study_aid(question, context, font_size=16, audio_output=False, simplify_text=False, theme="dark"):
    with open("decision_log.txt", "a") as f:
        f.write(f"Question: {question}, Simplified: {simplify_text}, Audio: {audio_output}, Font: {font_size}, Theme: {theme}\n")
    
    simplified_context = context
    if simplify_text and summarizer is not None:
        try:
            if len(context.split()) < 10:
                simplified_context = "Input too short to simplify."
            elif len(context.split()) > 512:
                simplified_context = "Input too long to simplify."
            else:
                summary = summarizer(context, max_length=100, min_length=50)[0]["summary_text"]
                simplified_context = summary
        except Exception as e:
            simplified_context = f"Error simplifying text: {str(e)}"
    
    answer = qa(question=question, context=simplified_context)["answer"]
    
    bg_color = "black" if theme == "dark" else "white"
    text_color = "white" if theme == "dark" else "black"
    output = f"<div style='font-size:{font_size}px; color:{text_color}; background-color:{bg_color}; padding:10px;'>"
    if simplify_text and simplified_context != context:
        output += f"<b>Simplified Context:</b> {simplified_context}<br>"
    output += f"<b>Answer:</b> {answer}</div>"
    
    # Add visual diagram for neural network questions
    diagram = None
    if "neural network" in question.lower():
        diagram = "https://upload.wikimedia.org/wikipedia/commons/thumb/4/46/Colored_neural_network.svg/300px-Colored_neural_network.svg.png"
        output += f"<br><img src='{diagram}' alt='Neural Network Diagram' style='width:300px; height:auto;'>"
    
    # Update stats
    score = update_score(10)
    questions_answered = update_questions_answered()
    progress = get_progress()
    motivation = get_motivational_message()
    
    # Sound effect for points earned
    sound_effect = "https://www.soundjay.com/buttons/sounds/beep-01a.mp3"
    
    if audio_output:
        tts = gTTS(text=answer, lang='en')
        tts.save("answer_audio.mp3")
        return output, "answer_audio.mp3", sound_effect, f"Your Score: {score} | {progress} | {motivation}"
    
    return output, None, sound_effect, f"Your Score: {score} | {progress} | {motivation}"

def submit_feedback(feedback):
    with open("feedback.txt", "a") as f:
        f.write(feedback + "\n")
    score = update_score(5)
    progress = get_progress()
    motivation = get_motivational_message()
    sound_effect = "https://www.soundjay.com/buttons/sounds/beep-01a.mp3"
    return f"Feedback submitted! Your Score: {score} | {progress} | {motivation}", sound_effect

def generate_quiz(context, theme="dark"):
    if tokenizer is None:
        return "Question generation not available: Tokenizer failed to load.", None, None, None
    
    try:
        # Simple rule-based question generation
        # Tokenize the context and extract key phrases (e.g., noun phrases)
        tokens = context.split()
        key_phrases = []
        current_phrase = []
        for token in tokens:
            if token in [",", ".", ":", ";", "!", "?", "and", "or", "but"]:
                if current_phrase:
                    key_phrases.append(" ".join(current_phrase))
                    current_phrase = []
            else:
                current_phrase.append(token)
        if current_phrase:
            key_phrases.append(" ".join(current_phrase))
        
        # Filter phrases that are likely to be meaningful (e.g., longer than 2 words)
        key_phrases = [phrase for phrase in key_phrases if len(phrase.split()) > 2]
        if not key_phrases:
            return "No suitable phrases found for question generation.", None, None, None
        
        # Generate a question using the first key phrase
        quiz_question = f"What is {key_phrases[0]}?"
        answer = qa(question=quiz_question, context=context)["answer"]
        
        bg_color = "black" if theme == "dark" else "white"
        text_color = "white" if theme == "dark" else "black"
        output = f"<div style='color:{text_color}; background-color:{bg_color}; padding:10px;'>"
        output += f"<b>Quiz Question:</b> {quiz_question}<br><b>Answer:</b> {answer}</div>"
        
        tts = gTTS(text=answer, lang='en')
        tts.save("quiz_audio.mp3")
        
        score = update_score(20)
        questions_answered = update_questions_answered()
        progress = get_progress()
        motivation = get_motivational_message()
        sound_effect = "https://www.soundjay.com/buttons/sounds/beep-01a.mp3"
        return output, "quiz_audio.mp3", sound_effect, f"Your Score: {score} | {progress} | {motivation}"
    except Exception as e:
        return f"Error generating quiz: {str(e)}", None, None, None

def read_logs():
    try:
        with open("decision_log.txt", "r") as f:
            return f.read()
    except FileNotFoundError:
        return "Decision log not found. It will be created once you start using the app."

with gr.Blocks(title="StudyBuddy: Accessible Study Aid for Neurodiverse Students") as app:
    gr.Markdown(
        """
        # StudyBuddy: Accessible Study Aid for Neurodiverse Students
        Ask questions about your college lecture notes with accessible text and audio outputs. No data is stored.
        """
    )
    
    with gr.Row():
        gr.Image(value=get_avatar(), label="Your Avatar", width=100, height=100)
        avatar_input = gr.Dropdown(choices=list(avatars.keys()), value="default", label="Choose Avatar")
        avatar_input.change(fn=set_avatar, inputs=avatar_input, outputs=gr.Image(label="Your Avatar", width=100, height=100))
    
    with gr.Tab("Ask a Question"):
        question_input = gr.Textbox(label="Question", placeholder="e.g., What is machine learning?")
        context_input = gr.Textbox(label="Context (Lecture Notes)", placeholder="Paste your notes here...")
        font_size_input = gr.Slider(12, 24, value=16, label="Font Size (px)")
        theme_input = gr.Dropdown(choices=["dark", "light"], value="dark", label="Theme")
        audio_output_input = gr.Checkbox(label="Generate Audio Output")
        simplify_text_input = gr.Checkbox(label="Simplify Text")
        with gr.Row():
            study_submit_btn = gr.Button("Get Answer")
            hint_btn = gr.Button("Get Hint")
        study_output_text = gr.HTML(label="Answer")
        study_output_audio = gr.Audio(label="Audio Narration")
        study_output_sound = gr.Audio(label="Sound Effect", visible=False)
        score_output = gr.Text(label="Score & Progress")
        hint_output = gr.Text(label="Hint")
        
        study_submit_btn.click(
            fn=study_aid,
            inputs=[question_input, context_input, font_size_input, audio_output_input, simplify_text_input, theme_input],
            outputs=[study_output_text, study_output_audio, study_output_sound, score_output]
        )
        hint_btn.click(
            fn=get_hint,
            inputs=context_input,
            outputs=hint_output
        )
    
    with gr.Tab("Quiz Me"):
        quiz_context_input = gr.Textbox(label="Context (Lecture Notes)", placeholder="Paste your notes here...")
        quiz_theme_input = gr.Dropdown(choices=["dark", "light"], value="dark", label="Theme")
        quiz_submit_btn = gr.Button("Generate Quiz Question")
        quiz_output_text = gr.HTML(label="Quiz Question and Answer")
        quiz_output_audio = gr.Audio(label="Audio Narration")
        quiz_output_sound = gr.Audio(label="Sound Effect", visible=False)
        quiz_score_output = gr.Text(label="Score & Progress")
        
        quiz_submit_btn.click(
            fn=generate_quiz,
            inputs=[quiz_context_input, quiz_theme_input],
            outputs=[quiz_output_text, quiz_output_audio, quiz_output_sound, quiz_score_output]
        )
    
    with gr.Tab("Submit Feedback"):
        feedback_input = gr.Textbox(label="Feedback", placeholder="Report issues or suggestions...")
        feedback_submit_btn = gr.Button("Submit Feedback")
        feedback_output = gr.Text(label="Feedback Status")
        feedback_sound = gr.Audio(label="Sound Effect", visible=False)
        
        feedback_submit_btn.click(
            fn=submit_feedback,
            inputs=feedback_input,
            outputs=[feedback_output, feedback_sound]
        )

    with gr.Tab("View Logs"):
        logs_output = gr.Textbox(label="Decision Logs", value=read_logs())

app.launch()