import gradio as gr
import pandas as pd
import torch
import time
import os
from datetime import datetime

from transformers import AutoModelForCausalLM, AutoTokenizer
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer
from datasets import load_dataset, Dataset

if torch.cuda.is_available():
    print("✅ GPU đã sẵn sàng:", torch.cuda.get_device_name(0))
else:
    print("❌ Không có GPU khả dụng, ứng dụng có thể gặp lỗi")

# ---------------------------
# 1. Load dữ liệu huấn luyện mẫu
# ---------------------------
df = pd.read_csv("819QA_data_for_training.csv")
df["question"] = df["question"].astype(str).str.strip()
df["answer"] = df["answer"].astype(str).str.strip()

# ---------------------------
# 2. Tạo TF-IDF vectorizer
# ---------------------------
vectorizer = TfidfVectorizer()
tfidf_matrix = vectorizer.fit_transform(df["question"])

# ---------------------------
# 3. Load mô hình và tokenizer từ Hugging Face Model Hub
# ---------------------------
model_path = "vulinhuit/chatbot_tvts_llama3.2-3B"
token = os.environ.get("HF_HUB_TOKEN")

tokenizer = AutoTokenizer.from_pretrained(model_path, use_auth_token=token)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="auto",
    use_auth_token=token
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ---------------------------
# 4. Hàm sinh phản hồi từ mô hình generative
# ---------------------------
def chat_with_model(prompt, max_length=256):
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding="max_length", max_length=256)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=50,
            pad_token_id=tokenizer.eos_token_id,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.2,
            do_sample=True
        )
    
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

# ---------------------------
# 5. Hàm lưu toàn bộ hội thoại lên Hugging Face Dataset khi kết thúc
# ---------------------------
def save_conversation(history):
    if not history:
        return "Không có hội thoại để lưu."
    
    new_data = {
        "timestamp": [],
        "user_input": [],
        "response": []
    }
    for user_msg, bot_msg in history:
        new_data["timestamp"].append(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
        new_data["user_input"].append(user_msg)
        new_data["response"].append(bot_msg)
    
    dataset_name = "vulinhuit/chatbot_history_llama"
    try:
        dataset = load_dataset(dataset_name)
        df_existing = dataset["train"].to_pandas()
        df_new = pd.DataFrame(new_data)
        df_combined = df_existing.append(df_new, ignore_index=True)
        new_dataset = Dataset.from_pandas(df_combined)
    except Exception as e:
        new_dataset = Dataset.from_dict(new_data)
    
    new_dataset.push_to_hub(dataset_name, split="train")
    return "Hội thoại đã được lưu lên Hugging Face Dataset."

# ---------------------------
# 6. Hàm xử lý tin nhắn của người dùng (bao gồm thông tin thời gian phản hồi)
# ---------------------------
def process_user_message(user_input, history):
    if not user_input.strip():
        return history, "Bạn chưa nhập câu hỏi nào!"
    
    start_time = time.time()
    gen_response = chat_with_model(user_input, max_length=256)
    gen_time = time.time() - start_time

    # Tính TF-IDF similarity giữa câu hỏi của người dùng và dữ liệu mẫu
    user_tfidf = vectorizer.transform([user_input])
    similarities = cosine_similarity(user_tfidf, tfidf_matrix)
    best_match_idx = similarities.argmax()
    best_answer = df.iloc[best_match_idx]["answer"]
    sim_score = similarities[0, best_match_idx]

    # Quyết định phản hồi cuối cùng
    final_response = best_answer if sim_score > 0.5 else gen_response

    # Tích hợp thông tin thời gian phản hồi vào câu trả lời
    final_message = f"🤖 ({gen_time:.2f}s): {final_response}"

    history.append((user_input, final_message))
    return history, final_message

    # Hàm thông báo chức năng lưu đang được xây dựng
def under_construction(chat_history):
    return "Chức năng đang trong quá trình xây dựng."

# ---------------------------
# 7. Tạo giao diện sử dụng gr.Blocks với title và mô tả
# ---------------------------
with gr.Blocks(title="Chatbot Tư vấn tuyển sinh HIAST") as demo:
    gr.Markdown("### Chatbot Tư vấn tuyển sinh HIAST\nChatbot sử dụng mô hình LLaMA3 kết hợp phản hồi sinh và dữ liệu tham chiếu.")
    chatbot = gr.Chatbot(label="Hội thoại")
    user_input = gr.Textbox(label="Nhập câu hỏi tại đây:", placeholder="Nhập câu hỏi của bạn...")
    send_btn = gr.Button("Gửi cho chatbot trả lời")
    end_btn = gr.Button("Kết thúc hội thoại và lưu thông tin lịch sử hội thoại!")
    notify = gr.Textbox(label="Thông báo", interactive=False)
    
    def user_send(user_input, chat_history):
        new_history, reply = process_user_message(user_input, chat_history)
        return new_history, "", new_history
    
    send_btn.click(user_send, inputs=[user_input, chatbot], outputs=[chatbot, user_input, chatbot])
    user_input.submit(user_send, inputs=[user_input, chatbot], outputs=[chatbot, user_input, chatbot])
    
    end_btn.click(under_construction, inputs=[chatbot], outputs=notify)

demo.launch()