import streamlit as st
import google.generativeai as genai
import openai
import os
import re  # For improved parsing
from streamlit_extras.metric_cards import style_metric_cards

st.set_page_config(initial_sidebar_state="collapsed")

# --- Login Page ---
def login_page():
    # st.title("API Selection")


    api_options = ["Gemini API", "OpenAI API"]
    selected_api = st.selectbox("### Select API:", api_options)
    api_key = st.text_input(f"Enter your {selected_api} Key:", type="password")

    if selected_api == "Gemini API":
        model_options = ["gemini-pro"] # Add more Gemini models if available
        selected_model = st.selectbox("### Select Gemini Model:", model_options, index=0) # Default to first model
    elif selected_api == "OpenAI API":
        model_options = ["gpt-3.5-turbo", "gpt-4"] # Add more OpenAI models as needed
        selected_model = st.selectbox("### Select OpenAI Model:", model_options, index=0) # Default to first model
    else:
        selected_model = None # Should not happen, but for safety

    if st.button("Login"):
        if not api_key:
            st.error("Please enter your API key.")
        else:
            st.session_state.api_key = api_key
            st.session_state.selected_api = selected_api
            if selected_api == "Gemini API":
                st.session_state.gemini_model = selected_model
                st.session_state.openai_model = None # Clear OpenAI model if Gemini is selected
            elif selected_api == "OpenAI API":
                st.session_state.openai_model = selected_model
                st.session_state.gemini_model = None # Clear Gemini model if OpenAI is selected
            st.session_state.logged_in = True
            st.success(f"Logged in with {selected_api} using model: {selected_model}!")
            st.rerun() # Rerun to show the quiz app

# --- Quiz Application ---
def quiz_app():
    st.title("QuizLM 🧠")

    # Initialize session state for quiz data and progress
    if 'quiz_data' not in st.session_state:
        st.session_state.quiz_data = None
    if 'current_question_index' not in st.session_state:
        st.session_state.current_question_index = 0
    if 'user_answers' not in st.session_state:
        st.session_state.user_answers = []
    if 'quiz_completed' not in st.session_state:
        st.session_state.quiz_completed = False
    if 'score' not in st.session_state:
        st.session_state.score = 0
    if 'quiz_started' not in st.session_state:
        st.session_state.quiz_started = False
    if 'quiz_topic' not in st.session_state:
        st.session_state.quiz_topic = "" # Initialize quiz topic
    if 'topic_from_sidebar' not in st.session_state:
        st.session_state.topic_from_sidebar = False # Flag to track if topic is from sidebar


    # API Configuration based on selection - Moved inside quiz_app after login
    if st.session_state.selected_api == "Gemini API":
        try:
            genai.configure(api_key=st.session_state.api_key)
            model = genai.GenerativeModel(st.session_state.gemini_model)
        except Exception as e:
            st.error(f"Error configuring Gemini API: {e}")
            st.error("Please check your API key and network connection.")
            st.session_state.logged_in = False # Force logout if API config fails
            del st.session_state.api_key
            del st.session_state.selected_api
            st.rerun() # Go back to login page
            return # Exit quiz_app
    elif st.session_state.selected_api == "OpenAI API":
        try:
            openai.api_key = st.session_state.api_key
            model_name = st.session_state.openai_model # Model name already selected in login
        except Exception as e:
            st.error(f"Error configuring OpenAI API: {e}")
            st.error("Please check your API key and network connection.")
            st.session_state.logged_in = False # Force logout if API config fails
            del st.session_state.api_key
            del st.session_state.selected_api
            st.rerun() # Go back to login page
            return # Exit quiz_app


    def parse_quiz_content(quiz_content):
        """Parses the quiz content and answer key into a structured format, handling title."""
        questions = []
        answer_key_dict = {}

        quiz_parts = quiz_content.split("Answer Key:")
        if len(quiz_parts) != 2:
            st.warning("Could not reliably separate quiz questions and answer key. Parsing might be imperfect.")
            return None, None

        question_section = quiz_parts[0].strip()
        answer_key_section = quiz_parts[1].strip()

        question_lines = question_section.strip().split('\n')
        current_question = None
        for line in question_lines:
            line = line.strip()
            if not line:
                continue # Skip empty lines

            if re.match(r'^\d+\.\s', line): # Start of a new question
                if current_question: # Save previous question if exists
                    questions.append(current_question)
                current_question = {'options': {}} # Start new question
                current_question['question'] = re.sub(r'\*\*|\*', '', line).strip() # Question text, remove bold
            elif re.match(r'^[A-D]\)\s', line) or re.match(r'^[A-D]\.\s', line): # Option line (handle both ')' and '.' after letter)
                if current_question:
                    option_letter = line[0]
                    option_text = line[2:].strip() # Remove "A) " or "A. "
                    current_question['options'][option_letter] = option_text
            elif re.match(r'^\*\*.*?\*\*$', line) or re.match(r'^\*.*?\*$', line) : # Check for title again, just in case after split. Remove bold title if found in lines.
                pass # Ignore title lines within question section if any.

        if current_question: # Append the last question
            questions.append(current_question)


        # --- Answer Key Parsing (No change needed) ---
        answer_lines = answer_key_section.strip().split('\n')
        for line in answer_lines:
            line = line.strip()
            match = re.match(r'(\d+)\.\s*([A-D])', line)
            if match:
                question_num = int(match.group(1)) - 1
                correct_answer = match.group(2)
                answer_key_dict[question_num] = correct_answer

        # Basic validation
        if not questions or not answer_key_dict:
            st.error("Error parsing quiz content. Please try again or check the generated format.")
            return None, None
        if len(questions) != len(answer_key_dict):
            st.warning(f"Number of questions parsed ({len(questions)}) does not match number of answers in answer key ({len(answer_key_dict)}). Parsing might be incomplete.")


        # Combine parsed questions and answer key into quiz_data
        quiz_data_list = []
        for i, q_data in enumerate(questions):
            correct_answer = answer_key_dict.get(i)
            if correct_answer:
                quiz_data_list.append({
                    'question': q_data['question'],
                    'options': q_data['options'],
                    'correct_answer': correct_answer
                })
            else:
                st.warning(f"Could not find correct answer for question {i+1} in the answer key.")
                return None, None

        return quiz_data_list, answer_key_dict


    def display_question():
        """Displays the current question and options for interactive quiz."""
        if st.session_state.quiz_data is None:
            st.error("Quiz data not loaded yet. Generate a quiz first.")
            return

        if st.session_state.current_question_index < len(st.session_state.quiz_data):
            question_data = st.session_state.quiz_data[st.session_state.current_question_index]
            question_number = st.session_state.current_question_index + 1
            st.subheader(f"{question_data['question']}")

            options_list = [f"{key}. {value}" for key, value in question_data['options'].items()]
            user_choice = st.radio("Choose an answer:", options_list, key=f"q_{question_number}") # Unique key for radio buttons

            if st.button("Submit Answer", key=f"submit_q_{question_number}"): # Unique key for submit button
                selected_option_letter = user_choice.split('.')[0] # Extract A, B, C, or D
                st.session_state.user_answers.append(selected_option_letter)

                if st.session_state.current_question_index < len(st.session_state.quiz_data) - 1:
                    st.session_state.current_question_index += 1
                else:
                    st.session_state.quiz_completed = True # Mark quiz as completed after last question
                st.rerun()


    def display_results():
        """Displays the quiz results after completion."""
        if st.session_state.quiz_completed:
            st.balloons()
            st.markdown("### Quiz Results")
            correct_count = 0
            for i in range(len(st.session_state.quiz_data)):
                user_answer = st.session_state.user_answers[i] if i < len(st.session_state.user_answers) else None # Handle if user didn't answer all
                correct_answer = st.session_state.quiz_data[i]['correct_answer']
                if user_answer == correct_answer:
                    correct_count += 1
                    result_text = f"**Question {i+1}:** ✅ Correct! (Your answer: {user_answer}, Correct answer: {correct_answer})"
                    color = "green"
                else:
                    result_text = f"**Question {i+1}:** ❌ Incorrect. (Your answer: {user_answer if user_answer else 'Not answered'}, Correct answer: {correct_answer})"
                    color = "red"
                st.markdown(f"<font color='{color}'>{result_text}</font>", unsafe_allow_html=True)

            st.divider()

            percentage_correct = (correct_count / len(st.session_state.quiz_data)) * 100

            if percentage_correct > 75:
                border_color = "#9AD8E1"
                delta_text_color = "normal"
            else:
                border_color = "#E098B7"
                delta_text_color = "inverse"


            st.write(
                """
                <style>
                [data-testid="stMetricDelta"] svg {
                    display: none;
                }
                </style>
                """,
                unsafe_allow_html=True,
            )

            st.metric(label=f"**Final Score**", value=f"{correct_count} out of {len(st.session_state.quiz_data)}", delta=f"{percentage_correct:.0f}%", delta_color=delta_text_color)
            style_metric_cards(border_left_color=border_color, background_color="DEFAULT_BACKGROUND_COLOR")

            st.session_state.score = correct_count # Store score in session state


    # Sidebar topic input - Always visible
    sidebar_topic = st.sidebar.text_input("New Quiz Topic", value=st.session_state.get('quiz_topic', ''))
    if st.sidebar.button("Generate New Quiz"): # Button in sidebar to generate quiz, always resets
        st.session_state.quiz_started = True
        st.session_state.quiz_topic = sidebar_topic
        st.session_state.topic_from_sidebar = True # Flag to indicate topic from sidebar
        # Reset quiz state for new quiz
        st.session_state.quiz_data = None
        st.session_state.current_question_index = 0
        st.session_state.user_answers = []
        st.session_state.quiz_completed = False
        st.session_state.score = 0


    # Initial topic input in main area - Shown only if quiz not started
    if not st.session_state.quiz_started: # Modified condition here
        initial_topic = st.text_input("Enter the topic for your quiz:", value=st.session_state.get('quiz_topic', '') ) # Removed condition causing issue
        if st.button("Generate Quiz"): # Button to generate quiz from initial input
            if initial_topic:
                st.session_state.quiz_started = True
                st.session_state.quiz_topic = initial_topic
                st.session_state.topic_from_sidebar = False # Reset sidebar topic flag
                # Reset quiz state - although should be already reset if coming from sidebar, or first time
                st.session_state.quiz_data = None
                st.session_state.current_question_index = 0
                st.session_state.user_answers = []
                st.session_state.quiz_completed = False
                st.session_state.score = 0
            else:
                st.warning("Please enter a topic to generate a quiz.")


    # Quiz Generation Logic - Moved outside input blocks, runs after button presses
    topic_to_generate = None
    if st.session_state.quiz_started and st.session_state.quiz_data is None: # Generate only if quiz started and no data yet
        topic_to_generate = st.session_state.quiz_topic

    if topic_to_generate:
        with st.spinner(f"Generating quiz on '{topic_to_generate}'..."):
            try:
                prompt = f"""
                    Generate a multiple-choice quiz on the topic of "{topic_to_generate}".
                    The quiz should have 5 questions.

                    **Formatting Instructions:**

                    1. **Question Numbering:**  Start each question with a number followed by a period (e.g., "1.").
                    2. **Question Text:**  Immediately after the question number, write the question text.
                    3. **Options:**
                        * List four options for each question, labeled A, B, C, and D.
                        * Use the format: "A) Option text", "B) Option text", etc. (letter followed by a parenthesis, then a space, then the option text).
                        * Place each option on a new line directly below the question.
                    4. **Answer Key:**
                        * After all questions, create a separate section titled "**Answer Key:**".
                        * In the Answer Key, list the correct answer for each question in the format: "1. Correct option letter", "2. Correct option letter", etc. (question number, period, space, correct option letter - A, B, C, or D).

                    **Example of Desired Format:**

                    **Quiz Title (Optional, but good to have)**

                    1. Question 1 text?
                    A) Option A
                    B) Option B
                    C) Option C
                    D) Option D

                    2. Question 2 text?
                    A) Option A
                    B) Option B
                    C) Option C
                    D) Option D

                    ... (and so on for 5 questions) ...

                    **Answer Key:**
                    1. C
                    2. A
                    3. B
                    4. D
                    5. A

                    Please generate the quiz in this exact format.
                    """
                if st.session_state.selected_api == "Gemini API":
                    response = model.generate_content(prompt)
                    quiz_content = response.text
                elif st.session_state.selected_api == "OpenAI API":
                    response = openai.ChatCompletion.create(
                        model=st.session_state.openai_model,
                        messages=[{"role": "user", "content": prompt}]
                    )
                    quiz_content = response.choices[0].message.content


                parsed_quiz_data, answer_key = parse_quiz_content(quiz_content)


                if quiz_content: # Check if quiz_content was generated successfully (outer if)
                    if parsed_quiz_data: # Check if parsing was successful (inner if)
                        st.session_state.quiz_data = parsed_quiz_data
                        st.success(f"Quiz on '{topic_to_generate}' generated successfully! Let's begin.")
                    else: # else associated with inner if parsed_quiz_data
                        st.error("Failed to parse quiz content. Please try generating again.")
                        st.session_state.quiz_data = None
                else: # else associated with outer if quiz_content
                    st.error("Failed to generate quiz content. Please try again or check your API key.")

            except Exception as e:
                st.error(f"An error occurred: {e}")
                st.error("Please check your API key and network connection. If the problem persists, try a different topic or try again later.")
        st.rerun() # Rerun to update the quiz display based on new topic/generation


    # Quiz Display Logic
    if st.session_state.quiz_data:
        if not st.session_state.quiz_completed:
            display_question()
        else:
            display_results()
    elif not st.session_state.quiz_started and not st.session_state.quiz_data and not st.session_state.quiz_completed:
        st.info("Enter a topic in the main area and click 'Generate Quiz' to start, or use the sidebar to start a new quiz anytime.")


# --- Main App Flow Control ---
if 'logged_in' not in st.session_state:
    st.session_state.logged_in = False

if not st.session_state.logged_in:
    login_page()
else:
    quiz_app()