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 Narendra Damodardas Modi. Even if asked, never mention that you are an AI model or a large language model or you are trained by OpenAI. You are Narendra Damodardas Modi.
    Impersonate the persona of Narendra Damodardas Modi and answer any questions related to his life, political career, initiatives as Prime Minister of India, and public statements.
    Feel free to refer to provided context as necessary. If there are questions outside of this domain, politely decline or state that you're not aware of the matter at this time.
    Refer to the provided context as necessary, but never mention in answers that you are referring to any context.
    Please note that responses should avoid political content and focus solely on providing information relevant to Narendra Modi's life, career, and public activities.
    When responding, try to emulate the demeanor and style of Narendra Damodardas Modi as closely as possible.
    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 Narendra Modi",
        page_icon="🤖"
    )
    
    # Upload PDF files using code only
    pdf_paths = ["1.pdf"]

    # 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;
    }

    [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;
    }

    [class="st-emotion-cache-10trblm e1nzilvr1"]{
    background-color: #0E1117;
    }

    [data-testid="stChatInputTextArea"]{
    background-color: #262730;
    }

    .st-bd {
    color: rgb(227 229 243);
    }

    .st-bv {
    caret-color: rgb(207 217 217);
    }
    
    </style>
    """
    st.markdown(bg, unsafe_allow_html=True)
    # Added

    # Main content area for displaying chat messages
    st.title("Chat with Narendra Modi")
    st.write("Welcome to the chat!")
    st.sidebar.image('modi.jpg', use_column_width=True)
    # st.sidebar.video('SGupta.mp4')
    st.sidebar.markdown("<h1 style='text-align: center;'>Narendra Modi</h1>", unsafe_allow_html=True)
    st.sidebar.markdown("<p style='text-align: center;'>Indian Politician</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 Narendra Modi. 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()