import sys
from unittest.mock import MagicMock

# Comprehensive workaround for missing _lzma in some Python builds
try:
    import lzma
except ImportError:
    mock_lzma = MagicMock()
    mock_lzma.FORMAT_XZ = 1
    mock_lzma.FORMAT_ALONE = 2
    mock_lzma.FORMAT_RAW = 3
    mock_lzma.CHECK_NONE = 0
    mock_lzma.CHECK_CRC32 = 1
    mock_lzma.CHECK_CRC64 = 4
    mock_lzma.CHECK_SHA256 = 10
    sys.modules["_lzma"] = MagicMock()
    sys.modules["lzma"] = mock_lzma

import mlx_whisper
import numpy as np
import sounddevice as sd
import queue
import torch
from silero_vad import load_silero_vad, get_speech_timestamps
from transformers import MarianMTModel, MarianTokenizer

# Parameters
WHISPER_MODEL = "mlx-community/whisper-small.en-mlx"
# List of language pairs (English to ...)
TARGET_LANGS = {
    "Spanish": "Helsinki-NLP/opus-mt-en-es",
    "French":  "Helsinki-NLP/opus-mt-en-fr",
    "Arabic":  "Helsinki-NLP/opus-mt-en-ar"
}

CHANNELS = 1
SAMPLERATE = 16000
BLOCK_SIZE = 512
VAD_THRESHOLD = 0.5 
BUFFER_LIMIT = SAMPLERATE * 30 
MIN_SILENCE_DURATION_MS = 500 

audio_queue = queue.Queue()

def callback(indata, frames, time, status):
    if status:
        print(status, file=sys.stderr)
    audio_queue.put(indata.copy())

def main():
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"Using device: {device}")

    print(f"Loading Whisper model '{WHISPER_MODEL}'...")
    
    # Dictionary to hold models and tokenizers
    translation_engines = {}
    
    for lang_name, model_id in TARGET_LANGS.items():
        print(f"Loading {lang_name} translation model ({model_id})...")
        tokenizer = MarianTokenizer.from_pretrained(model_id)
        model = MarianMTModel.from_pretrained(model_id).to(device)
        translation_engines[lang_name] = (model, tokenizer)

    print("Loading Silero VAD model...")
    vad_model = load_silero_vad()
    
    print("Models loaded.")

    print(f"\nStarting live transcription & multiple translations... (Press Ctrl+C to stop)")
    
    audio_buffer = []
    speech_started = False

    try:
        with sd.InputStream(samplerate=SAMPLERATE, channels=CHANNELS, callback=callback, blocksize=BLOCK_SIZE):
            while True:
                while not audio_queue.empty():
                    data = audio_queue.get()
                    audio_buffer.append(data.flatten())

                if len(audio_buffer) > 0:
                    current_audio = np.concatenate(audio_buffer)
                    audio_tensor = torch.from_numpy(current_audio)
                    
                    speech_timestamps = get_speech_timestamps(
                        audio_tensor, 
                        vad_model, 
                        sampling_rate=SAMPLERATE,
                        threshold=VAD_THRESHOLD,
                        min_silence_duration_ms=MIN_SILENCE_DURATION_MS
                    )

                    if len(speech_timestamps) > 0:
                        speech_started = True
                        last_end = speech_timestamps[-1]['end']
                        buffer_len_samples = len(current_audio)
                        
                        if (buffer_len_samples - last_end) > (SAMPLERATE * MIN_SILENCE_DURATION_MS / 1000) or buffer_len_samples > BUFFER_LIMIT:
                            
                            # 1. Transcribe once
                            result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL)
                            original_text = result['text'].strip()
                            
                            if original_text:
                                # Limit the input length to avoid memory spikes or model glitches
                                if len(original_text) > 250:
                                    original_text = original_text[:247] + "..."
                                
                                print(f"\n[EN]: {original_text}")
                                
                                # 2. Translate to all targets
                                for lang_name, (model, tokenizer) in translation_engines.items():
                                    inputs = tokenizer(original_text, return_tensors="pt", padding=True).to(device)
                                    with torch.no_grad():
                                        # Added max_new_tokens to prevent runaway generation
                                        translated_tokens = model.generate(**inputs, max_new_tokens=150)
                                    translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
                                    print(f"[{lang_name[:2].upper()}]: {translated_text}")
                            
                            audio_buffer = []
                            speech_started = False
                    
                    elif not speech_started and len(current_audio) > SAMPLERATE * 2:
                        audio_buffer = []

    except KeyboardInterrupt:
        print("\nStopped by user.")
    except Exception as e:
        print(f"\nError: {e}")

if __name__ == "__main__":
    main()