import os
from concurrent import futures

import langchain

import proto.document_query_pb2
import proto.document_query_pb2_grpc
import grpc
import proto.documents_pb2
import init
import doc_client
from langchain.llms.openai import OpenAI
from langchain.schema.document import Document
from langchain.embeddings import OpenAIEmbeddings
from langchain.chains.qa_with_sources import load_qa_with_sources_chain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.cache import InMemoryCache

langchain.llm_cache = InMemoryCache()

CHUNK_SIZE = 500


class AIServer(proto.document_query_pb2_grpc.DocumentQuery):
    def Query(self, target,
              options=(),
              channel_credentials=None,
              call_credentials=None,
              insecure=False,
              compression=None,
              wait_for_ready=None,
              timeout=None,
              metadata=None):


        print("新的请求：" + target.question)
        vec = init.text_to_vector(target.question)

        question = "Reply in spoken language:" + target.question

        search_param = {
            "data": [vec],
            "anns_field": "vector",
            "param": {"metric_type": "L2"},
            "limit": 5,
            "expr": "user_id == " + str(target.user_id),
            "output_fields": ["document_id", "user_id"],
        }

        res = init.collection.search(**search_param)

        # # 最多 5 个
        # if len(res[0]) > 5:
        #     res[0] = res[0][:5]


        # document_chunk_ids = []
        real_document = []

        for i in range(len(res[0])):
            _chunk_id = res[0][i].id
            print("正在获取分块 " + str(_chunk_id) + " 的内容...")

            try:
                _chunk_content = doc_client.stub.GetDocumentChunk(proto.documents_pb2.GetDocumentChunkByIdRequest(
                    id=_chunk_id
                ))

                _doc_content_full = _chunk_content.content

                doc_obj = Document(page_content=_doc_content_full, metadata={"source": "chunked content"})

                real_document.append(doc_obj)

            except Exception as e:
                print(e)

        print(real_document)

        print("正在调用 LLM...")
        chain = load_qa_with_sources_chain(OpenAI(temperature=0), chain_type="map_reduce",
                                           return_intermediate_steps=True,
                                           verbose=True)

        output = chain({"input_documents": real_document, "question": question}, return_only_outputs=False)
        print("回复：" + output["output_text"])

        return proto.document_query_pb2.QueryResponse(
            text=output["output_text"]
        )

    def Chunk(self,
              target,
              options=(),
              channel_credentials=None,
              call_credentials=None,
              insecure=False,
              compression=None,
              wait_for_ready=None,
              timeout=None,
              metadata=None):

        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=CHUNK_SIZE,
            chunk_overlap=20,
            length_function=len,
            add_start_index=True,
        )

        page_contents = text_splitter.create_documents([
            target.text
        ])

        texts = []

        for page_content in page_contents:
            texts.append(page_content.page_content)

        return proto.document_query_pb2.ChunkResponse(
            texts=texts
        )


def serve():
    _ADDR = os.getenv("BIND")
    if _ADDR is None:
        _ADDR = "[::]:50051"
    print("Listening on", _ADDR)

    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    proto.document_query_pb2_grpc.add_DocumentQueryServicer_to_server(AIServer(), server)

    server.add_insecure_port(_ADDR)
    server.start()
    server.wait_for_termination()


if __name__ == '__main__':
    serve()