import os
import sys
import subprocess
import gradio as gr
import time

# מתקין חבילות נדרשות אם חסרות
print("===== Application Startup at", os.popen('date "+%Y-%m-%d %H:%M:%S"').read().strip(), "=====")

try:
    import transformers
    import huggingface_hub
    import torch
except ImportError:
    print("מתקין חבילות נדרשות...")
    packages = [
        "transformers>=4.38.0",
        "huggingface_hub>=0.20.0",
        "torch>=2.0.0",
        "accelerate>=0.25.0"
    ]
    subprocess.check_call([sys.executable, "-m", "pip", "install"] + packages)
    print("התקנת החבילות הושלמה בהצלחה!")

# מייבא את הספריות אחרי התקנה במידת הצורך
try:
    import transformers
    from huggingface_hub import login
    import torch
    import gradio as gr
    print("יבוא הספריות הצליח!")
except ImportError as e:
    print(f"שגיאה ביבוא הספריות: {str(e)}")
    sys.exit(1)

# רשימת מודלים Qwen3-Coder
MODELS = {
    "Qwen/Qwen1.5-0.5B-Chat": "מודל Qwen קטן - תומך בעברית ומתאים לשיחה ותכנות",
    "Qwen/Qwen1.5-1.8B-Chat": "מודל Qwen בינוני - תומך בשפות רבות כולל עברית",
    "Qwen/Qwen1.5-4B-Chat": "מודל Qwen חזק יותר - תומך בשפות רבות וקידוד",
    "Qwen/Qwen1.5-7B-Chat": "מודל Qwen חזק מאוד - מתאים לשיחה ותכנות מתקדמת",
    "Qwen/Qwen1.5-7B-Coder": "מודל Qwen3-Coder - משופר לתכנות ומשימות קוד"
}

# מילון הגדרות אופטימליות לפי סוגי מודלים
MODEL_CONFIGS = {
    "qwen": {  # למודלי Qwen
        "max_new_tokens": 512,
        "do_sample": True,
        "temperature": 0.7,
        "top_p": 0.95,
        "repetition_penalty": 1.05,
        "no_repeat_ngram_size": 3
    },
    "default": {  # למודלים אחרים
        "max_new_tokens": 256,
        "do_sample": True,
        "temperature": 0.7,
        "top_p": 0.92,
        "repetition_penalty": 1.1,
        "no_repeat_ngram_size": 2
    }
}

# מתחבר לחשבון Hugging Face
token = os.environ.get("HF_TOKEN")
if token:
    login(token)
    print("התחברות ל-Hugging Face הצליחה!")
else:
    print("אזהרה: טוקן HF_TOKEN לא מוגדר בסביבת העבודה. ייתכן שלא תהיה גישה למודלים מוגבלים.")

# הגדרת משתנים גלובליים
default_model_name = "Qwen/Qwen1.5-7B-Coder"
current_model_name = default_model_name
generator = None

# מגדיר את generator כמשתנה גלובלי בכל הפונקציות
def init_globals():
    global generator, current_model_name
    # שם משתנים גלובליים מאותחלים
    
# מאתחל את המשתנים הגלובליים
init_globals()

def load_model(model_name, status_box=None):
    """טעינת מודל ועדכון סטטוס"""
    global generator, current_model_name
    
    # שחרור משאבים של מודל קודם אם קיים
    if generator is not None:
        import gc
        del generator
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    
    current_model_name = model_name
    
    if status_box is not None:
        status_box = gr.Markdown(f"טוען את המודל: {model_name}...")

    try:
        # Qwen מודלים מתאימים לצ'אט
        generator = transformers.pipeline(
            "text-generation",
            model=model_name,
            device_map="auto",
            torch_dtype="auto"
        )
        
        if status_box is not None:
            status_box = gr.Markdown(f"**המודל {model_name} נטען בהצלחה!**")
        
        return f"המודל {model_name} נטען בהצלחה!", status_box
    
    except Exception as e:
        error_msg = f"שגיאה בטעינת המודל {model_name}: {str(e)}"
        print(error_msg)
        
        if status_box is not None:
            status_box = gr.Markdown(f"**שגיאה:** {error_msg}")
        
        return error_msg, status_box

# טעינת מודל ברירת מחדל
print(f"טוען מודל ברירת מחדל {default_model_name}...")
_, _ = load_model(default_model_name)

def ask_model(prompt):
    global generator, current_model_name
    
    if generator is None:
        return "המודל לא נטען בהצלחה. נסה לטעון מודל תחילה."
    
    try:
        # טיפול במודלי Qwen
        outputs = generator(
            prompt,
            max_new_tokens=512,      # אורך תוצאה מקסימלי - הגדלנו עבור Qwen
            do_sample=True,          # דגימה אקראית במקום greedy
            temperature=0.7,         # איזון בין דיוק ויצירתיות
            top_p=0.95,              # נוקליוס דגימה - שומר על מגוון תשובות
            repetition_penalty=1.05, # מונע חזרות
            no_repeat_ngram_size=3,  # מניעת חזרה על ביגרמות (זוגות מילים)
            return_full_text=False   # מחזיר רק את הטקסט החדש שנוצר
        )
        
        # מחזיר את הטקסט שנוצר
        if isinstance(outputs, list) and len(outputs) > 0:
            return outputs[0]["generated_text"]
        else:
            return str(outputs)
    except Exception as e:
        return f"שגיאה בהפעלת המודל: {str(e)}"

# יצירת ממשק משתמש
with gr.Blocks() as demo:
    gr.Markdown("# מערכת צ'אט עם מודלי Qwen3-Coder")
    
    # אזור בחירת מודל וטעינה
    with gr.Row():
        with gr.Column(scale=3):
            # רשימה נפתחת לבחירת מודל
            model_dropdown = gr.Dropdown(
                choices=list(MODELS.keys()),
                value=default_model_name,
                label="בחר מודל",
                info="בחר מודל מהרשימה ולחץ על 'טען מודל'"             
            )
            
            # תיאור המודל הנבחר
            model_description = gr.Markdown(f"**תיאור המודל:** {MODELS[default_model_name]}")
            
        with gr.Column(scale=1):
            # כפתור טעינת מודל
            load_button = gr.Button("טען מודל", variant="primary")
    
    # אזור סטטוס טעינת מודל
    model_status = gr.Markdown(f"**מודל נוכחי:** {default_model_name}")
    
    # פונקציה להצגת תיאור המודל בעת בחירה
    def update_model_description(model_name):
        return f"**תיאור המודל:** {MODELS[model_name]}"
    
    model_dropdown.change(fn=update_model_description, inputs=model_dropdown, outputs=model_description)
    
    # פונקציה לטעינת המודל הנבחר
    def load_selected_model(model_name):
        message, _ = load_model(model_name)
        return message, f"**מודל נוכחי:** {model_name}"
    
    load_button.click(fn=load_selected_model, inputs=model_dropdown, outputs=[gr.Textbox(visible=False), model_status])
    
    gr.Markdown("---")
    gr.Markdown("### שאל את המודל")
    
    # אזור שאלות ותשובות
    with gr.Row():
        input_text = gr.Textbox(
            placeholder="כתוב כאן את השאלה שלך...", 
            lines=3, 
            label="שאלה/בקשה"
        )
        output_text = gr.Textbox(label="תשובת המודל", lines=10)
    
    submit_btn = gr.Button("שלח", variant="primary")
    submit_btn.click(fn=ask_model, inputs=input_text, outputs=output_text)
    
    # דוגמאות לשאלות
    gr.Markdown("### דוגמאות לשאלות:")
    examples = gr.Examples(
        [
            "ספר לי על בינה מלאכותית",
            "מה דעתך על השפה העברית?",
            "כתוב סיפור קצר על ירושלים",
            "כתוב לי קוד בפייתון שמחשב את מספרי פיבונאצ'י",
            "כתוב לי קוד ב-JavaScript שיוצר אפליקציית רשימת משימות",
            "הסבר לי על מערכות מבוססות טרנספורמרים"
        ],
        input_text
    )

    gr.Markdown("""
    ### הערות:
    - בחר מודל מהרשימה ולחץ על 'טען מודל' כדי להחליף את המודל הנוכחי
    - מודלי Qwen מתאימים במיוחד לצ'אט ותכנות
    - מומלץ להתחיל עם המודל הקטן יותר (0.5B) ולהשתמש במודלים הגדולים יותר רק עבור משימות מורכבות
    - הטעינה הראשונית של כל מודל עשויה לקחת מספר שניות עד דקות
    """)

# הפעלת הממשק עם הגדרות נוספות
demo.launch(show_error=True)