import streamlit as st
import os
from streamlit_chat import message
from PyPDF2 import PdfReader
import google.generativeai as genai
from langchain.prompts import PromptTemplate
from langchain import LLMChain
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.tools import DuckDuckGoSearchRun

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

llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest",
                             temperature=0.2)

search_engine = DuckDuckGoSearchRun()

def get_web_result(question):
    result = search_engine.invoke(question)
    return result.strip()


template = """You are a friendly chat assistant called "CRETA" having a conversation with a human and you are created by Pachaiappan [portfolio](https://mr-vicky-01.github.io/Portfolio/) an AI Specialist. If the question pertains to information that you do not have access to, respond with string 'search_query' only nothing else.
provided document:
{provided_docs}
previous_chat:
{chat_history}
Human: {human_input}
Chatbot:"""

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

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

template_2 = """read the web_result and Answer the fallowing question carefully, your response must be short and informative
web_result:
{web_result}
previous_chat:
{chat_history}
Question: {human_input}
Chatbot:"""


prompt2 = PromptTemplate(
    input_variables=["web_result", "chat_history", "human_input", "provided_docs"], template=template_2
)

llm_chain_2 = LLMChain(
    llm=llm,
    prompt=prompt2,
    verbose=True,
)

search_template = """Write a brief, user-friendly search query based on the details below. The response should be concise and ready for direct use on a search engine.
Chat History:
{chat_history}
Question: {human_input}
Search Query::"""

search_prompt = PromptTemplate(
    input_variables=["chat_history", "human_input"], template=search_template
)

search_llm = LLMChain(
    llm=llm,
    prompt=search_prompt,
    # verbose=True,
)

previous_response = ""
provided_docs = ""
def conversational_chat(query):
    global previous_response, provided_docs
    for i in st.session_state['history'][-5:]:
        if i is not None:
            previous_response += f"Human: {i[0]}\n Chatbot: {i[1]}\n"
    provided_docs = "".join(st.session_state["docs"])
    result = llm_chain.predict(chat_history=previous_response, human_input=query, provided_docs=provided_docs)

    if  'search_query'in result.strip():
        search_query = search_llm.predict(chat_history=previous_response, human_input=query)
        print(search_query)
        web_result = get_web_result(search_query.strip())
        result = llm_chain_2.predict(web_result= web_result,chat_history=previous_response, human_input=query, provided_docs=provided_docs)

    st.session_state['history'].append((query, result))
    
    # Keep only the last 5 history entries
    st.session_state['history'] = st.session_state['history'][-5:]
    
    return result


st.title("CRETA 🤖")
st.text("I am CRETA Your Friendly Assitant")

if 'history' not in st.session_state:
    st.session_state['history'] = []
    
# Initialize messages
if 'generated' not in st.session_state:
    st.session_state['generated'] = ["Hello ! Ask me anything"]

if 'past' not in st.session_state:
    st.session_state['past'] = [" "]
    
if 'docs' not in st.session_state:
    st.session_state['docs'] = []
    
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 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 CRETA 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("Paste the Documentation URL:")
    
    if st.button("Submit & Process"):
        if uploaded_files or uploaded_url:
            with st.spinner("Processing..."):
                if uploaded_files:
                    st.session_state["docs"] += get_pdf_text(uploaded_files)
                
                if uploaded_url:
                    st.session_state["docs"] += get_url_text(uploaded_url)
                
                st.success("Processing complete!")
        else:
            st.error("Please upload at least one PDF file or provide a URL.")
            
            
# Create containers for chat history and user input
response_container = st.container()
container = st.container()

# User input form
user_input = st.chat_input("Ask Your Questions 👉..")
with container:
    if user_input:
        output = conversational_chat(user_input)
        # answer = response_generator(output)
        st.session_state['past'].append(user_input)
        st.session_state['generated'].append(output)
        
        
# Display chat history
if st.session_state['generated']:
    with response_container:
        for i in range(len(st.session_state['generated'])):
            if i != 0:
                message(st.session_state["past"][i], is_user=True, key=str(i) + '_user', avatar_style="adventurer")
            message(st.session_state["generated"][i], key=str(i), avatar_style="bottts")