import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from deep_translator import GoogleTranslator
import torch

# انتخاب مدل پایدار
model_id = "google/gemma-7b-it"

# بارگذاری مدل و توکنایزر
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)
model.eval()

# تابع اصلی پیشنهاد موضوع پایان‌نامه
def generate_topics(field, major, keywords, audience, level):
    prompt = f"""<bos>[INST]Suggest 3 academic thesis topics based on the following:
Field: {field}
Specialization: {major}
Keywords: {keywords}
Target Audience: {audience}
Level: {level}[/INST]"""

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=256)
    english_output = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

    if not english_output:
        return "<div dir='rtl' style='color: red;'>❌ مشکلی در تولید موضوعات رخ داد. لطفاً دوباره تلاش کنید.</div>"

    translated_output = GoogleTranslator(source='en', target='fa').translate(english_output).strip()
    translated_output_html = translated_output.replace("\n", "<br>")

    html_output = (
        "<div dir='rtl' style='text-align: right; font-family: Tahoma, sans-serif; "
        "font-size: 16px; line-height: 1.8; min-height: 250px; max-height: 400px; "
        "overflow-y: auto; border: 1px solid #ccc; padding: 10px; background-color: #fdfdfd;'>"
        f"{translated_output_html}"
        "<br><br>📢 برای مشاوره و راهنمایی تخصصی با گروه مشاوره کاسپین تماس بگیرید:<br>"
        "<strong>021-88252497</strong></div>"
    )
    return html_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=gr.HTML(label="موضوعات پیشنهادی"),
    title="🎓 پیشنهادگر موضوع پایان‌نامه کاسپین",
    theme="default"
)

iface.launch()