import streamlit as st
from utils import *
import constants
import os
import base64

def get_base64(bin_file):
    with open(bin_file, 'rb') as f:
        data = f.read()
    return base64.b64encode(data).decode()

def set_background(png_file):
    bin_str = get_base64(png_file)
    page_bg_img = '''
    <style>
    .stApp {
    background-image: url("data:img.jpg;base64,%s");
    background-size: cover;
    }
    </style>
    ''' % bin_str
    st.markdown(page_bg_img, unsafe_allow_html=True)

st.set_page_config(
    page_title="TCE Chat Bot",
    initial_sidebar_state="collapsed"
)

hide_streamlit_style = """
    <style>
    #MainMenu {visibility: hidden;}
    .stDeployButton {display:none;}
    footer {visibility: hidden;}
    </style>
    """
st.markdown(hide_streamlit_style, unsafe_allow_html=True)

try:
    set_background('./15683.jpg')
except:
    st.warning("Background image not found, using default background.")

if 'HuggingFace_API_Key' not in st.session_state:
    st.session_state['HuggingFace_API_Key'] = os.environ.get("HF_TOKEN", "")
if 'Pinecone_API_Key' not in st.session_state:
    st.session_state['Pinecone_API_Key'] = os.environ.get("PINECONE_API", "")

st.title("🌐 TCE.edu Chat Assistant: Your Friendly Guide to Everything TCE! 🚀")

st.sidebar.title("🗝️")
load_button = st.sidebar.button("Load data to Pinecone", key="load_button")

if load_button:
    if st.session_state['HuggingFace_API_Key'] != "" and st.session_state['Pinecone_API_Key'] != "":
        with st.spinner("Loading data..."):
            site_data = get_website_data(constants.WEBSITE_URL)
            st.write("✅ Data fetched successfully!")
            chunks_data = split_data(site_data)
            st.write("✂️ Data split into manageable parts!")
            embeddings = create_embeddings()
            st.write("🧠 Model ready to understand your queries!")
            push_to_pinecone(st.session_state['Pinecone_API_Key'], constants.PINECONE_INDEX, embeddings, chunks_data)
            st.write("🚀 Data loaded into Pinecone for quick searching!")
            st.sidebar.success("🎉 Data successfully loaded into Pinecone!")
    else:
        st.sidebar.error("❌ Oops! Please provide your API keys.")

prompt = st.text_input('How can I help you today ❓', key="prompt")
document_count = st.slider('Number of results to show 🔗 - (0 LOW || 5 HIGH)', 0, 5, 2, step=1)
submit = st.button("Ask! 🔍")

if submit:
    if st.session_state['HuggingFace_API_Key'] != "" and st.session_state['Pinecone_API_Key'] != "":
        with st.spinner("Processing your query..."):
            embeddings = create_embeddings()
            st.write("🧠 Model ready to understand your queries!")
            index = pull_from_pinecone(st.session_state['Pinecone_API_Key'], constants.PINECONE_INDEX, embeddings)
            st.write("🔍 Database retrieval is done!")
            relavant_docs = get_similar_docs(index, prompt, document_count)
            
            st.success("🎉 Here are the search results:")
            st.write("📜 List of search results:")
            for document in relavant_docs:
                st.write("👉**Result : " + str(relavant_docs.index(document)+1) + "**")
                st.write("**Info:**: " + document.page_content)
                st.write("🔗 **Link**: " + document.metadata['source'])
    else:
        st.sidebar.error("❌ Oops! Please provide your API keys.")