from db.schema import 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()
if "VALIDATION_CODE" in os.environ:
    VALIDATION_CODE = os.getenv("VALIDATION_CODE")
if "DATA_REPO" in os.environ:
    REPO_NAME = os.getenv("DATA_REPO")
else:
    print("DATA_REPO not found in environment variables.")
if "LLAMA_DATA_FILES" in os.environ:
    DATA_FILES = os.getenv("LLAMA_DATA_FILES")
else:
    print("LLAMA_DATA_FILES not found in environment variables.")



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


def render_query_ratings(query_label, config, query_key, current_index, has_persona_alignment=False):
    """Helper function to render ratings for a given query."""
    st.markdown(f"### {query_label}")
    st.write(config[query_key])
    # columns = st.columns(3 if has_persona_alignment else 2)
    columns = st.columns(3)
    options = [0, 1, 2, 3, 4]

    persona_alignment_rating = None
    if has_persona_alignment:
        with columns[0]:
            persona_alignment_rating = st.radio(
                "Persona Alignment:", options=[0, 1, 2, 3, 4],
                format_func=lambda x: ["N/A", "Not Aligned", "Partially Aligned", "Aligned", "Unclear"][x],
                key=f"rating_{query_key}_persona_alignment_{current_index}"
            )

    with columns[1]:
        relevance_rating = st.radio("Relevance:",
                                    options, key=f"rating_{query_key}_relevance_{current_index}",
                                    format_func=lambda x:
                                    ["N/A", "Not Relevant", "Somewhat Relevant", "Relevant", "Unclear"][x], )
    with columns[2]:
        clarity_rating = st.radio("Clarity:",
                                  options=[0, 1, 2, 3],
                                  key=f"rating_{query_key}_clarity_{current_index}",
                                  format_func=lambda x: ["N/A", "Not Clear", "Somewhat Clear", "Very Clear"][x])

    return {
        "clarity": clarity_rating,
        "relevance": relevance_rating,
        "persona_alignment": persona_alignment_rating if has_persona_alignment else None
    }


def questions_screen(data):
    """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.text_area("", config['context'], height=300, disabled=False)

        # Render queries and collect ratings
        query_v_ratings = render_query_ratings("Query_v", config, "query_v", current_index)
        query_p0_ratings = render_query_ratings("Query_p0",
                                                config, "query_p0", current_index, has_persona_alignment=True)
        query_p1_ratings = render_query_ratings("Query_p1",
                                                config, "query_p1",
                                                current_index, has_persona_alignment=True)

        # Additional comments
        comment = st.text_area("Additional Comments (Optional):")
        if "persona_alignment" in query_v_ratings:
            query_v_ratings.pop("persona_alignment")
        # Collecting the response data
        response = Response(
            config_id=config["config_id"],
            query_v=query_v_ratings,
            query_p0=query_p0_ratings,
            query_p1=query_p1_ratings,
            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!")