from typing import Callable, Optional

import gradio as gr
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Zilliz
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains import RetrievalQAWithSourcesChain
import uuid
from project.llm.zhipuai_llm import ZhipuAILLM
from project.prompt.answer_by_private_prompt import (
    COMBINE_PROMPT,
    EXAMPLE_PROMPT,
    QUESTION_PROMPT,
)

chain: Optional[Callable] = None


def web_loader(file, openai_key, puzhiai_key, zilliz_uri, user, password):
    if not file:
        return "please upload file"
    loader = TextLoader(file)
    docs = loader.load()

    text_splitter = CharacterTextSplitter(chunk_size=1024, chunk_overlap=0)
    docs = text_splitter.split_documents(docs)
    embeddings = OpenAIEmbeddings(model="text-embedding-ada-002", openai_api_key=openai_key)

    if not embeddings:
        return "embeddings not"

    texts = [d.page_content for d in docs]

    docsearch = Zilliz.from_documents(
        docs,
        embedding=embeddings,
        ids=[str(uuid.uuid4()) for _ in range(len(texts))],
        connection_args={
            "uri": zilliz_uri,
            "user": user,
            "password": password,
            "secure": True,
        },
    )

    if not docsearch:
        return "docsearch not"

    global chain
    #chain = RetrievalQAWithSourcesChain.from_chain_type(
    #    ZhipuAILLM(model="glm-3-turbo", temperature=0.1, zhipuai_api_key=puzhiai_key),
    #    chain_type="refine",
    #    retriever=docsearch.as_retriever(),
    #)

    chain = RetrievalQAWithSourcesChain.from_llm(
        ZhipuAILLM(model="glm-3-turbo", temperature=0.1, zhipuai_api_key=puzhiai_key),
        EXAMPLE_PROMPT,
        QUESTION_PROMPT,
        COMBINE_PROMPT,
        retriever=docsearch.as_retriever(),
    )
    return "success to load data"


def query(question):
    global chain
    # "What is milvus?"
    if not chain:
        return "please load the data first"
    return chain(inputs={"question": question}, return_only_outputs=True).get(
        "answer", "fail to get answer"
    )


if __name__ == "__main__":
    block = gr.Blocks()
    with block as demo:
        gr.Markdown(
            """
        <h1><center>Langchain And Zilliz App</center></h1>
        
        v.2.27.13.58
        
        """
        )
        # url_list_text = gr.Textbox(
        #     label="url list",
        #     lines=3,
        #     placeholder="https://milvus.io/docs/overview.md",
        # )
        file = gr.File(label='请上传知识库文件',
                       file_types=['.txt', '.md', '.docx', '.pdf'])
        openai_key_text = gr.Textbox(label="openai api key", type="password", placeholder="sk-******")
        puzhiai_key_text = gr.Textbox(label="puzhi api key", type="password", placeholder="******")
        with gr.Row():
            zilliz_uri_text = gr.Textbox(
                label="zilliz cloud uri",
                placeholder="https://<instance-id>.<cloud-region-id>.vectordb.zillizcloud.com:<port>",
            )
            user_text = gr.Textbox(label="username", placeholder="db_admin")
            password_text = gr.Textbox(
                label="password", type="password", placeholder="******"
            )
        loader_output = gr.Textbox(label="load status")
        loader_btn = gr.Button("Load Data")
        loader_btn.click(
            fn=web_loader,
            inputs=[
                file,
                openai_key_text,
                puzhiai_key_text,
                zilliz_uri_text,
                user_text,
                password_text,
            ],
            outputs=loader_output,
            api_name="web_load",
        )

        question_text = gr.Textbox(
            label="question",
            lines=3,
            placeholder="What is milvus?",
        )
        query_output = gr.Textbox(label="question answer", lines=3)
        query_btn = gr.Button("Generate")
        query_btn.click(
            fn=query,
            inputs=[question_text],
            outputs=query_output,
            api_name="generate_answer",
        )

        demo.queue().launch(server_name="0.0.0.0", share=False)