import streamlit as st 

from groq import Groq
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings

def run_groq(prompt):

    client = Groq(api_key = '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')

    chat_completion = client.chat.completions.create(
        messages = [
            {
                'role' : 'user' , 
                'content' : prompt
            }
        ] , model = 'llama3-70b-8192'
    )

    return chat_completion.choices[0].message.content

vc = FAISS.load_local(
    'vc' , 
    embeddings = HuggingFaceEmbeddings(model_name = 'all-MiniLM-L6-v2') , 
    allow_dangerous_deserialization = True
)

def search(query) : 

    similar_docs = vc.similarity_search(query) 
    context = '\n'.join([doc.page_content for doc in similar_docs])

    prompt = f'''
    You are given a user query, some textual context and rules, all inside xml tags. You have to answer the query based on the context while respecting the rules.

    <context>
    {context}
    </context>

    <rules>
    - If you don't know, just say so.
    - If you are not sure, ask for clarification.
    - Answer in the same language as the user query.
    - If the context appears unreadable or of poor quality, tell the user then answer as best as you can.
    - If the answer is not in the context but you think you know the answer, explain that to the user then answer with your own knowledge.
    - Answer directly and without using xml tags.

    - The context will contain some youtube video ids,
    - Try to cite links using the ids each time you answer the query
    - When citing the links , use the format https://www.youtube.com/watch?v=<id>
    - DO not mention that you are using the context or the infromation is from the context or the information provided, treat the given context as your learned knowldge
    - The context will also have yotube video_title at the end after the video_id, only mention the video in the citation if the title matches with the query. ALso mention the video title while giving citations
    - Always provide the citations at the end with their title and link
    
    - Try to strucutre your response in markdown format, use points, tables, bold formats and much more to beautify the response
    </rules>

    <user_query>
    {query}
    </user_query>
    '''

    response = run_groq(prompt)

    return response

def check_prompt(prompt) : 

    try : prompt.replace('' , '') ; return True 
    except : return False

def check_mesaage() : 

    if 'messages' not in st.session_state : st.session_state.messages = []

check_mesaage()

for message in st.session_state.messages : 

    with st.chat_message(message['role']) : st.markdown(message['content'])

prompt = st.chat_input('Ask me anything')

if check_prompt(prompt) :

    with st.chat_message('user'): st.markdown(prompt)

    st.session_state.messages.append({
        'role' : 'user' , 
        'content' : prompt
    })

    if prompt != None or prompt != '' : 

        response = search(prompt)

        with st.chat_message('assistant') : st.markdown(response)

        st.session_state.messages.append({
            'role' : 'assistant' , 
            'content' : response
        })