import os
import gradio as gr
from transformers import pipeline, MarianMTModel, MarianTokenizer

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

# ⬇️ مدل تولید متن انگلیسی (پیشنهاد موضوع)
text_gen = pipeline(
    "text-generation",
    model="google/gemma-2b-it",  # نسخه سبک و سازگار
    token=token
)

# ⬇️ مدل ترجمه انگلیسی به فارسی
translation_model_name = "Helsinki-NLP/opus-mt-en-fa-opus"
translator_tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
translator_model = MarianMTModel.from_pretrained(translation_model_name)

def translate_to_persian(text):
    inputs = translator_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    translated = translator_model.generate(**inputs)
    return translator_tokenizer.decode(translated[0], skip_special_tokens=True)

def generate_topics(field, major, keywords, audience, level):
    # ساخت پرامپت برای مدل انگلیسی
    prompt = f"""
Suggest 3 thesis topics in the field of {field}, with a specialization in {major}, 
related to the keywords "{keywords}", and targeting the audience "{audience}". 
The academic level is {level}. Just list the topics briefly.
"""

    # تولید متن توسط مدل زبان
    raw_output = text_gen(prompt, max_new_tokens=250)[0]['generated_text']

    # حذف متن prompt تکراری (اگر مدل تکرار کرد)
    if raw_output.startswith(prompt.strip()):
        raw_output = raw_output[len(prompt.strip()):].strip()

    # ترجمه خروجی به فارسی
    translated_output = translate_to_persian(raw_output.strip())

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

    return final_output

# رابط Gradio
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()