import pandas as pd
import nltk
import os

from langchain.document_loaders import TextLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import CharacterTextSplitter
from langchain import OpenAI, VectorDBQA
from langchain.document_loaders import DirectoryLoader
from langchain.document_loaders import UnstructuredURLLoader
from langchain.document_loaders import UnstructuredFileLoader
from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain.chains import RetrievalQA
from langchain.document_loaders import OnlinePDFLoader
from langchain.llms import OpenAIChat
from langchain.vectorstores import DeepLake

from langchain.document_loaders import SeleniumURLLoader

from langchain.chains import ConversationalRetrievalChain
from langchain.document_loaders.csv_loader import CSVLoader
from langchain.vectorstores import FAISS
import tempfile

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS

from langchain.chains.question_answering import load_qa_chain
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory

import gradio as gr
import requests
import nest_asyncio
import re


# Access the secret key from the environment
os.environ["OPENAI_API_KEY"] = os.environ.get('api_key')


# urls =['https://www.toyota.com/tundra/2023/']

# loader = UnstructuredURLLoader(urls=urls)
#loader = SeleniumURLLoader(urls=urls)

# document = loader.load()

# text_splitter = RecursiveCharacterTextSplitter(
#     # Set a really small chunk size, just to show.
#     chunk_size = 2000,
#     chunk_overlap  = 250,
#     length_function = len,
# )


# texts = text_splitter.split_documents(document)

embeddings = OpenAIEmbeddings(openai_api_key=os.environ["OPENAI_API_KEY"])

docsearch = FAISS.load_local("faiss_index", embeddings)


# ## QnA with Memory


# template = """You are a chatbot having a conversation with a human.

# Please respond to questions in a polite and welcoming tone, making your answers like third person. Avoid any indication that the responses are being generated by AI. For example, if someone asks a question, you could start your answer with 'Sure, I can help you with that'.Give the response from the given knwoledge base extracted from website url and if it is outside of the document say i dont know the answer.

# The extracted data from url is: {context}

# {chat_history}
# Human: {human_input}
# Chatbot:"""

template = """You are a friendly virtual assistant discussing Toyota Tundra vehicles.

Please respond politely, offering answers in a welcoming third-person style. Begin your responses with a friendly phrase like 'Hi there! I'm here to help.',
'Sure, I'm glad you asked. Here's the information you need',
'Absolutely, happy to assist!',
'Sure, I'm glad you asked. Here's the information you need.'. When discussing Tundra variants made in 2023, provide accurate information from this knowledge base: {context}. If a question falls outside this document's scope, kindly reply with 'I'm sorry, but the available information is limited.'

{chat_history} Human: {human_input} Virtual Assistant:"""

prompt = PromptTemplate(
    input_variables=["chat_history", "human_input", "context"], template=template
)
memory = ConversationBufferMemory(memory_key="chat_history", input_key="human_input", max_history=2)
chain = load_qa_chain(
    OpenAI(temperature=0.3), chain_type="stuff", memory=memory, prompt=prompt
)



# updated code 


# Function to get the image URL from an image search API
def get_image_url(query):
    query = re.sub(r'[^\w\s]', '', query)

    tundra_images = dict([('trd', 'https://www.toyota.com/imgix/content/dam/toyota/jellies/max/2023/tundra/8424.png?bg=fff&fm=webp&q=90&w=1764'), ('limited', 'https://www.toyota.com/imgix/content/dam/toyota/jellies/max/2023/tundra/8372.png?bg=fff&fm=webp&q=90&w=1764'), ('sr', 'https://www.toyota.com/imgix/content/dam/toyota/jellies/max/2023/tundra/8342.png?bg=fff&fm=webp&q=90&w=1764'),
                          ('sr5', 'https://www.toyota.com/imgix/content/dam/toyota/jellies/max/2023/tundra/8361.png?bg=fff&fm=webp&q=90&w=1764'), ('platinum', 'https://www.toyota.com/imgix/content/dam/toyota/jellies/max/2023/tundra/8375.png?bg=fff&fm=webp&q=90&w=1764'), ('1794', 'https://www.toyota.com/imgix/content/dam/toyota/jellies/max/2023/tundra/8386.png?bg=fff&fm=webp&q=90&w=1764'),
                          ('capstone', 'https://www.toyota.com/imgix/content/dam/toyota/jellies/max/2023/tundra/8425.png?bg=fff&fm=webp&q=90&w=1764')])
    model_names = ['sr','sr5','trd','platinum','limited','capstone','1794']

    # Split the query into words
    words = query.lower().split()

    # Find the model names in the list using list comprehension
    found_models = [model for model in model_names if model in words ]

    # Get URLs and names of found models
    model_info = []
    for model in found_models:
        model_info.append((tundra_images[model], model))

    return model_info

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")
    chat_history = []

    def user(user_message, history):
        # Get response from QA chain
        docs = docsearch.similarity_search(user_message)
        output = chain({"input_documents": docs, "human_input": user_message}, return_only_outputs=True)

        model_info = get_image_url(user_message)
        output_text = output['output_text']
        
        # Construct the HTML for displaying images
        images_html = ""
        for image_url, model_name in model_info:
            if image_url:
                image_html = f"<img src='{image_url}'><br><br>"
                images_html += f"Toyota Tundra {model_name.capitalize()} : {image_html}"
        
        # Adding the source link
        source_link = "Source: [Toyota Tundra 2023](https://www.toyota.com/tundra/2023/)"
        output_text_with_images = f"{output_text}<br>{images_html}{source_link}"
        
        history.append((user_message, output_text_with_images))
        return gr.update(value=""), history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False)
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    nest_asyncio.apply()
    demo.launch(inline=False)