from db.schema import Feedback, Response
from db.crud import ingest
import pandas as pd
import streamlit as st
from datetime import datetime


# Load Q&A data
@st.cache_data
def load_data():
    return pd.read_csv("dummy_qa_data.csv")[:3]


# Function to validate Prolific ID (customize this based on your format requirements)
def validate_prolific_id(prolific_id: str) -> bool:
    return bool(prolific_id.strip())  # Example: Check non-empty

    # Function to reset the survey


def reset_survey():
    st.session_state.prolific_id = None
    st.session_state.current_index = 0
    st.session_state.responses = []
    st.session_state.completed = False


def initialization():
    """Initialize session state variables."""
    if "current_index" not in st.session_state:
        st.session_state.current_index = 0
    if "prolific_id" not in st.session_state:
        st.session_state.prolific_id = None
    if "responses" not in st.session_state:
        st.session_state.responses = []
    if "completed" not in st.session_state:
        st.session_state.completed = False


def prolific_id_screen():
    """Display the Prolific ID entry screen."""
    st.title("Welcome to the Feedback Survey")
    prolific_id_input = st.text_input("Enter your Prolific ID and press ENTER:")

    if prolific_id_input:
        if validate_prolific_id(prolific_id_input):
            st.session_state.prolific_id = prolific_id_input
            st.success("Prolific ID recorded! Now, proceed with the questions.")
            st.experimental_rerun()
        else:
            st.warning("Invalid Prolific ID. Please check and try again.")


def questions_screen(data):
    """Display the questions screen."""
    current_index = st.session_state.current_index

    # Display question and answer
    question = data.loc[current_index, "Question"]
    generated_answer = data.loc[current_index, "Generated Answer"]

    st.subheader(f"Question {current_index + 1}")
    st.markdown(f"***{question}***")
    st.subheader("Generated Answer:")
    st.write(generated_answer)

    # Prefill responses if they exist
    previous_rating = (
        st.session_state.responses[current_index].ans
        if len(st.session_state.responses) > current_index
        else 0
    )
    previous_feedback = (
        st.session_state.responses[current_index].feedback_text
        if len(st.session_state.responses) > current_index
        else ""
    )

    # Rating slider
    rating = st.slider("Rate the answer (1 = Worst, 5 = Best)", 1, 5, value=previous_rating,
                       key=f"rating_{current_index}")

    # Free-text feedback
    feedback_text = st.text_area("Any additional feedback?", value=previous_feedback, key=f"feedback_{current_index}")

    # Store or update the response
    response = Response(q_id=f"q{current_index + 1}", ans=rating, feedback_text=feedback_text)
    if len(st.session_state.responses) > current_index:
        st.session_state.responses[current_index] = response
    else:
        st.session_state.responses.append(response)

    # Navigation buttons
    navigation_buttons(data, rating, feedback_text)


def navigation_buttons(data, rating, feedback_text):
    """Display navigation buttons."""
    current_index = st.session_state.current_index

    col1, col2, col3 = st.columns([1, 1, 2])

    with col1:  # Back button
        if st.button("Back") and current_index > 0:
            st.session_state.current_index -= 1
            st.experimental_rerun()

    with col2:  # Next button
        if st.button("Next"):
            if not rating:
                st.warning("Please provide a rating before proceeding.")
            else:
                if current_index < len(data) - 1:
                    st.session_state.current_index += 1
                    st.experimental_rerun()
                else:
                    feedback = Feedback(
                        id=current_index + 1,
                        user_id=st.session_state.prolific_id,
                        time_stamp=datetime.now(),
                        responses=st.session_state.responses,
                    )
                    try:
                        ingest(feedback)
                        st.session_state.completed = True
                        st.experimental_rerun()
                    except Exception as e:
                        st.error(f"An error occurred while submitting feedback: {e}")


def reset_survey():
    """Reset the survey state to start over."""
    st.session_state.prolific_id = None
    st.session_state.current_index = 0
    st.session_state.responses = []
    st.session_state.completed = False


def ui():
    """Main function to control the survey flow."""
    data = load_data()
    initialization()

    if st.session_state.completed:
        st.title("Survey Completed")
        st.success("You have completed all questions and your answers have been recorded.")
        if st.button("Retake Survey"):
            reset_survey()
            st.experimental_rerun()
        return

    if st.session_state.prolific_id is None:
        prolific_id_screen()
    else:
        questions_screen(data)


if __name__ == "__main__":
    ui()