import json
from typing import Dict

from db.schema import Feedback, Response
import streamlit as st
from datetime import datetime
import os
from dotenv import load_dotenv
from views.nav_buttons import navigation_buttons

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


def survey_completed():
    """Display the survey completion message."""
    st.markdown("""
            <div class='exit-container'>
                <h1>You have already completed the survey! 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)
    st.session_state.show_questions = False
    st.session_state.completed = True
    st.session_state.start_new_survey = True
    # st.rerun()


def questions_screen(data):
    # TODO: refactor to avoid code duplication
    """Display the questions screen with split layout"""
    current_index = st.session_state.current_index
    try:
        config = data.iloc[current_index]

        # Progress bar
        progress = (current_index + 1) / len(data)
        st.progress(progress)
        st.write(f"Question {current_index + 1} of {len(data)}")
        st.subheader(f"Config ID: {config['config_id']}")

        # Context information
        st.markdown("### Context Information")

        with st.expander("Persona", expanded=True):
            st.write(config['persona'])

        with st.expander("Filters & Cities", expanded=True):
            st.write("**Filters:**", config['filters'])
            st.write("**Cities:**", config['city'])

        with st.expander("Full Context", expanded=False):
            st.write(config['context'])

        # Split layout for questions and ratings
        col11, col12, col13, col14 = st.columns([1, 1, 1, 1])  # Sub-columns for query ratings
        options = [0, 1, 2, 3, 4, 5]

        # Query_v and its ratings
        st.markdown("### Query_v")
        st.write(config['query_v'])
        col_v_1, col_v_2, col_v_3 = st.columns(3)
        with col_v_1:
            clarity_rating = st.radio("Clarity:", options, key=f"rating_v_clarity_{current_index}")

        with col_v_2:
            relevance_rating = st.radio("Relevance:", options, key=f"rating_v_relevance_{current_index}")

        with col_v_3:
            coverage_rating = st.radio("Coverage:", options, key=f"rating_v_coverage_{current_index}")

        query_v_ratings = {
            "clarity": clarity_rating,
            "relevance": relevance_rating,
            "coverage": coverage_rating,
        }

        # Query_p0 and its ratings
        st.markdown("### Query_p0")
        st.write(config['query_p0'])
        col_p0_1, col_p0_2, col_p0_3, col_p0_4 = st.columns(4)

        with col_p0_1:
            clarity_rating = st.radio("Clarity:", options, key=f"rating_p0_clarity_{current_index}")

        with col_p0_2:
            relevance_rating = st.radio("Relevance:", options, key=f"rating_p0_relevance_{current_index}")

        with col_p0_3:
            coverage_rating = st.radio("Coverage:", options, key=f"rating_p0_coverage_{current_index}")

        with col_p0_4:
            persona_alignment_rating = st.radio(
                "Persona Alignment:", options=[0, 1, 2, 3, 4],  # These are the values
                format_func=lambda x: ["N/A", "Not Aligned", "Partially Aligned", "Aligned", "Unclear"][x],
                key=f"rating_p0_persona_alignment_{current_index}"
            )

        # Collecting the ratings for query_p0
        query_p0_ratings = {
            "clarity": clarity_rating,
            "relevance": relevance_rating,
            "coverage": coverage_rating,
            "persona_alignment": persona_alignment_rating
        }

        # Query_p1 and its ratings
        st.markdown("### Query_p1")
        st.write(config['query_p1'])
        # Split the layout into 4 columns for query_p1 ratings
        col_p1_1, col_p1_2, col_p1_3, col_p1_4 = st.columns(4)

        with col_p1_1:
            clarity_rating_p1 = st.radio("Clarity:", options, key=f"rating_p1_clarity_{current_index}")

        with col_p1_2:
            relevance_rating_p1 = st.radio("Relevance:", options, key=f"rating_p1_relevance_{current_index}")

        with col_p1_3:
            coverage_rating_p1 = st.radio("Coverage:", options, key=f"rating_p1_coverage_{current_index}")

        with col_p1_4:
            persona_alignment_rating_p1 = st.radio(
                "Persona Alignment:", options=[0, 1, 2, 3, 4],  # These are the values
                format_func=lambda x: ["N/A", "Not Aligned", "Partially Aligned", "Aligned", "Unclear"][x],
                key=f"rating_p1_persona_alignment_{current_index}"
            )

        # Collecting the ratings for query_p1
        query_p1_ratings = {
            "clarity": clarity_rating_p1,
            "relevance": relevance_rating_p1,
            "coverage": coverage_rating_p1,
            "persona_alignment": persona_alignment_rating_p1
        }

        # Additional comments
        comment = st.text_area("Additional Comments (Optional):")

        # Collecting the response data
        response = Response(
            config_id=config["config_id"],
            query_v=query_v_ratings,  # Use the ratings dictionary for query_v
            query_p0=query_p0_ratings,  # Use the ratings dictionary for query_p0
            query_p1=query_p1_ratings,  # Use the ratings dictionary for query_p1
            comment=comment,
            timestamp=datetime.now().isoformat()
        )
        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, query_v_ratings["clarity"], query_p0_ratings["clarity"], query_p1_ratings["clarity"])
    except IndexError:
        print("Survey completed!")
        # st.stop()