import gradio as gr
from transformers import pipeline
import time

# بارگیری مدل فارسی
fa_model = pipeline(
    "text-generation",
    model="HooshvareLab/gpt2-fa",
    tokenizer="HooshvareLab/gpt2-fa"
)

def generate_response(message, history):
    # مرحله ۱: نمایش وضعیت پردازش
    history.append([message, None])
    yield history
    
    # مرحله ۲: تولید پاسخ واقعی
    response = fa_model(
        message,
        max_length=1000,
        num_return_sequences=1,
        pad_token_id=50256
    )[0]['generated_text']
    
    # مرحله ۳: بروزرسانی تاریخچه با پاسخ
    history[-1][1] = response
    yield history

def handle_upload(files):
    return f"✅ {len(files)} فایل آپلود شد!"

# طراحی رابط کاربری
css = """
/* CSS قبلی بدون تغییر */
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
    # هدر
    gr.HTML("""
    <div class="header">
        <h1>🧠 هوش مصنوعی فارسی پیشرفته</h1>
        <p>نسخه مبتنی بر GPT-2 فارسی</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(label="مکالمه", elem_classes="chat-container")
            msg = gr.Textbox(
                label="پیام خود را وارد کنید",
                placeholder="مثال: فلسفه زندگی چیست؟",
                lines=3
            )
            
            with gr.Row():
                submit_btn = gr.Button("ارسال پیام", elem_classes="btn-ai")
                upload_btn = gr.UploadButton(
                    "آپلود فایل",
                    file_types=["image", "pdf", "text"],
                    file_count="multiple"
                )

        with gr.Column(scale=1):
            gr.Markdown("## ⚙️ تنظیمات")
            temperature = gr.Slider(0, 1, value=0.85, label="سطح خلاقیت")
            max_length = gr.Slider(100, 2000, value=500, label="حداکثر طول پاسخ")

    # مدیریت رویدادها
    msg.submit(
        fn=generate_response,
        inputs=[msg, chatbot],
        outputs=[chatbot]
    )
    
    submit_btn.click(
        fn=generate_response,
        inputs=[msg, chatbot],
        outputs=[chatbot]
    )
    
    upload_btn.upload(
        fn=handle_upload,
        inputs=upload_btn,
        outputs=chatbot
    )

if __name__ == "__main__":
    app.launch(server_port=7860)