import json
from typing import Dict

from datasets import load_dataset

from db.schema import Feedback, Response
from db.crud import ingest, read, save_feedback
import pandas as pd
import streamlit as st
from datetime import datetime
import os
from dotenv import load_dotenv
from views.intro_screen import welcome_screen
from views.questions_screen import questions_screen, survey_completed
from views.continue_survey import continue_survey_screen
from css.layout import custom_css

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


# class SurveyState:
#     """Class to handle survey state management"""
#
#     def __init__(self):
#         pass
#
#     def save_state(self, username: str, current_state: Dict) -> None:
#         """Save current state to Firebase"""
#         try:
#             """Handles feedback submission to the database."""
#             feedback = Feedback(
#                 id=st.session_state.current_index + 1,
#                 user_id=st.session_state.username,
#                 time_stamp=datetime.now().isoformat(),
#                 responses=st.session_state.responses,
#             )
#             save_feedback(feedback)
#             st.success("Your progress has been saved! You can continue later.")
#             st.session_state.completed = True
#             st.rerun()
#         except Exception as e:
#             st.error(f"An error occurred while submitting feedback: {e}")
#

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 "survey_state" not in st.session_state:
    #     st.session_state.survey_state = SurveyState()


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("Start New 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 load_data():
    try:
        data = pd.read_csv("data/gemini_results_subset.csv")[:5]
        return data
    except Exception as e:
        repo_name = os.getenv("DATA_REPO")
        data_files = os.getenv("LLAMA_DATA_FILES")
        HF_TOKEN = os.getenv("HF_TOKEN")
        dataset = load_dataset(repo_name, token=True, data_files=data_files)
        dataset.set_format(type='pandas')  ## converting it into pandas
        df = dataset["train"][:]
        return df[:5]


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")
                    survey_completed()
                # 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()