#!/usr/bin/env python3
import csv
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.vectorstores import Pinecone
from langchain.chains import RetrievalQA
import os
import gradio as gr
import time
from fuzzywuzzy import fuzz
import pinecone
from getpass import getpass
from huggingface_hub import HfFileSystem


os.environ["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY")
YOUR_API_KEY = os.environ.get("YOUR_API_KEY")

YOUR_ENV = os.environ.get("YOUR_ENV")

index_name = os.environ.get("index_name")
pinecone.init(
	api_key=YOUR_API_KEY,
	environment=YOUR_ENV
)


model_name = 'text-embedding-ada-002'

embed = OpenAIEmbeddings(
	model=model_name,
	openai_api_key=os.environ["OPENAI_API_KEY"]
)
text_field = "text"
res = embed.embed_documents(text_field)

# switch back to normal index for langchain
index = pinecone.Index(index_name)

vectorstore = Pinecone(index, embed.embed_query, text_field)

llm = ChatOpenAI(
	openai_api_key=os.environ["OPENAI_API_KEY"],
	model_name='gpt-4',
	temperature=0.5 ,
	max_tokens=850
)

qa = RetrievalQA.from_chain_type(
	llm=llm,
	chain_type="stuff",
	retriever=vectorstore.as_retriever()
)

inputs = gr.inputs.Textbox(lines=7, label="Frage:")
outputs = gr.outputs.Textbox(label="Antwort")

query = inputs

def answer_question(query): 
	result = vectorstore.similarity_search(query, k=3)
	llm = ChatOpenAI(
		openai_api_key=os.environ["OPENAI_API_KEY"],
		model_name='gpt-4',
		temperature=0.5 ,
		max_tokens=850
	)
	qa = RetrievalQA.from_chain_type(
		llm=llm,
		chain_type="stuff",
		retriever=vectorstore.as_retriever()
	)
	answer = qa.run(query)
	return {"answer": answer}  # Return the result as a dictionary with an "answer" key

with gr.Blocks() as demo:
	chatbot = gr.Chatbot()
	msg = gr.Textbox()
	clear = gr.Button("Clear")
	
	def respond(message, chat_history):
		#message = message[2:]
		bot_response = answer_question(message)
		if isinstance(bot_response, dict) and "answer" in bot_response:
			bot_message = bot_response["answer"]
		else:
			bot_message = "Sorry, I couldn't generate a response."
				
		# Save to CSV file
		#with open(csv_file, "a", newline='') as f:
		#	writer = csv.writer(f)
		#	writer.writerow([message, bot_message])
			
		chat_history.append((message, bot_message))
		time.sleep(1)
		return "", chat_history
	
	msg.submit(respond, [msg, chatbot], [msg, chatbot])
	clear.click(lambda: None, None, chatbot, queue=False)

with gr.Blocks() as demo:
	instructions = gr.Markdown("## Willkommen zum KFO Abrechnungs-Bot\n\n\nBitte stelle deine Frage in der Textbox unten.")
	chatbot = gr.Chatbot()
	msg = gr.Textbox()
	clear = gr.Button("Clear")
	
	msg.submit(respond, [msg, chatbot], [msg, chatbot])
	clear.click(lambda: None, None, chatbot, queue=False)
	
	
if __name__ == "__main__":
	demo.launch(share=False, inbrowser=True)