from db.schema import Feedback, Response
from db.crud import save_feedback, read
import streamlit as st
from datetime import datetime
import os
from dotenv import load_dotenv

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


def submit_feedback(current_index):
    """Handles feedback submission to the database."""
    feedback = Feedback(
        id=current_index + 1,
        user_id=st.session_state.username,
        time_stamp=datetime.now().isoformat(),
        responses=st.session_state.responses,
    )
    try:
        save_feedback(feedback)
        st.session_state.completed = True
        st.rerun()
    except Exception as e:
        st.error(f"An error occurred while submitting feedback: {e}")


def flatten_ratings(response):
    all_ratings = []
    for model_ratings in response.model_ratings.values():
        all_ratings.extend(model_ratings.query_v_ratings.values())
        all_ratings.extend(model_ratings.query_p0_ratings.values())
        all_ratings.extend(model_ratings.query_p1_ratings.values())
    return all_ratings


def navigation_buttons(data, response: Response):
    """Display navigation buttons."""

    current_index = st.session_state.current_index

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

    with col1:  # Back button #TODO fix: only gets ratings for the session, not from previous session
        if st.button("Back", disabled=st.session_state.current_index == 0):
            if current_index > 0:
                st.session_state.previous_ratings[
                    data.iloc[st.session_state.current_index]['config_id']] = response.model_ratings
                st.session_state.current_index -= 1
                st.rerun()
            else:
                st.warning("You are at the beginning of the survey, can't go back.")
            # st.rerun()

    with col2:  # Next button
        if st.button("Next", disabled=current_index == len(data) - 1):
            all_ratings = flatten_ratings(response)
            if any(rating == 0 for rating in all_ratings):
                st.warning("Please provide all ratings before proceeding.")
            elif current_index < len(data) - 1:
                # Store the ratings with proper structure
                config_id = data.iloc[current_index]['config_id']
                if isinstance(response.model_ratings, dict):
                    st.session_state.previous_ratings[config_id] = {
                        model: {
                            'query_v_ratings': ratings.query_v_ratings,
                            'query_p0_ratings': ratings.query_p0_ratings,
                            'query_p1_ratings': ratings.query_p1_ratings
                        }
                        for model, ratings in response.model_ratings.items()
                    }
                else:
                    st.session_state.previous_ratings[config_id] = {
                        model: {
                            'query_v_ratings': getattr(ratings, 'query_v_ratings', {}),
                            'query_p0_ratings': getattr(ratings, 'query_p0_ratings', {}),
                            'query_p1_ratings': getattr(ratings, 'query_p1_ratings', {})
                        }
                        for model, ratings in response.model_ratings.items()
                    }
                st.session_state.current_index += 1
                st.rerun()

    with col3:  # Submit button
        if st.button("Submit & Resume Later"):
            all_ratings = flatten_ratings(response)
            if any(rating == 0 for rating in all_ratings):
                st.warning("Please provide all ratings before submitting.")
            else:
                submit_feedback(current_index)
    with col4:  # Save & Resume Later button
        if st.button("Exit & Resume Later"):
            submit_feedback(current_index)