import json
import streamlit as st
import google.generativeai as gai
import vertexai
import textwrap
import datetime
from database.bigquery_operations import *
from collections import defaultdict
from utils import *

if 'quiz_id' not in st.session_state:
    st.session_state.quiz_id = ''
    
def generate_quiz_questions(transcript, num_questions, difficulty):
    prompt = textwrap.dedent(f"""
        Please generate {num_questions} {difficulty} multiple-choice questions based on the following lecture transcript. Return the questions in a JSON format without any Markdown or additional formatting, such as code block annotations. Each question should include a unique question text, four choices, and the correct answer indicated. Make sure that the question is a string, choices is a list of strings, and the answer is also a string. Ensure the questions comprehensively cover the main points of the lecture.

        Here is the format I need:

        {{
          "questions": [
            {{
              "question": "What is the main concept of ...?",
              "choices": ["This", "That", "Maybe this", "Maybe that"],
              "answer": "That"
            }},
            ...
          ]
        }}

        Transcript: "{transcript}"
    """)

    
    gai.configure(api_key=st.secrets.gai_api_key)
    model = gai.GenerativeModel('gemini-pro')
    response = model.generate_content(prompt)

    while not validate_response(response.text):
        st.write('Invalid response, trying again...')
        response = model.generate_content(prompt)

    return response.text

def validate_response(raw_text):
    raw_text_cleaned = raw_text.strip("```").replace("json", "").strip()
    return json_validator(raw_text_cleaned, ["question", "choices", "answer"], outer_key="questions")

# Once validated, this will be used to parse the prompt response into a json.
def parse_response(raw_text, quiz_name):
    raw_text_cleaned = raw_text.strip("```").replace("json", "").strip()
    data = json.loads(raw_text_cleaned)
    questions = []
    for item in data['questions']:
        question = item['question']
        choices = item['choices']
        answer = item['answer']
        
        questions.append({
            "question": question,
            "choices": choices,
            "answer": answer
        })
    quiz_id = insert_quiz_data(questions, st.session_state.lecture_id, quiz_name)
    handle_quiz_selection(quiz_id)

def display_quiz(organized_quiz):
    quiz_completed = any(q['quiz_completed'] for q in organized_quiz.values())
    
    if quiz_completed:
        st.write("This quiz has already been completed!")
        return

    with st.form("Quiz Form"):
        user_responses = {}
        for question_id, question_info in organized_quiz.items():
            st.subheader(question_info['question_text'])
            options = [answer['answer_text'] for answer in question_info['answers']]
            answer_ids = [answer['answer_id'] for answer in question_info['answers']]
            user_choice = st.radio(
                "Select your answer:",
                options,
                key=question_id
            )
            selected_answer_id = answer_ids[options.index(user_choice)]
            user_responses[question_id] = selected_answer_id

        if st.form_submit_button("Submit Quiz"):
            with st.spinner('Processing quiz...'):
                st.write('Submitting your quiz...')
                submit_quiz_responses(organized_quiz, user_responses, st.session_state.quiz_id)
                st.write('Grading quiz and processing feedback...')
                feedback_text = generate_feedback(st.session_state.quiz_id)
                print(feedback_text)
                insert_feedback(st.session_state.quiz_id, feedback_text)
                st.experimental_rerun()

def display_feedback():
    correct_answers = 0
    quiz_results_text = []

    st.header("Quiz Summary and Feedback")
    st.subheader("Review Your Answers")

    quiz_results = get_quiz_results(st.session_state.quiz_id)
    for item in quiz_results:
        question = item['question_text']
        correct_answer = item['correct_answer_text']
        user_response = item['user_answer_text']

        st.text(question)
        if correct_answer == user_response:
            correct_answers += 1
            st.success(f"Correct! Your answer: {user_response}")
            result_text = "Correct"
        else:
            st.error(f"Wrong. Your answer: {user_response} | Correct answer: {correct_answer}")
            result_text = "Incorrect"

        # Append data for download in a text format
        quiz_results_text.append(f"Question: {question}")
        quiz_results_text.append(f"Your Answer: {user_response}")
        quiz_results_text.append(f"Correct Answer: {correct_answer}")
        quiz_results_text.append(f"Result: {result_text}")
        quiz_results_text.append("")  # Add a blank line for better readability

    total_questions = len(quiz_results)
    feedback = get_feedback(st.session_state.quiz_id)

    st.subheader("Feedback")
    st.write(feedback)  # Assuming feedback is a single string

    # Include summary and feedback in the download
    quiz_results_text.append(f"Total Correct Answers: {correct_answers} out of {total_questions}")
    quiz_results_text.append("\nDetailed Feedback:")
    quiz_results_text.append(feedback)

    # Join all feedback to form a plain text formatted string
    download_str = "\n".join(quiz_results_text)

    # Create a download button for plain text
    st.download_button(
        label="Download Quiz Results",
        data=download_str,
        file_name="quiz_results.txt",
        mime="text/plain"
    )

def generate_feedback(quiz_id):
    print("Generating feedback...")
    feedback_data = get_quiz_results(quiz_id)
    gai.configure(api_key=st.secrets.gai_api_key)
    prompt = textwrap.dedent(f"""
        Based on a detailed review of a user's quiz performance, please generate personalized honest feedback. The feedback should directly address the user with "You", highlighting both strengths and areas needing improvement. Provide specific, actionable steps for improvement that are directly related to the quiz questions and material. Aim to guide the user towards a deeper understanding of the subject matter. The feedback should be completely honest.

        Here is the data:

        {json.dumps(feedback_data)}
        """)
    
    model = gai.GenerativeModel('gemini-pro')
    response = model.generate_content(prompt)
    return response.text


def handle_quiz_selection(quiz_id):
    if st.session_state.get('quiz_id') != quiz_id:
        st.session_state.quiz_id = quiz_id

def organize_quiz_data(quiz_data):
    organized_data = {}
    
    for question_id, question_info in quiz_data.items():
        organized_data[question_id] = {
            "question_text": question_info['question_text'],
            "answers": question_info['answers'],
            "correct_answer_id": None,
            "selected_answer_id": None, 
            "quiz_completed": False  
        }

    return organized_data

def submit_quiz_responses(organized_quiz, user_responses, quiz_id):
    for question_id, selected_answer_id in user_responses.items():
        insert_user_response(question_id, selected_answer_id)

    mark_quiz_as_completed(quiz_id)

check_login(st)
st.title("Lecture Quiz Generator")

# Sidebar for navigation and quiz selection
with st.sidebar:
    st.page_link("pages/dashboard.py", label="Dashboard")
    st.page_link("pages/quiz.py", label="Quiz Page")
    st.page_link("pages/chat.py", label="Chat Page")
    if st.button("New Quiz"):
        st.session_state.quiz_id = ''

    # Finds quizzes and displays their name.
    if 'lecture_id' in st.session_state:
        quizzes = get_quizzes(st.session_state.lecture_id)
        for quiz in quizzes:
            if st.button(quiz['Quiz_name'], key=quiz['Quiz_id']):
                handle_quiz_selection(quiz['Quiz_id'])


# Main quiz interaction area
if 'lecture_id' not in st.session_state:
    st.write("No lecture selected.")
elif st.session_state.quiz_id != '':
    completed = check_quiz_completion(st.session_state.quiz_id)
    if not completed:
        quiz_data = get_quiz_data(st.session_state.quiz_id)
        print(f"quiz data: {quiz_data}")
        organized_quiz = organize_quiz_data(quiz_data)
        display_quiz(organized_quiz)
    else:
        display_feedback()
else:
    quiz_name = st.text_input("Name your quiz:")
    difficulty = st.selectbox("Select difficulty", ["Easy", "Medium", "Hard"])
    num_questions = st.slider("Select number of questions", 5, 20, 10)
    if st.button("Generate Quiz"):
        if quiz_name == '':
            st.error("Please input a quiz name!")
        else:
            with st.status('Generating quiz...'):
                st.write('Generating questions...')
                response = generate_quiz_questions(st.session_state.transcript, num_questions, difficulty)
                st.write('Parsing response...')
                parse_response(response, quiz_name)
                st.experimental_rerun()