import os
import gradio as gr
import json
from tqdm import tqdm
from langchain_community.vectorstores import FAISS
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import google.generativeai as genai

from helpers import (
    list_docx_files, get_splits, get_json_splits_only, get_answer
)

os.environ["USER_AGENT"] = "gradio-hf-space"

if "GOOGLE_API_KEY" not in os.environ:
    os.environ["GOOGLE_API_KEY"] = "AIzaSyDJ4vIKuIBIPNHATLxnoHlagXWbsAz-vRs"
key = "AIzaSyDJ4vIKuIBIPNHATLxnoHlagXWbsAz-vRs"

###

# Cấu hình API key cho Google GenAI
genai.configure(api_key=key)

### loop cho tất cả các file tài liệu và cộng lại thành final_all_split (nhét vào 1 database)

folder_path = 'syllabus_nct_word_format/'
docx_files = list_docx_files(folder_path)

all_splits = []  # Khởi tạo danh sách lưu kết quả
for i, file_path in enumerate(tqdm(docx_files, desc="Đang xử lý", unit="file")):
    output_json_path = f"output_{i}.json"
    splits = get_splits(file_path, output_json_path)
    all_splits += splits

# output_json_path = "output_{i}.json"
# splits = get_splits(docx_files, output_json_path)
# all_splits += splits

FAQ_path = 'syllabus_nct_word_format/FAQ.json'
FAQ_splits = get_json_splits_only(FAQ_path)
all_splits += FAQ_splits

# web_path = 'https'
# web_splits = await get_urls_splits(url='https://nct.neu.edu.vn/')
# all_splits += web_splits

# Lưu vào vectorstore với nhúng từ GenAI (Cần chỉ định model)
embedding = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
vectorstore = FAISS.from_documents(documents=all_splits, embedding=embedding)

###

institutions = ['Tất cả'] + ['Trường Công Nghệ']
categories = ['Tất cả'] + ['Đề án', 'Chương trình đào tạo']

with gr.Blocks() as demo:
    with gr.Row():
        category1 = gr.Dropdown(choices = institutions, label="Trường", value = 'Tất cả')
        category2 = gr.Dropdown(choices = categories, label="Bạn quan tâm tới", value = 'Tất cả')

    chat_interface = gr.ChatInterface(get_answer,
                                      textbox=gr.Textbox(placeholder="Đặt câu hỏi tại đây",
                                                        container=False,
                                                        autoscroll=True,
                                                        scale=7),
                                      type="messages",
                                      # textbox=prompt,
                                      # additional_inputs=[category1, category2]
                                      )

if __name__ == "__main__":
    demo.launch(share=True)