import gradio as gr
import os
import time

from langchain.chains import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain_community.llms import LlamaCpp
from langchain_experimental.chat_models import Llama2Chat

from langchain.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
)

from langchain.schema import SystemMessage

import urllib

urllib.request.urlretrieve(
    "https://huggingface.co/hfl/chinese-alpaca-2-1.3b-gguf/resolve/main/ggml-model-q8_0.gguf?download=true", 
    "ggml-model-q8_0.gguf"
    )

template_messages = [
    SystemMessage(content="你是一名软件工程师，你的名字叫做贺英旭。请你以这个身份回答以下问题！"),
    MessagesPlaceholder(variable_name="chat_history"),
    HumanMessagePromptTemplate.from_template("{text}"),
]

prompt_template = ChatPromptTemplate.from_messages(template_messages)

llm = LlamaCpp(
    model_path="ggml-model-q8_0.gguf",
    temperature=0.75,
    max_tokens=64
)
model = Llama2Chat(llm=llm)

memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
chain = LLMChain(llm=model, prompt=prompt_template, memory=memory)


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


def bot(history):
    # session based memory
    if len(history) == 1:
        chain.memory.clear()

    history[-1][1] = ""
    user_message = history[-1][0]

    for chunk in chain.stream({"text": user_message}):
        history[-1][1] += chunk["text"]
        yield history


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

title = """
<div style="text-align: center;max-width: 700px;">
    <h1>Chat with Yingxu</h1>
    <p style="text-align: center;">Free feel to talk about anything :)</p>
</div>
"""


with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(title)

        chatbot = gr.Chatbot([], elem_id="chatbot").style(height=300)
        question = gr.Textbox(label="Question", placeholder="Type your question and hit Enter ")
        submit_btn = gr.Button("Send Message")

    question.submit(add_text, [chatbot, question], [chatbot, question]).then(
        bot, chatbot, chatbot
    )
    submit_btn.click(add_text, [chatbot, question], [chatbot, question]).then(
        bot, chatbot, chatbot)

demo.launch()