import librosa
import torch
import numpy as np
import langid  # Language detection library
from transformers import Wav2Vec2ForCTC, AutoProcessor

ASR_SAMPLING_RATE = 16_000
MODEL_ID = "facebook/mms-1b-all"
# openai/whisper-large-v3-turbo
#ASR_SAMPLING_RATE = 16_000
#MODEL_ID = "openai/whisper-large-v3-turbo"

# Load MMS Model
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID)
model.eval()

def detect_language(text):
    """Detects language using langid (fast & lightweight)."""
    lang, _ = langid.classify(text)
    return lang if lang in ["en", "sw"] else "en"  # Default to English

def transcribe_auto(audio_data=None):
    if not audio_data:
        return "<<ERROR: Empty Audio Input>>"
    
    # Process Microphone Input
    if isinstance(audio_data, tuple):
        sr, audio_samples = audio_data
        audio_samples = (audio_samples / 32768.0).astype(np.float32)
        if sr != ASR_SAMPLING_RATE:
            audio_samples = librosa.resample(audio_samples, orig_sr=sr, target_sr=ASR_SAMPLING_RATE)
    
    # Process File Upload Input
    else:
        if not isinstance(audio_data, str):
            return "<<ERROR: Invalid Audio Input>>"
        audio_samples = librosa.load(audio_data, sr=ASR_SAMPLING_RATE, mono=True)[0]

    inputs = processor(audio_samples, sampling_rate=ASR_SAMPLING_RATE, return_tensors="pt")

    # **Step 1: Transcribe without Language Detection**
    with torch.no_grad():
        outputs = model(**inputs).logits
        ids = torch.argmax(outputs, dim=-1)[0]
        raw_transcription = processor.decode(ids)

    # **Step 2: Detect Language from Transcription**
    detected_lang = detect_language(raw_transcription)
    lang_code = "eng" if detected_lang == "en" else "swh"

    # **Step 3: Reload Model with Correct Adapter**
    processor.tokenizer.set_target_lang(lang_code)
    model.load_adapter(lang_code)

    # **Step 4: Transcribe Again with Correct Adapter**
    with torch.no_grad():
        outputs = model(**inputs).logits
        ids = torch.argmax(outputs, dim=-1)[0]
        final_transcription = processor.decode(ids)

    return f"{final_transcription}"