import os
import sys
import logging
import warnings
import torch
from typing import Optional, Dict, Any

# הגדרת logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# התעלמות מאזהרות לא רלוונטיות
warnings.filterwarnings("ignore", category=UserWarning)

try:
    import gradio as gr
    from pyannote.audio import Pipeline
except ImportError as e:
    logger.error(f"שגיאה בטעינת הספריות הנדרשות: {str(e)}")
    sys.exit(1)

def initialize_pipeline() -> Optional[Pipeline]:
    """אתחול מודל הדיאריזציה עם טיפול שגיאות"""
    try:
        hf_token = os.getenv('HF_TOKEN')
        if not hf_token:
            raise ValueError("לא נמצא HF_TOKEN במשתני הסביבה")

        pipeline = Pipeline.from_pretrained(
            "pyannote/speaker-diarization@2.1",
            use_auth_token=hf_token
        )

        # העברה ל-GPU אם זמין
        if torch.cuda.is_available():
            pipeline = pipeline.to(torch.device("cuda"))
            logger.info("המודל הועבר ל-GPU")
        else:
            logger.info("רץ על CPU")

        return pipeline

    except Exception as e:
        logger.error(f"שגיאה באתחול המודל: {str(e)}")
        return None

def process_audio(audio_path: str, min_speakers: int = None, max_speakers: int = None) -> Dict[str, Any]:
    """עיבוד קובץ האודיו וזיהוי דוברים"""
    try:
        # בדיקות תקינות
        if not audio_path or not os.path.exists(audio_path):
            return {"error": "לא נבחר קובץ אודיו תקין"}

        file_size_mb = os.path.getsize(audio_path) / (1024 * 1024)
        if file_size_mb > 100:  # הגבלת גודל ל-100MB
            return {"error": f"גודל הקובץ ({file_size_mb:.1f}MB) גדול מדי. המקסימום הוא 100MB"}

        # אתחול המודל
        pipeline = initialize_pipeline()
        if pipeline is None:
            return {"error": "שגיאה באתחול המודל"}

        # הגדרת פרמטרים לזיהוי
        diarization_options = {}
        if min_speakers and min_speakers > 0:
            diarization_options['min_speakers'] = min_speakers
        if max_speakers and max_speakers > 0:
            diarization_options['max_speakers'] = max_speakers

        # זיהוי דוברים
        diarization = pipeline(audio_path, **diarization_options)

        # עיבוד התוצאות
        results = []
        speakers = set()
        total_duration = 0

        for turn, _, speaker in diarization.itertracks(yield_label=True):
            segment = {
                "start": turn.start,
                "end": turn.end,
                "duration": turn.duration,
                "speaker": speaker
            }
            results.append(segment)
            speakers.add(speaker)
            total_duration += turn.duration

        # יצירת פלט מעוצב
        output = [
            "תוצאות זיהוי הדוברים:",
            "=" * 30,
            ""
        ]

        for segment in results:
            output.append(
                f"[{segment['start']:.1f}s -> {segment['end']:.1f}s] "
                f"{segment['speaker']}"
            )

        output.extend([
            "",
            "=" * 30,
            "סיכום:",
            f"• מספר דוברים שזוהו: {len(speakers)}",
            f"• משך כולל: {total_duration:.1f} שניות",
            f"• גודל הקובץ: {file_size_mb:.1f}MB",
        ])

        if min_speakers or max_speakers:
            constraints = []
            if min_speakers:
                constraints.append(f"מינימום {min_speakers} דוברים")
            if max_speakers:
                constraints.append(f"מקסימום {max_speakers} דוברים")
            output.append(f"• הגבלות: {', '.join(constraints)}")

        return {"text": "\n".join(output)}

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

def create_gradio_interface() -> gr.Interface:
    """יצירת ממשק המשתמש"""
    return gr.Interface(
        fn=lambda audio, min_spk, max_spk: process_audio(audio, min_spk, max_spk).get("text", "שגיאה בעיבוד"),
        inputs=[
            gr.Audio(
                label="קובץ אודיו",
                type="filepath"
            ),
            gr.Number(
                label="מינימום דוברים (אופציונלי)",
                value=None,
                minimum=0,
                step=1
            ),
            gr.Number(
                label="מקסימום דוברים (אופציונלי)",
                value=None,
                minimum=0,
                step=1
            )
        ],
        outputs=gr.Textbox(
            label="תוצאות הזיהוי",
            lines=10
        ),
        title="זיהוי דוברים בהקלטות",
        description="""
        העלה קובץ אודיו לזיהוי הדוברים השונים והזמנים שלהם.
        
        הערות:
        • תומך בפורמטים: WAV, MP3, FLAC
        • גודל קובץ מקסימלי: 100MB
        • מומלץ להשתמש בהקלטות באיכות טובה
        • אם ידוע מספר הדוברים, ציון שלו יכול לשפר את הדיוק
        """,
        article="""
        שים לב:
        1. הזיהוי הכי מדויק בהקלטות ללא רעשי רקע
        2. לפעמים המודל יכול לזהות את אותו דובר פעמיים
        3. הזמנים מוצגים בשניות
        """,
        examples=[
            ["example.wav", 2, 4],
            ["interview.mp3", 2, 2]
        ]
    )

if __name__ == "__main__":
    try:
        # מידע על הסביבה
        logger.info(f"Python version: {sys.version}")
        logger.info(f"PyTorch version: {torch.__version__}")
        logger.info(f"Space ID: {os.getenv('SPACE_ID', 'unknown')}")
        logger.info(f"GPU available: {torch.cuda.is_available()}")
        
        if torch.cuda.is_available():
            logger.info(f"GPU model: {torch.cuda.get_device_name(0)}")

        # יצירה והפעלת הממשק
        demo = create_gradio_interface()
        demo.launch(
            share=True
        )
        
    except Exception as e:
        logger.error(f"שגיאה קריטית: {str(e)}")
        sys.exit(1)