import gradio as gr
from PIL import Image
from io import BytesIO
import base64
import traceback

from backend import process_image

def inference(image: Image.Image, gemini_api_key: str):
    """
    פונקציה שמבצעת זיהוי וטשטוש נשים בתמונה,
    ומעדכנת את סרגל ההתקדמות בהתאם.
    """
    if not gemini_api_key.strip():
        raise gr.Error("אנא הכנס מפתח API של Gemini על מנת להמשיך")

    progress = gr.Progress()  # אובייקט לעדכון ההתקדמות

    def progress_callback(fraction, description=""):
        """
        פונקציה פנימית שתיקרא מ-backend בכל שלב.
        fraction - ערך בין 0 ל-1 (לדוגמה 0.3 = 30%)
        description - מלל להסבר השלב
        """
        progress(fraction, desc=description)

    try:
        # כעת נקרא ל-process_image עם אפשרות לעדכן התקדמות
        encoded_image = process_image(image, gemini_api_key, progress_callback=progress_callback)
        decoded_image = Image.open(BytesIO(base64.b64decode(encoded_image)))
        return decoded_image, gr.update(value="", visible=False)  # החזרת תמונה ו-Textbox מוסתר
    except Exception as e:
        # טיפול בשגיאה והחזרת הודעה ב-Textbox
        error_message = f"שגיאה: {type(e).__name__}\n"
        error_message += f"הודעה: {e}\n\n"
        error_message += "Traceback:\n"
        error_message += traceback.format_exc()
        return None, gr.update(value=error_message, visible=True)  # החזרת Textbox גלוי עם שגיאה

title_str = "🤖 זיהוי וטשטוש נשים בתמונה"
description_str = """
<div style='text-align: center; direction: rtl'>
    <p>
        ברוכים הבאים לכלי לזיהוי וטשטוש נשים בתמונה!
        <br>
        העלו תמונה, הזינו את מפתח ה־API של Gemini,
        ולחצו על "הרץ" כדי לנתח את התמונה ולטשטש אוטומטית נשים.
    </p>
    <p>
        שימו לב: נדרש מפתח API תקין של Gemini כדי להשתמש בכלי זה.
        <br>
        הכלי משתמש בטכנולוגיות מתקדמות כמו YOLO, SAM2 ו-Gemini.
    </p>
</div>
"""

# נתיבים לתמונות דוגמה
EXAMPLE_IMAGES = ["example_images/example.jpg", "example_images/example2.jpg", "example_images/example3.jpg"]

with gr.Blocks(
    title=title_str,
    css="style.css"  # קישור לקובץ CSS
) as demo:
    gr.Markdown(f"<h1 style='text-align: center;'>{title_str}</h1>")
    gr.Markdown(description_str)

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="🖼️ בחרו תמונה לניתוח")
            api_key_input = gr.Textbox(
                label="🔑 מפתח API של Gemini",
                placeholder="הכניסו את מפתח ה-API  כאן",
                type="password"
            )
            submit_button = gr.Button("🚀 הרץ", variant="primary")
            gr.Examples(
                examples=EXAMPLE_IMAGES,
                inputs=image_input,
                label="👇 דוגמאות",
                # cache_examples=True # caching examples speeds up start time, but uses more memory
            )
        with gr.Column():
            image_output = gr.Image(type="pil", label="🖼️ תוצאה לאחר טשטוש")
            error_output = gr.Textbox(label="📜 שגיאות", visible=False, lines=5)

    submit_button.click(
        fn=inference,
        inputs=[image_input, api_key_input],
        outputs=[image_output, error_output]
    )

if __name__ == "__main__":
    demo.launch()