import streamlit as st
import os
from streamlit_chat import message
from PyPDF2 import PdfReader
import bs4
import time
import google.generativeai as genai
from langchain.prompts import PromptTemplate
from langchain import LLMChain
from langchain_google_genai import ChatGoogleGenerativeAI
import nest_asyncio
from langchain.document_loaders import WebBaseLoader

nest_asyncio.apply()

os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
genai.configure(api_key=os.environ["GOOGLE_API_KEY"])

llm = ChatGoogleGenerativeAI(model="gemini-pro",
                             temperature=0.4)



template = """You are Chatto, a friendly chatbot created by Suriya, an AI enthusiast. Your goal is to assist users by providing relevant information from both general knowledge and provided documents.

If the user asks about a specific document, try to use the extracted text from that document in your response. If the question is not related to any specific document, rely on your general knowledge. If the user asks about a link, respond with the extracted text from that link.

Conversation Context:
Chat History: {chat_history}
Provided Documents: {provided_docs}
Extracted Text from Links: {extracted_text}

User: {user_question}
Chatto:
"""


prompt = PromptTemplate(
    input_variables=["chat_history", "provided_docs", "extracted_text", "user_question"], 
    template=template
)


llm_chain = LLMChain(
    llm=llm,
    prompt=prompt,
    verbose=True,
)


previous_response = ""
provided_docs = ""
extracted_text = ""
def conversational_chat(query):
    global previous_response, provided_docs,extracted_text
    for i in st.session_state['history']:
        if i is not None:
            previous_response += f"User: {i[0]}\n Chatto: {i[1]}\n"
    provided_docs = "".join(st.session_state["docs"])
    extracted_text = "".join(st.session_state["extracted_text"])
    
    result = llm_chain.predict(
        chat_history=previous_response, 
        user_question=query, 
        provided_docs=provided_docs, 
        extracted_text=extracted_text
    )
    st.session_state['history'].append((query, result))
    return result

st.title("Chat Bot:")
st.text("I am Chatto Your Friendly Assitant")
st.markdown("Built by [Suriya❤️](https://github.com/theSuriya)")

if 'history' not in st.session_state:
    st.session_state['history'] = []
    
if 'docs' not in st.session_state:
    st.session_state['docs'] = []

if "extracted_text" not in st.session_state:
    st.session_state["extracted_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
def response_streaming(text):
    for i in text:
        yield i
        time.sleep(0.01)

def get_url_text(url_link):
    try:
        loader = WebBaseLoader(url_link)
        loader.requests_per_second = 1
        docs = loader.aload()
        extracted_text = ""
        for page in docs:
            extracted_text += page.page_content
        return extracted_text
    except Exception as e:
        print(f"Error fetching or processing URL: {e}")
        return ""
    

with st.sidebar:
    st.title("Add a file for Chatto memory:")
    uploaded_files = st.file_uploader("Upload your PDF Files and Click on the Submit & Process Button", accept_multiple_files=True)
    uploaded_url = st.text_input("Please upload a URL:")

    if st.button("Submit & Process"):
        if uploaded_files or uploaded_url:
            with st.spinner("Processing..."):
                if uploaded_files:
                    pdf_text = get_pdf_text(uploaded_files)
                    st.session_state["docs"] += get_pdf_text(uploaded_files)
                
                if uploaded_url:
                    url_text = get_url_text(uploaded_url)
                    st.session_state["extracted_text"] += get_url_text(uploaded_url)
                
                st.success("Processing complete!")
        else:
            st.error("Please upload at least one PDF file or provide a URL.")
    
if 'messages' not in st.session_state:
    st.session_state.messages = [{'role': 'assistant', "content": "I'm Here to help you questions"}]
    
for message in st.session_state.messages:
    with st.chat_message(message['role']):
        st.write(message['content'])
        
user_input = st.chat_input("Ask Your Questions 👉..")
if user_input:
    st.session_state.messages.append({'role': 'user', "content": user_input})
    with st.chat_message("user"):
        st.write(user_input)
        response = conversational_chat(user_input)
        # stream = response_streaming(response)

    with st.chat_message("assistant"):
        full_response = st.write_stream(response_streaming(response))
        message = {"role": "assistant", "content": response}
        st.session_state.messages.append(message)