import streamlit as st
from youtube_transcript_api import YouTubeTranscriptApi
import re
import tempfile
import os
import warnings
from groq import Groq

# Set up Groq client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# Supported file types for Groq API
SUPPORTED_FILE_TYPES = ["mp3", "mp4", "mpeg", "mpga", "m4a", "wav", "webm"]

# Function to transcribe audio using Groq Whisper API
def transcribe_audio(file_path):
    file_extension = os.path.splitext(file_path)[1][1:]
    if file_extension not in SUPPORTED_FILE_TYPES:
        return f"Error: Unsupported file type '{file_extension}'. Please upload a valid file."
    
    try:
        with open(file_path, "rb") as file:
            transcription = client.audio.transcriptions.create(
                file=(file_path, file.read()),
                model="whisper-large-v3",
            )
        return transcription.text
    except Exception as e:
        return f"Error during transcription: {e}"

# Function to get transcript from YouTube
def get_transcript(url):
    try:
        video_id_match = re.search(r"(?:v=|\/)([0-9A-Za-z_-]{11}).*", url)
        if not video_id_match:
            return "Error: Invalid YouTube URL"
        video_id = video_id_match.group(1)
        transcript = YouTubeTranscriptApi.get_transcript(video_id)
        transcript_text = ' '.join([entry['text'] for entry in transcript])
        return transcript_text
    except Exception as e:
        return str(e)

# Function to summarize text using Groq API
def summarize_text(text):
    try:
        response = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": f"Summarize the following text:\n\n{text}"
                }
            ],
            model="llama3-8b-8192",
        )
        summary = response.choices[0].message.content.strip()
        return summary
    except Exception as e:
        return f"Error summarizing text: {e}"

# Function to generate quiz questions using Groq API
def generate_quiz_questions(text):
    try:
        response = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": f"Generate quiz questions for the following text:\n\n{text}"
                }
            ],
            model="llama3-8b-8192",
        )
        quiz_questions = response.choices[0].message.content.strip()
        return quiz_questions
    except Exception as e:
        return f"Error generating quiz questions: {e}"

# Function to parse quiz questions from generated text
def parse_quiz_questions(quiz_text):
    questions = []
    question_blocks = quiz_text.split("\n\n")
    current_question = None
    current_choices = []
    correct_answer = None
    
    for block in question_blocks:
        lines = block.strip().split("\n")
        if lines:
            if re.match(r'^\d+\.', lines[0]):  # This line is a question number
                if current_question and current_choices and correct_answer:
                    questions.append({
                        "question": current_question,
                        "choices": current_choices,
                        "correct_answer": correct_answer
                    })
                current_question = lines[0]
                current_choices = lines[1:5]
                correct_answer = lines[-1].split(": ")[-1].strip() if len(lines) > 5 else None
            else:  # This line is an answer
                correct_answer = lines[-1].split(": ")[-1].strip()

    # Add the last question if it exists
    if current_question and current_choices and correct_answer:
        questions.append({
            "question": current_question,
            "choices": current_choices,
            "correct_answer": correct_answer
        })
    
    return questions

# Function to generate explanation for quiz answers using Groq API
def generate_explanation(question, correct_answer, user_answer):
    try:
        response = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": f"Explain why the correct answer to the following question is '{correct_answer}' and not '{user_answer}':\n\n{question}"
                }
            ],
            model="llama3-8b-8192",
        )
        explanation = response.choices[0].message.content.strip()
        return explanation
    except Exception as e:
        return f"Error generating explanation: {e}"

# Function to check answers and provide feedback
def check_answers(questions, user_answers):
    feedback = []
    correct_count = 0
    for i, question in enumerate(questions):
        correct_answer = question['correct_answer']
        user_answer = user_answers.get(f"question_{i+1}", "")
        if user_answer == correct_answer:
            feedback.append({
                "question": question['question'],
                "user_answer": user_answer,
                "correct_answer": correct_answer,
                "status": "Correct"
            })
            correct_count += 1
        else:
            explanation = generate_explanation(question['question'], correct_answer, user_answer)
            feedback.append({
                "question": question['question'],
                "user_answer": user_answer,
                "correct_answer": correct_answer,
                "status": "Incorrect",
                "explanation": explanation
            })
    return feedback

# Function to handle uploaded files
def handle_uploaded_file(uploaded_file):
    file_path = tempfile.mktemp(suffix=os.path.splitext(uploaded_file.name)[1])
    with open(file_path, "wb") as f:
        f.write(uploaded_file.read())
    return file_path

# Streamlit app layout and functionality
st.title("YouTube Transcript Quiz Generator")
st.markdown("**Instructions:** Paste a YouTube link or upload a media file to generate a quiz.")

option = st.selectbox("Choose input type", ("YouTube URL", "Upload audio/video file"))

if "generated_quiz" not in st.session_state:
    st.session_state.generated_quiz = False

if option == "YouTube URL":
    url = st.text_input("YouTube URL", value="")
    if url:
        if st.button("Generate Quiz"):
            transcript_text = get_transcript(url)
            if "Error" not in transcript_text:
                summary = summarize_text(transcript_text)
                quiz_text = generate_quiz_questions(transcript_text)
                questions = parse_quiz_questions(quiz_text)

                if not questions:
                    st.error("No valid quiz questions could be generated.")
                else:
                    st.session_state.summary = summary
                    st.session_state.questions = questions
                    st.session_state.user_answers = {}
                    st.session_state.generated_quiz = True
            else:
                st.error(transcript_text)

if option == "Upload audio/video file":
    uploaded_file = st.file_uploader("Choose an audio or video file", type=SUPPORTED_FILE_TYPES)
    if uploaded_file:
        if st.button("Generate Quiz"):
            tmp_file_path = handle_uploaded_file(uploaded_file)
            transcript_text = transcribe_audio(tmp_file_path)
            os.remove(tmp_file_path)
            if "Error" not in transcript_text:
                summary = summarize_text(transcript_text)
                quiz_text = generate_quiz_questions(transcript_text)
                questions = parse_quiz_questions(quiz_text)

                if not questions:
                    st.error("No valid quiz questions could be generated.")
                else:
                    st.session_state.summary = summary
                    st.session_state.questions = questions
                    st.session_state.user_answers = {}
                    st.session_state.generated_quiz = True
            else:
                st.error(transcript_text)

if st.session_state.generated_quiz:
    st.write("## Summary")
    st.write(st.session_state.summary)

    st.write("## Quiz Questions")
    for i, question in enumerate(st.session_state.questions):
        st.write(f"### Question {i+1}")
        st.write(question['question'])
        st.session_state.user_answers[f"question_{i+1}"] = st.radio(
            label="",
            options=question['choices'],
            key=f"question_{i+1}"
        )

    if st.button("Submit Answers"):
        if "questions" in st.session_state and st.session_state.questions:
            with st.spinner('Processing your answers...'):
                feedback = check_answers(st.session_state.questions, st.session_state.user_answers)
                st.write("## Feedback")
                for i, item in enumerate(feedback):
                    with st.expander(f"Question {i+1} Feedback"):
                        st.write(f"### {item['question']}")
                        st.write(f"**Your answer:** {item['user_answer']}")
                        st.write(f"**Correct answer:** {item['correct_answer']}")
                        if item['status'] == "Incorrect":
                            st.write(f"**Explanation:** {item['explanation']}")
        else:
            st.write("Please generate the quiz first.")