import os
import re
import time
import torch
import torchaudio
import warnings
from pydub import AudioSegment

warnings.filterwarnings("ignore")

# =========================================================================
# פתרון חסין לטעינה על CPU
# =========================================================================
original_load = torch.load

def safe_load(*args, **kwargs):
    if 'map_location' not in kwargs:
        kwargs['map_location'] = 'cpu'
    if 'weights_only' in kwargs:
        kwargs['weights_only'] = False
    return original_load(*args, **kwargs)

torch.load = safe_load
# =========================================================================

try:
    from chatterbox.mtl_tts import ChatterboxMultilingualTTS
    from phonikud_onnx import Phonikud
    from phonikud import lexicon
except ImportError:
    print("❌ שגיאה: חסרות ספריות. וודא שאתה ב-venv_fix והרצת pip install.")
    exit()

class TTSManager:
    def __init__(self, reference_audio_path="new_teacher_reference.wav"):
        print("⚙️  טוען מנוע 'המורה למתמטיקה' (ZipVoice Mode)...")
        self.device = "cpu"
        self.ref_audio_path = reference_audio_path
        self.phonikud_model = Phonikud("phonikud-1.0.int8.onnx")
        self.tts_model = ChatterboxMultilingualTTS.from_pretrained(device=self.device)
        print("✅ המנוע מוכן.")

    def _split_to_sentences(self, text):
        return [s.strip() for s in re.split(r'(?<=[.!?]) +', text) if s.strip()]

    def generate_audio(self, text, output_filename="welcome_speech_v2.wav"):
        sentences = self._split_to_sentences(text)
        print(f"📝 מעבד {len(sentences)} משפטים קצרים (למניעת 'נשימות')...")
        
        combined = AudioSegment.empty()
        silence = AudioSegment.silent(duration=250) 

        for i, sentence in enumerate(sentences):
            temp_file = f"temp_part_{i}.wav"
            print(f"🎙️  משפט {i+1}/{len(sentences)}: {sentence[:30]}...")
            
            # שלב 1: ניקוד אוטומטי
            voweled = self.phonikud_model.add_diacritics(sentence)
            voweled = re.sub(fr"[{lexicon.NON_STANDARD_DIAC}]", "", voweled)
            
            # שלב 2: תיקון כפוי לנקבה! מחליף "מורֶה" ב-"מורָה"
            voweled = voweled.replace("מוֹרֶה", "מוֹרָה").replace("מּוֹרֶה", "מּוֹרָה")
            
            # ייצור הקול
            wav = self.tts_model.generate(
                voweled, 
                language_id="he", 
                audio_prompt_path=self.ref_audio_path,
                cfg_weight=0.90 
            )
            torchaudio.save(temp_file, wav, self.tts_model.sr)
            combined += AudioSegment.from_wav(temp_file) + silence
            os.remove(temp_file)

        combined.export(output_filename, format="wav")
        print(f"✨ הושלם! הקובץ מחכה ב: {output_filename}")

if __name__ == "__main__":
    REFERENCE_FILE = "new_teacher_reference.wav"
    
    if not os.path.exists(REFERENCE_FILE):
        print(f"❌ שגיאה: קובץ דגימה '{REFERENCE_FILE}' לא נמצא.")
    else:
        tts = TTSManager(reference_audio_path=REFERENCE_FILE)
        
        # הטקסט פוצל למשפטים קצרים עם נקודות למניעת רעשי נשימה ועומס על המודל
        script = (
            "היי! ברוכים הבאים למורה למתמטיקה. איזה כיף שהצטרפתם אליי! "
            "אני כאן כדי להפוך את המתמטיקה לחוויה ברורה וחזותית. "
            "במסך כאן תמצאו שני כפתורים עיקריים. "
            "לחצו על הכפתור הירוק פתרי לי תרגיל, אם אתם צריכים פתרון מלא ומודרך מאפס. "
            "או על הכפתור הוורוד בדקי לי תשובה, כדי שאעבור על הדרך שאתם עשיתם ואגיד לכם אם צדקתם. "
            "חשוב שתדעו, אני תמיד כאן בשבילכם. "
            "אם סיימתי לפתור ועדיין משהו לא מובן, פשוט תלחצו על הכפתור שאל את המורה. "
            "אני אשמח להסביר הכל שוב עד שהכל יהיה ברור ומסודר בראש. "
            "ועוד טיפ קטן ממני. "
            "אם אתם יושבים ללמוד למבחן, הכי כדאי לעבוד עם מחשב ולא עם הטלפון, כי שם הכל הרבה יותר ברור. "
            "כדי להתחבר, פשוט נכנסים לאתר מהמחשב ומאשרים את הכניסה דרך הטלפון שלכם. "
            "שיהיה המון בהצלחה!"
        )
        
        tts.generate_audio(script, "welcome_speech_v2.wav")