import torch
from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import gradio as gr
import re

# 📌 مدل‌های مختلف
# مدل BERT برای semantic search
bert_model_name = "HooshvareLab/bert-fa-base-uncased"
bert_tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
bert_model = AutoModel.from_pretrained(bert_model_name)

# مدل generative برای تولید پاسخ (می‌توانید از مدل‌های دیگر استفاده کنید)
# اگر مدل generative فارسی در دسترس نداشته باشید، از همین BERT استفاده می‌کنیم
try:
    # مثال: استفاده از مدل generative فارسی
    gen_model_name = "HooshvareLab/bert-fa-base-uncased"  # جایگزین کنید
    gen_tokenizer = AutoTokenizer.from_pretrained(gen_model_name)
    gen_model = AutoModel.from_pretrained(gen_model_name)
    has_generative = False  # فعلاً False چون BERT generative نیست
except:
    has_generative = False

# 📄 دیتاست گسترده‌تر
faq_data = {
    "پایتخت ایران کجاست؟": "تهران پایتخت ایران است.",
    "زبان رسمی ایران چیست؟": "فارسی زبان رسمی ایران است.",
    "واحد پول ایران چیست؟": "ریال واحد پول ایران است.",
    "چه زمانی انتخاب واحد شروع می‌شود؟": "معمولاً انتخاب واحد در پایان شهریور یا بهمن ماه شروع می‌شود.",
    "چه معدلی برای گرفتن 24 واحد لازم است؟": "برای گرفتن 24 واحد حداقل معدل 17 لازم است.",
    "ساعت کاری ادارات چیست؟": "ساعت کاری ادارات معمولاً از 8 صبح تا 4 عصر است.",
    "چگونه می‌توانم درخواست پاسپورت بدهم؟": "برای درخواست پاسپورت باید به اداره گذرنامه مراجعه کنید.",
    "فصل‌های سال در ایران چیست؟": "فصل‌های سال شامل بهار، تابستان، پاییز و زمستان است.",
}

# اضافه کردن دانش عمومی
general_knowledge = {
    "چگونه": "این سوال درباره نحوه انجام کاری است. برای پاسخ دقیق‌تر، لطفاً سوال خود را کامل‌تر بپرسید.",
    "چرا": "این سوال درباره دلیل چیزی است. برای پاسخ بهتر، موضوع مشخصی را بیان کنید.",
    "چیست": "این سوال تعریف چیزی را می‌خواهد. لطفاً موضوع مورد نظر را دقیق‌تر بیان کنید.",
    "کجا": "این سوال درباره مکان است. برای پاسخ دقیق‌تر، موضوع خاصی را مشخص کنید.",
    "کی": "این سوال درباره زمان است. لطفاً موضوع مورد نظر را دقیق‌تر بیان کنید.",
}

# ترکیب دیتاست‌ها
all_data = {**faq_data, **general_knowledge}
questions = list(all_data.keys())
answers = list(all_data.values())

# 📄 تولید embedding
def get_embedding(text):
    inputs = bert_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
    with torch.no_grad():
        outputs = bert_model(**inputs)
        emb = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
    return emb

# محاسبه embedding برای همه سوالات
faq_embeddings = [get_embedding(q) for q in questions]

# 📄 تشخیص نوع سوال
def detect_question_type(question):
    question_lower = question.lower()
    
    # الگوهای سوال
    patterns = {
        'definition': ['چیست', 'چی هست', 'تعریف', 'معنی'],
        'how': ['چگونه', 'چطور', 'چه طور'],
        'why': ['چرا', 'به چه دلیل'],
        'when': ['چه زمانی', 'کی', 'چه وقت'],
        'where': ['کجا', 'در کجا', 'کدام مکان'],
        'who': ['کی', 'چه کسی', 'کدام فرد'],
        'greeting': ['سلام', 'درود', 'صبح بخیر', 'ظهر بخیر', 'عصر بخیر', 'شب بخیر'],
        'thanks': ['ممنون', 'متشکرم', 'سپاس', 'تشکر'],
        'yes_no': ['آیا', 'مگر', 'آیا که']
    }
    
    for q_type, words in patterns.items():
        for word in words:
            if word in question_lower:
                return q_type
    
    return 'general'

# 📄 تولید پاسخ برای سوالات عمومی
def generate_general_answer(question, question_type):
    if question_type == 'greeting':
        return "سلام! چطور می‌تونم کمکتون کنم؟"
    
    elif question_type == 'thanks':
        return "خواهش می‌کنم! اگر سوال دیگری دارید، بپرسید."
    
    elif question_type == 'definition':
        return "برای تعریف دقیق‌تر این موضوع، لطفاً سوال خود را کامل‌تر بپرسید تا بتوانم پاسخ مناسبی ارائه دهم."
    
    elif question_type == 'how':
        return "برای راهنمایی دقیق‌تر درباره نحوه انجام این کار، لطفاً جزئیات بیشتری از سوال خود ارائه دهید."
    
    elif question_type == 'why':
        return "برای توضیح دلایل، لطفاً موضوع مورد نظر را دقیق‌تر مشخص کنید تا بتوانم پاسخ مناسبی ارائه دهم."
    
    elif question_type == 'when':
        return "برای اطلاع از زمان دقیق، لطفاً موضوع خاصی را مشخص کنید تا بتوانم راهنمایی کنم."
    
    elif question_type == 'where':
        return "برای اطلاع از مکان، لطفاً موضوع مورد نظر را دقیق‌تر بیان کنید."
    
    elif question_type == 'yes_no':
        return "برای پاسخ به این سوال، لطفاً موضوع را کامل‌تر مطرح کنید."
    
    else:
        # پاسخ عمومی هوشمند
        return "سوال جالبی پرسیده‌اید. متأسفانه در حال حاضر اطلاعات کاملی در این زمینه ندارم، اما اگر سوال خود را دقیق‌تر مطرح کنید، شاید بتوانم کمک بیشتری کنم."

# 📄 تابع اصلی پاسخ
def answer_question(user_question):
    # حذف فاصله‌های اضافی
    user_question = user_question.strip()
    
    if not user_question:
        return "لطفاً سوال خود را بنویسید."
    
    # تشخیص نوع سوال
    question_type = detect_question_type(user_question)
    
    # جستجوی semantic در دیتاست
    user_emb = get_embedding(user_question)
    sims = [cosine_similarity([user_emb], [emb])[0][0] for emb in faq_embeddings]
    best_idx = int(np.argmax(sims))
    best_score = sims[best_idx]
    
    # تنظیم threshold بر اساس نوع سوال
    if question_type in ['greeting', 'thanks']:
        threshold = 0.3
    else:
        threshold = 0.65
    
    if best_score > threshold:
        return answers[best_idx]
    else:
        # تولید پاسخ برای سوالات خارج از دیتاست
        return generate_general_answer(user_question, question_type)

# 📄 تابع اضافه کردن دانش جدید
def add_knowledge(question, answer):
    if question and answer:
        global questions, answers, faq_embeddings, all_data
        
        # اضافه کردن به دیتاست
        all_data[question] = answer
        questions.append(question)
        answers.append(answer)
        
        # محاسبه embedding جدید
        new_emb = get_embedding(question)
        faq_embeddings.append(new_emb)
        
        return f"دانش جدید اضافه شد: {question} -> {answer}"
    else:
        return "لطفاً هم سوال و هم پاسخ را وارد کنید."

# 📄 رابط Gradio
with gr.Blocks(title="🤖 دستیار فارسی هوشمند") as demo:
    gr.Markdown("## 🤖 دستیار فارسی هوشمند (پاسخ به سوالات داخل و خارج دیتاست)")
    
    with gr.Tab("💬 پرسش و پاسخ"):
        inp = gr.Textbox(label="سؤال خود را بنویسید", placeholder="مثال: سلام، پایتخت ایران کجاست؟")
        out = gr.Textbox(label="پاسخ", lines=3)
        btn = gr.Button("پاسخ بده", variant="primary")
        btn.click(fn=answer_question, inputs=inp, outputs=out)
    
    with gr.Tab("📚 افزودن دانش"):
        gr.Markdown("### افزودن سوال و پاسخ جدید به دیتاست")
        new_q = gr.Textbox(label="سوال جدید")
        new_a = gr.Textbox(label="پاسخ جدید", lines=2)
        add_btn = gr.Button("اضافه کن", variant="secondary")
        add_result = gr.Textbox(label="نتیجه")
        add_btn.click(fn=add_knowledge, inputs=[new_q, new_a], outputs=add_result)
    
    with gr.Tab("ℹ️ راهنما"):
        gr.Markdown("""
        ### نحوه استفاده:
        1. **سوالات معمولی**: مثل "پایتخت ایران کجاست؟"
        2. **احوال‌پرسی**: مثل "سلام" یا "صبح بخیر"
        3. **تشکر**: مثل "ممنون" یا "متشکرم"
        4. **سوالات عمومی**: حتی اگر در دیتاست نباشد، پاسخ مناسب می‌دهد
        
        ### ویژگی‌ها:
        - پاسخ به سوالات داخل دیتاست با دقت بالا
        - پاسخ هوشمند به سوالات خارج دیتاست
        - تشخیص نوع سوال (تعریف، چگونه، چرا، کجا، کی)
        - قابلیت افزودن دانش جدید
        """)

demo.launch()