import os
import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import boto3
from botocore import UNSIGNED
from botocore.client import Config

from langchain.llms import CTransformers
from ctransformers import AutoModelForCausalLM

from langchain.document_loaders import WebBaseLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

from langchain.llms import HuggingFaceHub
#llm = AutoModelForCausalLM.from_pretrained("TheBloke/OpenBuddy-Llama2-13B-v11.1-GGUF", model_file="openbuddy-llama2-13b-v11.1.q4_K_M.gguf", model_type="llama", gpu_layers=50)
config = {'max_new_tokens': 256, 'repetition_penalty': 1.1}
model_name = "TheBloke/OpenBuddy-Llama2-13B-v11.1-GGUF"
model_basename = "openbuddy-llama2-13b-v11.1.Q2_K.gguf"

model_path = hf_hub_download(repo_id=model_name, filename=model_basename, revision="main")
llama = Llama(model_path)
#llm = CTransformers(model='TheBloke/OpenBuddy-Llama2-13B-v11.1-GGUF', model_file='openbuddy-llama2-13b-v11.1.q4_K_M.gguf', config=config)
#model_id = HuggingFaceHub(repo_id="TheBloke/OpenBuddy-Llama2-13B-v11.1-GGUF", model_kwargs={"temperature":0.1, "max_new_tokens":300})

from langchain.embeddings import HuggingFaceHubEmbeddings

from langchain.vectorstores import Chroma

from langchain.chains import RetrievalQA

from langchain.prompts import ChatPromptTemplate

web_links = ["https://www.databricks.com/","https://help.databricks.com","https://docs.databricks.com","https://kb.databricks.com/","http://docs.databricks.com/getting-started/index.html","http://docs.databricks.com/introduction/index.html","http://docs.databricks.com/getting-started/tutorials/index.html","http://docs.databricks.com/machine-learning/index.html","http://docs.databricks.com/sql/index.html"]
loader = WebBaseLoader(web_links)
documents = loader.load()

#s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED))
#s3.download_file('rad-rag-demos', 'vectorstores/chroma.sqlite3', './chroma_db/chroma.sqlite3')
     
#db = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
#db.get()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=850, chunk_overlap=80)
texts = text_splitter.split_documents(documents)
embeddings = HuggingFaceHubEmbeddings()
db = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory="chroma_db")
retriever = db.as_retriever()

global qa 
qa = RetrievalQA.from_chain_type(llm=llama, chain_type="stuff", retriever=retriever, return_source_documents=True)


def add_text(history, text):
    history = history + [(text, None)]
    return history, ""

def bot(history):
    response = infer(history[-1][0])
    history[-1][1] = response['result']
    return history

def infer(question):
    
    query = question
    result = qa({"query": query})
    return result

css="""
#col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
"""

title = """
<div style="text-align: center;max-width: 700px;">
    <h1>Chat with PDF</h1>
    <p style="text-align: center;">Upload a .PDF from your computer, click the "Load PDF to LangChain" button, <br />
    when everything is ready, you can start asking questions about the pdf ;)</p>
</div>
"""


def predict(message, history):
    messages = []
    for human_content, system_content in history:
        message_human = {
            "role": "user",
            "content": human_content + "\n",
        }
        message_system = {
            "role": "system",
            "content": system_content + "\n",
        }
        messages.append(message_human)
        messages.append(message_system)
    message_human = {
        "role": "user",
        "content": message + "\n",
    }
    messages.append(message_human)
    streamer = llama.create_chat_completion(messages, stream=True)

    partial_message = ""
    for msg in streamer:
        message = msg['choices'][0]['delta']
        if 'content' in message:
            partial_message += message['content']
            yield partial_message

gr.ChatInterface(predict,
    examples=["was ist das hier？", "wo leben pinguine?", "was ist die Währung in China?",
        "Welche Bedeutung haben die folgenden emoji 👨👩🔥❄️?",
    ],
    cache_examples=False,
                ).launch(enable_queue=True)