import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import json
import os
# 🔹 مدل embedding (برای تشخیص شباهت)
embedding_model_name = "HooshvareLab/bert-fa-base-uncased"
embedding_tokenizer = AutoTokenizer.from_pretrained(embedding_model_name)
embedding_model = AutoModel.from_pretrained(embedding_model_name)

# 🔹 مدل تولید (برای پاسخ جدید)
#gen_model_name = "HooshvareLab/PersianMind"
gen_model_name = "universitytehran/PersianMind-v1.0"
gen_tokenizer = AutoTokenizer.from_pretrained(gen_model_name)
gen_model = AutoModelForCausalLM.from_pretrained(gen_model_name)

# 🔹 مسیر فایل دیتابیس
DATA_FILE = "faq_data.json"
ADMIN_PASSWORD = "admin123"  # رمز عبور ادمین

# بارگذاری سوالات از فایل یا مقدار پیش‌فرض
def load_faq_data():
    if os.path.exists(DATA_FILE):
        with open(DATA_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    else:
        return {
            "زمان انتخاب واحد چه موقع است؟": "معمولاً پایان شهریور و بهمن است.",
            "چه زمانی می‌توان حذف و اضافه انجام داد؟": "حدود یک هفته پس از شروع ترم تحصیلی است.",
            "چه معدلی برای انتخاب ۲۴ واحد لازم است؟": "حداقل معدل 17 نیاز است.",
            "تا چه زمانی امکان حذف اضطراری وجود دارد؟": "تا هفته هشتم ترم مجاز است.",
            "چگونه می‌توان مهمان شد؟": "با موافقت دانشگاه مبدا و مقصد انجام می‌شود.",
        }

# ذخیره‌سازی در فایل
def save_faq_data():
    with open(DATA_FILE, "w", encoding="utf-8") as f:
        json.dump(faq_dict, f, ensure_ascii=False, indent=2)

# پایگاه دانش و embedding‌ها
faq_dict = load_faq_data()
faq_questions = list(faq_dict.keys())

# تابع تولید embedding
def get_embedding(text):
    inputs = embedding_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = embedding_model(**inputs)
    return outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()

# ساخت embedding اولیه
faq_embeddings = [get_embedding(q) for q in faq_questions]

# تابع تولید پاسخ با PersianMind
def generate_with_persianmind(prompt):
    inputs = gen_tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512)
    with torch.no_grad():
        output_ids = gen_model.generate(
            inputs.input_ids,
            max_length=200,
            do_sample=True,
            top_p=0.9,
            temperature=0.8,
            pad_token_id=gen_tokenizer.eos_token_id
        )
    answer = gen_tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return answer

# پاسخ‌دهی اصلی
def student_bot(user_question):
    try:
        user_emb = get_embedding(user_question)
        sims = [cosine_similarity([user_emb], [faq_emb])[0][0] for faq_emb in faq_embeddings]
        best_idx = int(np.argmax(sims))
        best_score = sims[best_idx]

        if best_score > 0.6:
            return faq_dict[faq_questions[best_idx]]
        else:
            return generate_with_persianmind(user_question)
    except Exception as e:
        return f"❗️خطا: {str(e)}"

# افزودن سؤال جدید با رمز ادمین
def add_faq(new_q, new_a, password):
    if password != ADMIN_PASSWORD:
        return "⛔️ دسترسی فقط برای ادمین مجاز است."

    if new_q.strip() == "" or new_a.strip() == "":
        return "⚠️ لطفاً سؤال و پاسخ را وارد کنید."

    if new_q in faq_dict:
        return "⚠️ این سؤال قبلاً ثبت شده است."

    faq_dict[new_q] = new_a
    faq_questions.append(new_q)
    faq_embeddings.append(get_embedding(new_q))
    save_faq_data()
    return "✅ سؤال جدید با موفقیت افزوده شد و ذخیره شد."

# رابط گرافیکی Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 ایجنت راهنمای دانشجویان")

    with gr.Tab("🟢 پرسیدن سؤال"):
        user_input = gr.Textbox(label="سؤال شما")
        response = gr.Textbox(label="پاسخ")
        ask_btn = gr.Button("پاسخ را دریافت کن")
        ask_btn.click(fn=student_bot, inputs=user_input, outputs=response)

    with gr.Tab("🔐 افزودن سؤال (فقط ادمین)"):
        new_q = gr.Textbox(label="سؤال جدید")
        new_a = gr.Textbox(label="پاسخ مربوط")
        password = gr.Textbox(label="رمز عبور", type="password")
        result = gr.Textbox(label="وضعیت")
        add_btn = gr.Button("افزودن سؤال")
        add_btn.click(fn=add_faq, inputs=[new_q, new_a, password], outputs=result)

demo.launch()