from utils.loaders import load_data
from db.crud import read
import streamlit as st
import os
from dotenv import load_dotenv
from views.intro_screen import welcome_screen
from views.questions_screen import questions_screen, display_completion_message
from views.continue_survey import continue_survey_screen
from css.layout import custom_css
st.set_page_config(layout="wide")


load_dotenv()
VALIDATION_CODE = os.getenv("VALIDATION_CODE")


def initialization():
    """Initialize session state variables."""
    if "current_index" not in st.session_state:
        st.session_state.current_index = 0
    if "username" not in st.session_state:
        st.session_state.username = None
    if "responses" not in st.session_state:
        st.session_state.responses = []
    if "completed" not in st.session_state:
        st.session_state.completed = False
    if "show_questions" not in st.session_state:
        st.session_state.show_questions = False
    if "survey_continued" not in st.session_state:
        st.session_state.survey_continued = None
    if "start_new_survey" not in st.session_state:
        st.session_state.start_new_survey = False
    if 'ratings' not in st.session_state:
        st.session_state.ratings = {}
    if 'previous_ratings' not in st.session_state:
        st.session_state.previous_ratings = {}


def exit_screen():
    """Display exit screen"""
    st.markdown("""
        <div class='exit-container'>
            <h1>Thank you for participating!</h1>
            <p>Your responses have been saved successfully.</p>
            <p>You can safely close this window or start a new survey.</p>
        </div>
    """, unsafe_allow_html=True)

    # if st.button("Resume Survey"):
    #     reset_survey()
    #     st.rerun()


def reset_survey():
    """Reset the survey state to start over."""
    st.session_state.responses = []
    st.session_state.completed = True
    st.session_state.start_new_survey = True


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

    if st.session_state.completed and not st.session_state.start_new_survey:
        exit_screen()
        return

    if st.session_state.username is None:
        welcome_screen()
    else:
        # Check if user progress exists in Firebase
        saved_state = read(st.session_state.username)
        if saved_state:
            # If there's saved progress and the survey has not been continued, show continue screen
            if "survey_continued" not in st.session_state or not st.session_state.survey_continued:
                continue_survey_screen(data)
            else:
                if st.session_state.current_index >= len(data):
                    # If all questions have been answered, show the exit screen
                    print("survey completed")
                    display_completion_message()
                # Otherwise, show questions from where they left off
                questions_screen(data)
        else:
            # If no saved progress (new user), start with the questions screen
            questions_screen(data)


if __name__ == "__main__":
    ui()