import os
import tempfile
import gradio as gr
from pyannote.audio import Pipeline

# שליפת Hugging Face Token מה-Secret
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
    raise ValueError("HF_TOKEN is missing. Please set it in the Secrets section.")

# טעינת מודל pyannote לזיהוי דוברים
try:
    pipeline = Pipeline.from_pretrained(
        "pyannote/speaker-diarization",
        use_auth_token=hf_token
    )
except Exception as e:
    raise RuntimeError(f"Failed to load the pipeline: {e}")

# פונקציה לזיהוי דוברים
def diarize(audio):
    try:
        # שמירת האודיו לקובץ זמני
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
            temp_audio.write(audio.read())
            temp_audio_path = temp_audio.name

        # עיבוד האודיו עם pyannote
        diarization = pipeline(temp_audio_path)

        # עיבוד התוצאה לזיהוי דוברים
        result = []
        for turn, _, speaker in diarization.itertracks(yield_label=True):
            result.append(f"{speaker}: {turn.start:.1f}s - {turn.end:.1f}s")

        # מחיקת הקובץ הזמני לאחר השימוש
        os.remove(temp_audio_path)

        return "\n".join(result)

    except Exception as e:
        return f"Error during diarization: {e}"

# יצירת ממשק Gradio
interface = gr.Interface(
    fn=diarize,
    inputs=gr.inputs.Audio(source="upload", type="file"),
    outputs="text",
    title="Speaker Diarization",
    description="Upload an audio file (WAV, MP3, etc.) to detect speakers and their timestamps."
)

# הפעלת הממשק
if __name__ == "__main__":
    interface.launch()