import os
import re
import gradio as gr
from transformers import pipeline

# گرفتن توکن از متغیر محیطی
token = os.environ.get("HF_TOKEN")

# استفاده از مدل نسبتاً سبک و قابل اجرا در Hugging Face Spaces
pipe = pipeline(
    "text-generation",
    model="google/gemma-2b-it",
    token=token
)

# تابع تشخیص و ترجمه کلمات انگلیسی رایج در حوزه روانشناسی
EN_FA_DICT = {
    "behavior": "رفتار",
    "self-esteem": "عزت نفس",
    "confidence": "اعتماد به نفس",
    "learning": "یادگیری",
    "attachment": "دلبستگی",
    "communication": "ارتباط",
    "stress": "استرس",
    "parent": "والد",
    "trauma": "آسیب روانی"
}

def replace_english_words(text):
    for eng, fa in EN_FA_DICT.items():
        text = re.sub(rf"\b{eng}\b", fa, text, flags=re.IGNORECASE)
    return text

def generate_topics(field, major, keywords, audience, level):
    prompt = f"""
۳ موضوع پایان‌نامه در رشته {field} با گرایش {major} پیشنهاد بده که به کلیدواژه‌های "{keywords}" مربوط باشه و جامعه هدف آن "{audience}" باشد. مقطع: {level}.
موضوعات را فارسی بنویس.
"""
    output = pipe(prompt, max_new_tokens=250)[0]['generated_text']

    # حذف prompt در صورتی که مدل آن را تکرار کرده باشد
    if output.startswith(prompt.strip()):
        output = output[len(prompt.strip()):].strip()

    # حذف فاصله‌های اضافی و جایگزینی کلمات انگلیسی
    output = replace_english_words(output.strip())

    # افزودن پیام پایانی تبلیغاتی
    output += "\n\nبرای مشاوره و راهنمایی تخصصی با گروه مشاوره کاسپین تماس بگیرید:\n02188252497"

    return output

iface = gr.Interface(
    fn=generate_topics,
    inputs=[
        gr.Textbox(label="رشته"),
        gr.Textbox(label="گرایش"),
        gr.Textbox(label="کلیدواژه‌ها"),
        gr.Textbox(label="جامعه هدف"),
        gr.Dropdown(choices=["کارشناسی ارشد", "دکتری"], label="مقطع")
    ],
    outputs="text",
    title="پیشنهادگر هوشمند موضوع پایان‌نامه کاسپین 🎓"
)

iface.launch()