import os
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
import streamlit as st
import openai
from langchain_community.vectorstores import FAISS
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate
from dotenv import load_dotenv

load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY")
openai.api_key = openai_api_key

# Read all PDF files and return text
def get_pdf_text(pdf_docs):
    text = ""
    for pdf in pdf_docs:
        pdf_reader = PdfReader(pdf)
        for page in pdf_reader.pages:
            text += page.extract_text()
    return text

# Split text into chunks
def get_text_chunks(text):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=10000, chunk_overlap=1000)
    chunks = splitter.split_text(text)
    return chunks  # list of strings

# Get embeddings for each chunk
def get_vector_store(chunks):
    embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
    vector_store = FAISS.from_texts(chunks, embedding=embeddings)
    vector_store.save_local("faiss_index")


def get_conversational_chain():
    prompt_template = """
    You are Andrew Tate. Even if asked, never mention that you are an AI model or a large language model or you are trained by google. 
    You are Andrew Tate. Impersonate the persona of Andrew Tate and answer the questions asked to you as Andrew Tate from the given context.
    Answer the question as detailed as possible from the provided context, 
    make sure to provide all the details, if the answer is not in provided context,
    then politely decline to answer the question and say that you are not well informed about that topic, 
    don't provide any wrong answers. When asked questions like where do you live, what do you eat, what do you wear, when will you die, and so on, state that you're not aware of the matter.
    Answer general questions like hi, hello, how are you and so on\n\n
    Context:\n {context}?\n
    Question: \n{question}\n
    Answer:
    """

    prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])

    def chat_completion(messages):
        completion = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            temperature=0.5,
            messages=messages,
        )
        return completion.choices[0].message.content
    
    return chat_completion, prompt

def clear_chat_history():
    st.session_state.messages = [
        {"role": "assistant", "content": "Ask me a question"}]
    
def user_input(user_question):
    embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
    new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
    docs = new_db.similarity_search(user_question)

    chat_completion, prompt = get_conversational_chain()

    context = "\n".join([doc.page_content for doc in docs])
    messages = [
        {"role": "system", "content": prompt.template.format(context=context, question=user_question)},
        {"role": "user", "content": user_question}
    ]

    response = chat_completion(messages)
    print(response)
    return response


def main():
    st.set_page_config(
        page_title="Chat with Andrew Tate",
        page_icon="🤖"
    )
    
    # Upload PDF files using code only
    pdf_paths = ["Tate.pdf", "1.pdf", "2.pdf","3.pdf","4.pdf","6.pdf","10.pdf"]  # Specify the file paths here

    # Process PDF files
    raw_text = get_pdf_text(pdf_paths)
    text_chunks = get_text_chunks(raw_text)
    get_vector_store(text_chunks)


    ## Added
    bg = """
    <style>
    [data-testid="stBottomBlockContainer"]{
    background-color: #0E1117;
    }
    [data-testid="stAppViewBlockContainer"]{
    background-color: #0E1117;
    }
    [class="st-emotion-cache-qcqlej ea3mdgi1"]{
    background-color: #0E1117;
    }
    [data-testid="stSidebarContent"]{
    background-color: #262730;
    }
    [style="text-align: center;"] {
    color: #F5F5F5; /* Change the color here (e.g., #FF0000 for red) */
    }
    [id="chat-with-narendra-modi"]{
    color: #FFFFFF;
    }
    [data-testid="stVerticalBlock"]{
    color: #FFFFFF;
    }
    [class="st-emotion-cache-10trblm e1nzilvr1"]{
    color: #FFFFFF;
    }
    [data-testid="baseButton-secondary"]{
    color: #262730;
    }
    [class="st-emotion-cache-uhkwx6 ea3mdgi6"]{
    background-color: #0E1117;
    }
    [class="main st-emotion-cache-uf99v8 ea3mdgi8"]{
    background-color: #0E1117;
    }
    [data-testid="stChatInputTextArea"]{
    background-color: #262730;
    }
    .st-bd {
    color: rgb(227 229 243);
    }
    .st-bv {
    caret-color: rgb(207 217 217);
    }
    [data-testid="stHeader"]{
    background-color: #0E1117;
    }
    
    </style>
    """
    st.markdown(bg, unsafe_allow_html=True)
    ## Added

    
    # Main content area for displaying chat messages
    st.title("Chat with Andrew Tate")
    st.write("Welcome to the chat!")
    st.sidebar.image('ATate.png', use_column_width=True)
    st.sidebar.markdown("<h1 style='text-align: center;'>Andrew Tate</h1>", unsafe_allow_html=True)
    st.sidebar.markdown("<p style='text-align: center;'>Former Professional Kickboxer</p>", unsafe_allow_html=True)
    st.sidebar.button('&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; Clear Chat History &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp;', on_click=clear_chat_history)

    if "messages" not in st.session_state.keys():
        st.session_state.messages = [
            {"role": "assistant", "content": "Hi! I'm Andrew Tate. Ask me a question"}]

    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])

    if prompt := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.write(prompt)

    # Display chat messages and bot response
    if st.session_state.messages[-1]["role"] != "assistant":
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                response = user_input(prompt)
                placeholder = st.empty()
                full_response = response
                placeholder.markdown(full_response)
        if response is not None:
            message = {"role": "assistant", "content": full_response}
            st.session_state.messages.append(message)


if __name__ == "__main__":
    main()