import gradio as gr
from faster_whisper import WhisperModel
from pydub import AudioSegment
import os
import tempfile
from transformers import pipeline

# הגדרת המודל לתמלול
model = WhisperModel("ivrit-ai/faster-whisper-v2-d4")

# הגדרת pipeline לסיכום
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

def summarize_audio_or_video(file_path):
    try:
        # בדיקה אם הקובץ הוא וידאו והמרת וידאו לאודיו במידת הצורך
        if file_path.endswith((".mp4", ".mov", ".avi", ".mkv")):
            audio_file = convert_video_to_audio(file_path)
        else:
            audio_file = file_path

        # תמלול האודיו
        segments, _ = model.transcribe(audio_file, language="he")
        transcript = " ".join([segment.text for segment in segments])

        # סיכום התמלול
        summary = summarizer(transcript, max_length=50, min_length=25, do_sample=False)[0]["summary_text"]
        
        # מחיקת קובץ האודיו במידת הצורך (אם היה וידאו)
        if audio_file != file_path:
            os.remove(audio_file)

        return summary

    except Exception as e:
        return f"שגיאה בעיבוד הקובץ: {str(e)}"

def convert_video_to_audio(video_file):
    # יצירת קובץ אודיו זמני
    temp_audio = tempfile.mktemp(suffix=".wav")
    video = AudioSegment.from_file(video_file)
    video.export(temp_audio, format="wav")
    return temp_audio

# הגדרת ממשק Gradio
interface = gr.Interface(
    fn=summarize_audio_or_video,
    inputs=gr.Audio(type="filepath"),
    outputs="text",
    title="ממיר אודיו/וידאו לסיכום",
    description="העלה קובץ אודיו או וידאו של מרצה וקבל סיכום קצר של התוכן."
)

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