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 = """<p style='text-align: right; direction: rtl'>
    העלה תמונה, הכנס את מפתח ה־API של Gemini,<br>
    ולחץ על "הרץ" כדי לזהות ולטשטש נשים בתמונה באופן אוטומטי
</p>
"""

# נתיב לתמונת דוגמה
EXAMPLE_IMAGE = "example_images/example.jpg"

demo = gr.Interface(
    fn=inference,
    inputs=[
        gr.Image(type="pil", label="בחר תמונה לניתוח או גרור אותה לכאן"),
        gr.Textbox(
            label="מפתח API של Gemini",
            placeholder="הכנס את מפתח ה-API  כאן",
            type="password"
        )
    ],
    outputs=[
        gr.Image(type="pil", label="תוצאה סופית"),
        gr.Textbox(label="שגיאות", visible=False)  # הוספת רכיב להצגת שגיאות
    ],
    title=title_str,
    description=description_str,
    examples=[
        [EXAMPLE_IMAGE]  # תמונה בלבד, ללא מפתח API
    ],
    flagging_mode="never",
    theme=gr.themes.Default()  # עיצוב קליל לממשק
)


if __name__ == "__main__":
    # ניתן להגדיר share=True אם רוצים לשתף מחוץ לרשת המקומית
    demo.launch()