from pymilvus import MilvusClient, connections
from pymilvus import Collection, DataType, FieldSchema, CollectionSchema, utility

from sentence_transformers import SentenceTransformer
import re

import gradio as gr
import openai

import requests
import json

def connect_milvus(endpoint, token):
    connections.connect(
        "default",
        uri=CLUSTER_ENDPOINT,
        token=TOKEN
    )
    
    client = MilvusClient(
        uri=CLUSTER_ENDPOINT,
        token=TOKEN
    )
    #print(client)
    
def create_vectordb(COLLECTION_NAME, EMBEDDING_DIM=1024):
    
    COLLECTION_NAME = "van_ban_phap_luat"
    check_collection = utility.has_collection(COLLECTION_NAME)
    if check_collection:
        drop_result = utility.drop_collection(COLLECTION_NAME)
        print('drop completed!')
        
    chunk_id = FieldSchema(
        name="chunk_id",
        dtype=DataType.INT64,
        is_primary=True,
        description="Chunk ID"
    )

    chunk_ref = FieldSchema(
        name="chunk_ref",
        dtype=DataType.VARCHAR,
        max_length=512,
        description="Chunk ref")

    chunk_text = FieldSchema(
        name="chunk_text",
        dtype=DataType.VARCHAR,
        max_length=4096,
        description="Chunk text")

    chunk_embedding = FieldSchema(
        name="chunk_embedding",
        dtype=DataType.FLOAT_VECTOR,
        dim=EMBEDDING_DIM,
        description="Chunk Embedding")

    schema = CollectionSchema(
        fields=[chunk_id, chunk_ref, chunk_text, chunk_embedding],
        auto_id=False,
        description="Vector Store Chunk using multilingual-e5-large")

    collection = Collection(
        name=COLLECTION_NAME,
        schema=schema
    )
    
    entities = [ids, rules, chunks, embeddings]
    ins_resp = collection.insert(entities)
    ins_resp
    
    collection.flush()
    
    index_params = {
        "index_type": "IVF_FLAT",
        "metric_type": "COSINE",  # L2
        "params": {}
    }
    collection.create_index(
        field_name=chunk_embedding.name,
        index_params=index_params
    )
    
    collection.load()
    #return None   

def load_vectordb(COLLECTION_NAME):
    
    collection = Collection(
        name=COLLECTION_NAME,
        #schema=schema
    )
    #print(collection)
    #print(collection.has_index())
    return collection

def load_model(model_name):
    model = SentenceTransformer(model_name)

    return model

def search_chunks(query, topk=5):
        # search
    search_params = {
        "metric_type": "COSINE",  # L2
        "params": {"level": 2}
    }

    collection = load_vectordb(COLLECTION_NAME)

    embed_query = model.encode(query)
    results = collection.search(
        [embed_query],
        #anns_field=chunk_embedding.name,
        anns_field="chunk_embedding",
        param=search_params,
        limit=topk,
        guarantee_timestamp=1,
        output_fields=['chunk_ref', 'chunk_text'] #
    )

    refs, relevant_chunks = [], []

    pattern = r"'chunk_ref': '([^']*)', 'chunk_text': '([^']*)'"

    #print(len(results[0]))

    for sample in results[0]:
        #print(sample)
        #print(type(sample))
        matches = re.findall(pattern, str(sample))

        # Lặp qua các kết quả tìm được và thêm vào list
        for match in matches:
            refs.append(match[0])
            relevant_chunks.append(match[1])

    return refs, relevant_chunks

    #return results[0]
    #print(len(results[0]))
    #return results

CLUSTER_ENDPOINT = "https://in03-d63abe0e8a8f47b.api.gcp-us-west1.zillizcloud.com"
TOKEN = "6529d1f59d5e3d38d6135ec9ddf5820a9c38e3db6ca22c53b3aa2ad9c9148e29ef0f11c312bee71f1544da350f85320b598a30f3"
COLLECTION_NAME = "van_ban_phap_luat"
MODEL_NAME = 'qminh369/datn-dense_embedding'
#query = "Cơ quan quản lý giáo dục và đào tạo phải làm gì để đảm bảo pháp luật về giao thông đường bộ được đưa vào chương trình giảng dạy?" # 7
#query = "Trách nhiệm và quản lý hoạt động giao thông đường bộ được phân công và phân cấp như thế nào?" # 4
#query = "Những trường hợp cho phép chở người trên xe ô tô chở hàng?" # 21

connect_milvus(CLUSTER_ENDPOINT, TOKEN)
#load_vectordb(COLLECTION_NAME)
model = load_model(MODEL_NAME)

model_gpt_3_5 = "gpt-3.5-turbo"
openai.api_key = "sk-proj-P6SydZ6cYprqmpMH6BqdT3BlbkFJT4hjCdGZpGZPfcpBqi7q"

def write_prompt(doc_relevant, question):
    
    PROMPT = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n{question}\n\n### Input:\n{doc_relevant}"
    
    input_prompt = PROMPT.format_map(  
    {"doc_relevant": doc_relevant, "question": question}  
)
    return input_prompt

def answer(question):

    refs, relevant_chunks = search_chunks(question)
    #print(results)
    system_msg = "Hãy trả lời câu hỏi sau dựa trên thông tin được cung cấp. Nếu thông tin được cung cấp không liên quan dến câu hỏi thì trả về câu trả lời 'Không có thông tin'. Lưu ý các câu trả lời đều sử dụng Tiếng Việt"
    user_msg = write_prompt(relevant_chunks, question)
    
    response = openai.ChatCompletion.create(model=model_gpt_3_5,
                                            messages=[{"role": "system", "content": system_msg},
                                             {"role": "user", "content": user_msg}])
    
    response = response["choices"][0]["message"]["content"]
    
    ref = "\n" + "Trích dẫn từ: " + refs[0]
    
    response = response + ref
    
    return response.strip()

def chatbot(question, history=[]):
    output = answer(question)
    history.append((question, output))
    return history, history

demo = gr.Interface(
    fn=chatbot,
    inputs=["text", "state"],
    outputs=["chatbot", "state"])

demo.queue().launch(share=True)