import streamlit as st
import random
import time
import hmac
import os

st.header(" Scientific Claim Verification ")
st.caption("Team UMBC-SBU-UT")

def check_password():
    """Returns `True` if the user had a correct password."""

    def login_form():
        """Form with widgets to collect user information"""
        with st.form("Credentials"):
            st.text_input("Username", key="username")
            st.text_input("Password", type="password", key="password")
            st.form_submit_button("Log in", on_click=password_entered)

    def password_entered():
        """Checks whether a password entered by the user is correct."""

        stored_password = os.getenv(st.session_state["username"])
        if stored_password == st.session_state["password"]:
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # Remove credentials from session
            del st.session_state["username"]
            return

          # If authentication fails
        st.session_state["password_correct"] = False

    # Return True if the username + password is validated.
    if st.session_state.get("password_correct", False):
        return True

    # Show inputs for username + password.
    login_form()
    if "password_correct" in st.session_state:
        st.error("😕 User not known or password incorrect")
    return False


def select_models():
    """Returns only when a valid option is selected from both dropdowns."""

    #placeholders
    retriever_options = ["Choose one...", "Simple", "Trained", "No Retriever"]
    reasoner_options = ["Choose one...", "Claude Sonnet", "GPT-4o", "o3-mini"]

    #selectboxes

    retriever = st.selectbox(
        "Select the Retriever Model",
        retriever_options,
        key="retriever"
    )

    reasoner = st.selectbox(
        "Select the Reasoner Model",
        reasoner_options,
        key="reasoner"
    )

    #next button

    if st.button("Next"):
        # Check that both selections are not the placeholder.
        if retriever == "Choose one..." or reasoner == "Choose one...":
            st.info("Please select both a retriever and a reasoner.")
            return None, None
        else:
            # Store the valid selections in session state
            st.session_state["selected_models"] = (retriever, reasoner)
            return retriever, reasoner
    else:
        st.info("Click 'Next' once you have made your selections.")
        return None, None

if not check_password():
    st.stop()


if "selected_models" not in st.session_state:
    selected_retriever, selected_reasoner = select_models()
    # If valid selections are returned, store them and reset the change flag.
    if selected_retriever is not None and selected_reasoner is not None:
        st.session_state.selected_models = (selected_retriever, selected_reasoner)
        st.rerun()
    else:
        st.stop()  # Halt further execution until valid selections are made.
else:
    selected_retriever, selected_reasoner = st.session_state.selected_models

#START OF AGENTIC DEMO

column1, column2 = st.columns(2)
column1.caption(f"Retriever Selected: {selected_retriever}")
column2.caption(f"Reasoner Selected: {selected_reasoner}")


if st.button("Change Selection", key="change_selection_btn"):
    st.session_state.pop("selected_models", None)
    st.session_state.pop("retriever", None)
    st.session_state.pop("reasoner", None)
    st.rerun()

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "Let's start verifying the claims here! 👇"}]

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

def retriever(query: str):
    """Simulate a 'retriever' step, searching for relevant information."""
    with st.chat_message("assistant"):
        placeholder = st.empty()
        text=""
        
        if selected_retriever == "Simple":
            message = "Using the simple retriever to search for documents related to your query..."
        elif selected_retriever == "Trained":
            message = "Using the trained retriever to fetch detailed documents relevant to your query..."
        else:
            message = "No retriever selected. Skipping document retrieval."
       

        for chunk in message.split():
            text += chunk + " "
            time.sleep(0.05)
            # Add a blinking cursor to simulate typing
            placeholder.markdown(text + "▌")
        placeholder.markdown(text)
    # You could return retrieved info here.
    return message

def reasoner(info: list[str]):
    """Simulate a 'reasoner' step, thinking about how to answer."""
    with st.chat_message("assistant"):
        placeholder = st.empty()
        text=""

        if selected_reasoner == "Claude Sonnet":
            message = "Using Claude Sonnet to reason and verify the claim..."
        elif selected_reasoner == "GPT-4o":
            message = "Using GPT-4o to analyze and verify the claim in detail..."
        else:
            message = "Using o3-mini to quickly analyze the claim..."

        for chunk in message.split():
            text += chunk + " "
            time.sleep(0.05)
            # Add a blinking cursor to simulate typing
            placeholder.markdown(text + "▌")
        placeholder.markdown(text)
    # You could return reasoning info here.
    return message

# Accept user input
if prompt := st.chat_input("Type here"):
    # Add user message to chat history
    prompt= prompt + "  \n"+ "  \n"+ f"Retriever: {selected_retriever}, Reasoner: {selected_reasoner}"
    st.session_state.messages.append({"role": "user", "content": prompt})
    # Display user message in chat message container
    with st.chat_message("user"):
        st.markdown(prompt)


    retrieved_documents=retriever(prompt)
    reasoning = reasoner(retrieved_documents)

    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        assistant_response = random.choice(
            [
                "The claim is correct.",
                "The claim is incorrect.",
            ]
        )

        # Simulate stream of response with milliseconds delay
        for chunk in assistant_response.split():
            full_response += chunk + " "
            time.sleep(0.05)
            # Add a blinking cursor to simulate typing
            message_placeholder.markdown(full_response + "▌")
        message_placeholder.markdown(full_response)
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": full_response})