import openai
import gradio as gr
from datetime import datetime
import json
import chromadb
import pandas as pd
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
from FlagEmbedding import BGEM3FlagModel
import re

# Đặt API key của bạn ở đây
api_key = "sk-optimi-ai-zR7o9I8TWt6vOwOHJFEpT3BlbkFJlGLNXidFLiV3E10jQv1p"
openai.api_key = api_key

# Khởi tạo mô hình embedding


# Đọc dữ liệu từ JSONL
data_path = "/home/acer/Documents/ABI/Lich_optimi/nhansu_optimi.xlsx"
data = pd.read_excel(data_path)
data['index'] = range(len(data))

# Hàm truy vấn cơ sở dữ liệu để tìm tên đầy đủ của người tham gia
def get_full_name(participant):
    list_name = []
    for i in range(len(data)):
        if participant in data['Họ và tên'][i]:
            #print(data['Họ và tên'][i])
            list_name.append(data['Họ và tên'][i])
    return list_name
def replace_participants(reply, dict_replace):
    for key, names in dict_replace.items():
        if names:
        # Join names with comma for display
            names_str = ','.join(names)
        # Replace the key with the list of names
            reply = reply.replace(key, names_str)
        else:
            reply = reply.replace(key,'')
    parts = reply.split('Người tham gia:')
    if len(parts) > 1:
        participants = parts[1].strip().split(',')
        cleaned_participants = [p.strip() for p in participants if p.strip()]
        reply = parts[0] + 'Người tham gia: ' + ', '.join(cleaned_participants)
    return reply
# Định nghĩa hệ thống prompt
def create_system_prompt():
    localtime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    system_prompt = f"Tách yêu cầu sau thành 5 trường thông tin gồm: Tên cuộc họp (Ngắn gọn dưới 10 từ bao quát về mục đích cuộc họp), Ngày họp (DD/MM/YYYY), Giờ học bắt đầu (HH/MM), Giờ học kết thúc (HH/MM) và Người tham gia (Chỉ trích dẫn tên).\
    Chú ý xác định thời điểm hiện tại là: {localtime} từ đó xác định ngày họp cụ thể bao gồm ngày, tháng, năm.\
    Ví dụ: Tách yêu cầu 'họp với team AI, anh Hoa và giám đốc Dương design search ngày mai 10h sáng đến 3 giờ chiều.' thành:\
    Tên cuộc họp: Design search.\
    Ngày họp: 15/06/2024.\
    Giờ bắt đầu: 10:00.\
    Giờ kết thúc: 15:00.\
    Người tham gia: Hoa,Dương.\
    Bắt đầu tách yêu cầu sau: "
    return system_prompt

# Hàm dự đoán cho Gradio
def predict(message, history):
    # Tạo prompt hệ thống
    system_prompt = create_system_prompt()
    
    # Chuẩn bị lịch sử hội thoại
    history_openai_format = [{"role": "system", "content": system_prompt}]
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human})
        history_openai_format.append({"role": "assistant", "content": assistant})
    history_openai_format.append({"role": "user", "content": message})
    
    # Gọi API của OpenAI
    response = openai.ChatCompletion.create(
        model='gpt-3.5-turbo',
        messages=history_openai_format,
        temperature=1.0
    )
    
    # Trích xuất phản hồi từ mô hình
    reply = response.choices[0].message['content']
    #print(reply)
    #print(type(reply))
    # Tách các tên người tham gia từ phản hồi
    # Tách các tên người tham gia từ phản hồi
    start_index = reply.find("Người tham gia:") + len("Người tham gia:")
    end_index = reply.find("\n", start_index)

# Trích xuất danh sách người tham gia
    participants_str = reply[start_index:end_index].strip()

# Tách danh sách thành từng phần tử
    participants = [participant.strip() for participant in participants_str.split(",")]
    participants = [participant.rstrip('.') for participant in participants_str.split(",")]
    print(participants)
    # Thay thế các tên người tham gia bằng tên đầy đủ
    replace_dict = {participant: get_full_name(participant) for participant in participants}
    print(replace_dict)
    # replace_dict = {key: value for key, value in replace_dict.items() if value != []}
    # print(replace_dict)
    reply = replace_participants(reply,replace_dict)
    print(reply)
    # Cập nhật lịch sử hội thoại
    history.append((message, reply))
    
    return reply

# Khởi tạo giao diện Gradio
iface = gr.ChatInterface(fn=predict, title="Tách Thông Tin Cuộc Họp")

# Chạy Gradio
iface.launch(share=True)