import logging
import os
import pprint
import uuid
from typing import List

import chromadb
import gradio as gr
import requests
import zhipuai
from bs4 import BeautifulSoup
from dotenv import load_dotenv, find_dotenv
# Import langchain stuff
from langchain.chains import ConversationalRetrievalChain
from langchain.document_loaders import DirectoryLoader
from langchain.memory import ConversationBufferMemory
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.chat_models import ChatOpenAI
from langchain_community.document_loaders import AsyncChromiumLoader, AsyncHtmlLoader
from langchain_community.document_transformers import BeautifulSoupTransformer
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores.chroma import Chroma
from langchain_core.documents import Document
from langchain_core.prompts import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
from llms.zhipuai_llm import ZhipuAILLM
from langchain.chains import create_extraction_chain

_ = load_dotenv(find_dotenv())  # 读取并加载环境变量，来自 .env 文件

os.environ["http_proxy"] = os.environ["PROXY"]
os.environ["https_proxy"] = os.environ["PROXY"]
os.environ["no_proxy"] = os.environ["NO_PROXY"]
os.environ['CURL_CA_BUNDLE'] = ''

# 填写控制台中获取的 APIKey 信息
zhipuai.api_key = os.environ["ZHIPUAI_API_KEY"]
# # LLM Model
llm = ZhipuAILLM(model="chatglm_turbo", temperature=0.9, top_p=0.1, zhipuai_api_key=zhipuai.api_key)

# Log setup
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
)
logger = logging.getLogger(__name__)

COLLECTION_NAME = "webfaq_en"
PERSIST_DIRECTORY = "./database/cncbi/en/"
PATH_TO_SFT_JSON_FILES = './sft/'
REF_WEBSITE_LINK = ["https://www.cncbinternational.com/personal/e-banking/inmotion/en/support/index.html"]

CHROMA_CLIENT = chromadb.PersistentClient(path=PERSIST_DIRECTORY)
CHROMA_COLLECTION = CHROMA_CLIENT.get_or_create_collection(name=COLLECTION_NAME)


RAG_TEMPLATE = """You are a customer service agent of China CITIC Bank International, and please respond to the quesCHROMA_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2" # or use multilingual sentence-transformers/LaBSEtion at the end. If the question is not related to the bank's customer service, you have to decline answering and politely inform the user that you are only tuned to bank customer service. Do not make up the answer from your general knowledge, and if you cannot find reference information from the below Frequently Asked Questions and Answers, just refer the customer to the customer hotline at 22876767.

Frequently Asked Questions and Answers:
{context}

Chat history:
{chat_history}

Question: {question}

Helpful Answer:"""


class QAPair:
    def __init__(self, question, answers):
        self.question = question
        self.answers = answers

    def __str__(self):
        return f'question: {self.question} , answers: {"; ".join(self.answers)}'


def scrape_webpages(urls):
    faq_listings = {}
    for url in urls:
        logger.info("fetching page " + url)
        loader = requests.get(url)
        soup = BeautifulSoup(loader.content, 'html.parser')

        q_listings = {}
        a_listings = {}
        qa_listings = {}

        faq_content = soup.find('div', class_='faq-contain')
        logger.debug("faq_content")
        logger.debug(faq_content)
        q_items = faq_content.find_all(class_='faq-question-wrapper')
        a_items = faq_content.find_all(class_='faq-answer-wrapper')

        k = 0
        for q_item in q_items:
            logger.debug("q_item on key = " + str(k))
            logger.debug(q_item)
            questions = q_item.find_all('p')
            for question in questions:
                if len(question.text.strip()) > 0:
                    q_listings.setdefault(k, []).append(question.text.strip())
            k = k + 1

        k = 0
        for a_item in a_items:
            logger.debug("a_item on key = " + str(k))
            logger.debug(a_item)
            answers = a_item.find_all(['p', 'li'])
            for answer in answers:
                if len(answer.text.strip()) > 0:
                    a_listings.setdefault(k, []).append(answer.text.strip())
            k = k + 1

        for q in q_listings:
            qa_listings[q] = {(tuple(q_listings[q]), tuple(a_listings[q]))}

        logger.debug(qa_listings)
        faq_listings.setdefault(url, []).append(qa_listings)
    return faq_listings


# extracted_content = scrape_with_playwright(REF_WEBSITE_LINK)
# logger.info(extracted_content)

def extract_docs(urls):
    my_docs: List[Document] = list()
    for k, v in scrape_webpages(urls).items():
        logger.info("parsing page " + k)
        for doc in v:
            for pair in doc:
                questions = list(doc[pair])[0][0][0]
                answers = list(doc[pair])[0][1]
                qa_pair = QAPair(questions.strip(), answers)
                my_docs.append(Document(page_content=str(qa_pair), metadata={"source": k}))
    return my_docs


# 初始化加载器
text_splitter = CharacterTextSplitter(chunk_size=1024, chunk_overlap=0)
# 切割加载的 document
split_docs = text_splitter.split_documents(extract_docs(REF_WEBSITE_LINK))

# RAG VectorSearch: 将 document 通过 openai 的 embeddings 对象计算 embedding 向量信息并临时存入 Chroma 向量数据库，用于后续匹配查询
logger.info("building vector database index ...")
embeddings = HuggingFaceEmbeddings(model_name=CHROMA_EMBEDDING_MODEL)

if CHROMA_COLLECTION.count() > 0:
    vectorstore = Chroma(client=CHROMA_CLIENT,
                         embedding_function=embeddings,
                         collection_name=COLLECTION_NAME,
                         persist_directory=PERSIST_DIRECTORY)
else:
    vectorstore = Chroma.from_documents(split_docs,
                                        embedding=embeddings,
                                        collection_name=COLLECTION_NAME,
                                        persist_directory=PERSIST_DIRECTORY
                                        )
    vectorstore.persist()

chroma_retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
custom_question_prompt = PromptTemplate(input_variables=["context", "question", "chat_history"], template=RAG_TEMPLATE)


def querying(query, history):
    # 定义内存记忆
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

    if history:
        logger.debug("chat history:")
        logger.debug(history)
        for itemset in history:
            logger.debug("input:" + itemset[0] + "; output: " + itemset[1])
            msg_human = itemset[0]
            msg_bot = itemset[1]
            memory.save_context({"input": msg_human}, {"output": msg_bot})

    qa_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=chroma_retriever,
        memory=memory,
        verbose=True,
        combine_docs_chain_kwargs={"prompt": custom_question_prompt}
    )
    logger.info("memory:")
    logger.debug(memory.chat_memory.messages)
    logger.debug("question: " + query)

    result = qa_chain({"question": query})
    logger.debug("answer: " + result["answer"].strip())
    return result["answer"].strip().replace("\\n", "</br>")


# Launch the interface
# gr.ChatInterface(querying).launch(share=False)
gr.ChatInterface(querying, title="This is an AI chatbot for customer service").launch(share=False,
                                                                                      server_name="0.0.0.0",
                                                                                      server_port=7865)

langchain.text_splitter.CharacterTextSplitter langchain_community.vectorstores.chroma.Chroma.from_documents langchain.chains.ConversationalRetrievalChain.from_llm langchain.memory.ConversationBufferMemory langchain_community.vectorstores.chroma.Chroma langchain_core.prompts.PromptTemplate langchain_community.embeddings.HuggingFaceEmbeddings
[((1392, 1490), 'llms.zhipuai_llm.ZhipuAILLM', 'ZhipuAILLM', ([], {'model': '"""chatglm_turbo"""', 'temperature': '(0.9)', 'top_p': '(0.1)', 'zhipuai_api_key': 'zhipuai.api_key'}), "(model='chatglm_turbo', temperature=0.9, top_p=0.1,\n    zhipuai_api_key=zhipuai.api_key)\n", (1402, 1490), False, 'from llms.zhipuai_llm import ZhipuAILLM\n'), ((1500, 1665), 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.INFO', 'format': '"""%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s"""', 'datefmt': '"""%Y-%m-%d %H:%M:%S"""'}), "(level=logging.INFO, format=\n    '%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s'\n    , datefmt='%Y-%m-%d %H:%M:%S')\n", (1519, 1665), False, 'import logging\n'), ((1680, 1707), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (1697, 1707), False, 'import logging\n'), ((1940, 1989), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'path': 'PERSIST_DIRECTORY'}), '(path=PERSIST_DIRECTORY)\n', (1965, 1989), False, 'import chromadb\n'), ((5176, 5231), 'langchain.text_splitter.CharacterTextSplitter', 'CharacterTextSplitter', ([], {'chunk_size': '(1024)', 'chunk_overlap': '(0)'}), '(chunk_size=1024, chunk_overlap=0)\n', (5197, 5231), False, 'from langchain.text_splitter import CharacterTextSplitter\n'), ((5489, 5545), 'langchain_community.embeddings.HuggingFaceEmbeddings', 'HuggingFaceEmbeddings', ([], {'model_name': 'CHROMA_EMBEDDING_MODEL'}), '(model_name=CHROMA_EMBEDDING_MODEL)\n', (5510, 5545), False, 'from langchain_community.embeddings import HuggingFaceEmbeddings\n'), ((6235, 6333), 'langchain_core.prompts.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['context', 'question', 'chat_history']", 'template': 'RAG_TEMPLATE'}), "(input_variables=['context', 'question', 'chat_history'],\n    template=RAG_TEMPLATE)\n", (6249, 6333), False, 'from langchain_core.prompts import PromptTemplate\n'), ((1084, 1097), 'dotenv.find_dotenv', 'find_dotenv', ([], {}), '()\n', (1095, 1097), False, 'from dotenv import load_dotenv, find_dotenv\n'), ((5599, 5733), 'langchain_community.vectorstores.chroma.Chroma', 'Chroma', ([], {'client': 'CHROMA_CLIENT', 'embedding_function': 'embeddings', 'collection_name': 'COLLECTION_NAME', 'persist_directory': 'PERSIST_DIRECTORY'}), '(client=CHROMA_CLIENT, embedding_function=embeddings, collection_name\n    =COLLECTION_NAME, persist_directory=PERSIST_DIRECTORY)\n', (5605, 5733), False, 'from langchain_community.vectorstores.chroma import Chroma\n'), ((5828, 5958), 'langchain_community.vectorstores.chroma.Chroma.from_documents', 'Chroma.from_documents', (['split_docs'], {'embedding': 'embeddings', 'collection_name': 'COLLECTION_NAME', 'persist_directory': 'PERSIST_DIRECTORY'}), '(split_docs, embedding=embeddings, collection_name=\n    COLLECTION_NAME, persist_directory=PERSIST_DIRECTORY)\n', (5849, 5958), False, 'from langchain_community.vectorstores.chroma import Chroma\n'), ((6388, 6461), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (6412, 6461), False, 'from langchain.memory import ConversationBufferMemory\n'), ((6814, 6987), 'langchain.chains.ConversationalRetrievalChain.from_llm', 'ConversationalRetrievalChain.from_llm', ([], {'llm': 'llm', 'retriever': 'chroma_retriever', 'memory': 'memory', 'verbose': '(True)', 'combine_docs_chain_kwargs': "{'prompt': custom_question_prompt}"}), "(llm=llm, retriever=chroma_retriever,\n    memory=memory, verbose=True, combine_docs_chain_kwargs={'prompt':\n    custom_question_prompt})\n", (6851, 6987), False, 'from langchain.chains import ConversationalRetrievalChain\n'), ((3193, 3210), 'requests.get', 'requests.get', (['url'], {}), '(url)\n', (3205, 3210), False, 'import requests\n'), ((3226, 3270), 'bs4.BeautifulSoup', 'BeautifulSoup', (['loader.content', '"""html.parser"""'], {}), "(loader.content, 'html.parser')\n", (3239, 3270), False, 'from bs4 import BeautifulSoup\n'), ((7372, 7450), 'gradio.ChatInterface', 'gr.ChatInterface', (['querying'], {'title': '"""This is an AI chatbot for customer service"""'}), "(querying, title='This is an AI chatbot for customer service')\n", (7388, 7450), True, 'import gradio as gr\n')]