import streamlit as st
from llm import build_rag_chain
from setup import setup

st.set_page_config(layout="wide")

st.title("LinuxGPT")

l_col, r_col = st.columns((3, 1))

if "trigger" not in st.session_state:
    st.session_state["trigger"] = False

def on_enter():
    st.session_state["trigger"] = True

with r_col:
    submit_button, openai_models = setup()

# chat input goes here:
with l_col:
    user_question = st.text_area(
        "Ask about Linux fundamentals",
        on_change=on_enter,
    )
    # BELOW IS TEMPORARY, JUST FOR DEMOS!
    api_key = st.secrets["api_key"]
    if (submit_button or st.session_state["trigger"]) and api_key and user_question:
        rag_chain = build_rag_chain(api_key=api_key)

        formatted_history = [
            {"role": "user", "content": item["question"]} if idx % 2 == 0
            else {"role": "assistant", "content": item["answer"]}
            for idx, item in enumerate(st.session_state.get("history", []))
        ]
        with st.spinner("Thinking..."):
            # invoke answer
            result = rag_chain.invoke({"input": user_question, "chat_history": formatted_history})
            answer = result['answer']

            st.header("LinuxGPT says:")
            st.write(answer)

            # add to chat history automatically
            if 'history' not in st.session_state:
                st.session_state['history'] = []
            st.session_state['history'].append({"question": user_question, "answer": answer})

    else:
        st.write("Please provide an input. No API key is needed for demoing")

# clear history
if st.button("Clear History"):
    st.session_state['history'] = []
    st.write("Chat history cleared")