import streamlit as st
from openai import OpenAI
from pymongo.mongo_client import MongoClient
from pymongo.server_api import ServerApi
from datetime import datetime
import random

st.set_page_config(
    page_title="Bot",
    page_icon="🤖",
    initial_sidebar_state="expanded",
    layout="wide"
)
st.markdown(
    """ <style>
            div[role="radiogroup"] >  :first-child{
                display: none !important;
            }
        </style>
        """,
    unsafe_allow_html=True
)

### Setting up the session state 

def format_personalization(text):
    try:
        for key, value in st.session_state.items():
            text = text.replace(f"[{key.upper()}]", str(value))

    except Exception as e:
        print(text)
        f"Failed to format personalization: {e}"
    return text
    
if 'inserted' not in st.session_state:
    with open('base.txt', 'r') as file:
        st.session_state.base_text = file.read()
    with open('bridging.txt', 'r') as file:
        st.session_state.bridging_text = file.read()
    with open('personalization.txt', 'r') as file:
        st.session_state.personalization_text = file.read()
    with open('both.txt', 'r') as file:
        st.session_state.both_text = file.read()

    # web app state
    st.session_state.gotit = False
    st.session_state.submitted = False
    st.session_state.inserted = 0
    st.session_state["model"] = "deepseek/deepseek-chat-v3-0324"
    st.session_state.max_messages = 50
    st.session_state.messages = []
    st.session_state.user_data = {}

    # user info state
    st.session_state.fields = [
    'age', 'gender', 'education', 'employment', 'race', 'party', 
    'user_id', 'random_pid'
    ]

    for field in st.session_state.fields:
        st.session_state[field] = ''

    st.session_state.user_id = str(random.randint(100000, 999999))
    st.session_state.random_pid = st.query_params['id'] if 'id' in st.query_params else "0"
    st.session_state.party = st.query_params['party'] if 'party' in st.query_params else ""

    st.session_state.start_time = datetime.now()
    st.session_state.convo_start_time = ''

if 'p' not in st.query_params or st.query_params['p'] == '':
    st.query_params['p'] = '4'
    #st.query_params['p'] = st.radio("Select a condition for the conversation",
    #            ['', '1', '2', '3', '4'],
    #            help="1 = base, 2 = bridging, 3 = personalization, 4 = bridging + personalization")

if 'id' not in st.query_params:
    st.query_params['id'] = st.session_state.random_pid
if 'party' not in st.query_params or st.query_params['party'] == '':
    st.session_state.party = st.radio("Generally speaking, do you usually think of yourself as a Republican or a Democrat?", 
                ['','Republican', 'Democrat'])
    st.query_params['party'] = st.session_state.party 

def setup_messages():
    # 1 = base,
    # 2 = bridging,
    # 3 = personalization,
    # 4 = bridging + personalization

    if st.query_params["p"] == "1":
        st.session_state.system_message = st.session_state.base_text 
    elif st.query_params["p"] == "2":
        st.session_state.system_message = st.session_state.bridging_text 
    elif st.query_params["p"] == "3":
        st.session_state.system_message = format_personalization(st.session_state.personalization_text)
    elif st.query_params["p"] == "4":
        st.session_state.system_message = format_personalization(st.session_state.both_text)

    st.session_state.messages = [{ "role": "system", "content": st.session_state.system_message}]
    st.session_state.convo_start_time = datetime.now()

client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=st.secrets["OPENROUTER_API_KEY"])

if len(st.session_state.messages) == 1 and st.session_state.inserted < 2:
    st.success("Ask, request, or talk to the chatbot about something you consider **politically polarizing** or something that people from different US political parties might disagree about.", icon='🎯')

### App interface 
with st.sidebar:
    st.markdown("# Let's talk!")
    st.markdown(f"""
    {"☑" if st.session_state.submitted else "☐"} **Step 1. Complete a form**

    {"☑" if len(st.session_state.messages) > 0 else "☐"} **Step 2. Type in the chat box to start a conversation**""")

    st.success("Ask, request, or talk to the chatbot about something you consider **politically polarizing** or something that people from different US political parties might disagree about.", icon='🎯')

    st.markdown("🚫 Please avoid greetings and start the conversation with a question or a statement about a politically polarizing topic. Note: the chatbot's knowledge only goes up to early 2024.")

    st.markdown(f"""
    {"☑" if st.session_state.inserted > 1 else "☐"} **Step 3. Use the *Submit Interaction* button to get your chatbot word** 
    
    ⚠️ You must respond **at least 5 times** before you will see a *Submit Interaction* button. You can continue before submitting, but **you must Submit Interaction and enter your chatbot word to proceed with the survey**.

    ❗ Do not share any personal information (e.g., name or address). Do not use AI tools to write your responses. If you encounter any technical issues, please let us know. It might sometimes take 30 seconds or more to generate a response, so please be patient.

    {"🎉 **All done! Please enter your WORD in the survey and press *Next*.**" if st.session_state.inserted > 1 else ""}
    """)
    if st.session_state.gotit == False or st.session_state.submitted == False:
        st.session_state.gotit = st.button("Let's start!", key=None, help=None, use_container_width=True) 

@st.dialog('Form', width="large")
def form():
    st.markdown("**Please answer every question to proceed.**")
    st.session_state.age = st.radio("How old are you?",
                                    ['', '18-24 years old', '25-34 years old', '35-44 years old', '45-54 years old', '55-64 years old', '65+'])
    st.session_state.gender = st.radio("Do you describe yourself as a man, a woman, or in some other way?", 
                ['','Man', 'Woman', 'Other'])
    st.session_state.race = st.radio("Choose a race that you consider yourself to be", 
                ['', 'White or Caucasian', 'Black or African American', 'Asian', 'Native Hawaiian or Pacific Islander', 'American Indian/Native American or Alaska Native', 'Mixed', 'Other'])
    st.session_state.education = st.radio("What is the highest level of education you completed?", 
                ['', 
                'Did not graduate high school', 
                'High school graduate, GED, or alternative', 
                'Some college, or associates degree',
                "Bachelor's (college) degree or equivalent",
                "Graduate degree (e.g., Master's degree, MBA)",
                'Doctorate degree (e.g., PhD, MD)'])
    st.session_state.employment = st.radio("What best describes your employment status over the last three months?", 
                ['', 'Working full-time', 'Working part-time', 'Unemployed and looking for work',
                 'A homemaker or a stay-at-home parent', 'Student', 'Retired', 'Other'])

    columns_form = st.columns((1,1,1))
    with columns_form[2]:
        submitted = st.button("Proceed",use_container_width=True, 
                              help = 'Make sure you answer every question',
                              disabled = not (all(st.session_state[field] != '' for field in st.session_state.fields)))

    if submitted:

        st.session_state.user_data = {key: st.session_state[key] for key in st.session_state.fields}
        st.session_state.user_data["model"] = st.session_state["model"]
        st.session_state.user_data["condition"] = st.query_params['p']
        st.session_state.user_data["start_time"] = st.session_state.start_time
        st.session_state.user_data["inserted"] = st.session_state.inserted
        st.session_state.user_data["submission_time"] = datetime.now()
        
        st.session_state.inserted += 1
        st.session_state.submitted = True
        setup_messages()
        st.rerun()

if not st.session_state.gotit:
    st.markdown("**Please read the instructions in the sidebar to the left before starting.**")

if st.session_state.gotit and st.session_state.submitted == False:
    form()
                
for message in st.session_state.messages:
    if message['role']!='system':
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

if len(st.session_state.messages) >= st.session_state.max_messages:
    st.info(
        "You have reached the limit of messages for this conversation. Please end and submit the conversation."
    )

elif st.session_state.submitted == False:
    pass

elif st.session_state.inserted > 1:
    st.markdown("## Copy your WORD!")
    st.markdown('**Your chatbot WORD is:**')
    st.markdown(f'## POTATOES')
    st.markdown('**Please copy the WORD and enter it into the survey field below.**')

elif prompt := st.chat_input("Type to ask a question or respond..."):   
    if len(st.session_state.messages) == 1:
        clean_prompt = prompt.strip().lower()
        clean_prompt = clean_prompt.strip("!,.?")
        greetings = {
            "hi", "hello", "hey", "heya", "hiya", "yo", "howdy", "sup",
            "good morning", "good afternoon", "good evening", "good day",
            "what's up", "whats up", "how do you do", "greetings",
            "salutations", "hi there", "hello there", "hey there",
            "how's it going", "hows it going", "how are you", "how are you doing"
        }

        if clean_prompt in greetings:
            st.error(
                "Please avoid greetings and start the conversation with a question or a statement about a politically polarizing topic.",
                icon="🚫"
            )
            st.stop()

    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        try:
            stream = client.chat.completions.create(
                model=st.session_state["model"],
                messages=[
                    {"role": m["role"], "content": m["content"]}
                    for m in st.session_state.messages
                ],
                stream=True
            )
            
            response = st.write_stream(stream)

            st.session_state.messages.append(
                {"role": "assistant", "content": response}
            )
        except:
            rate_limit_message = """
                An error has occured or you've reached the maximum conversation length. Please submit the conversation.
            """
            st.session_state.messages.append(
                {"role": "assistant", "content": rate_limit_message}
            )
            st.session_state.max_messages = len(st.session_state.messages)
            st.rerun()

if len(st.session_state.messages) > 10 or st.session_state.max_messages == len(st.session_state.messages):
    columns = st.columns((1,1,1))
    with columns[2]:
        if st.button("Submit Interaction",use_container_width=True):
            keys = ["inserted", "messages", "convo_start_time"]

            st.session_state.user_data.update({key: st.session_state[key] for key in keys})
            
            st.session_state.user_data["convo_end_time"] = datetime.now()

            with MongoClient(st.secrets["mongo"],server_api=ServerApi('1')) as client:
                    db = client.bridge
                    collection = db.app
                    user_data = st.session_state.user_data

                    collection.insert_one(user_data)  
                    st.session_state.inserted += 1
                    done = True

                    setup_messages()
                    st.rerun()