src/input_agent.py

import whisper
import os
import torch
import warnings
import gc  # Garbage Collector for memory management
import tempfile
import re
import time
from pydub import AudioSegment # 🟢 NEW: Add pydub

# Suppress FP16 warnings on CPU
warnings.filterwarnings("ignore", message="FP16 is not supported on CPU; using FP32 instead")

class AgentInput:
    TURBO_MODEL_ID = "openai/whisper-large-v3-turbo"
    DISTIL_MODEL_ID = "distil-whisper/distil-large-v3"
    MMS_MODEL_ID = "facebook/mms-1b-all"
    QWEN_ASR_MODEL_ID = "Qwen/Qwen3-ASR-1.7B"

    def __init__(self, device="cpu"):
        print(f"👂 Agent 1 (Input) Online: Preparing Whisper on {device}...")
        self.device = device
        self.loaded_models = {}
        self.loaded_pipelines = {}
        self.turbo_disabled = False
        self.turbo_slow_threshold_s = float(os.environ.get("ASR_TURBO_SLOW_THRESHOLD_S", "18"))
        
        # USE 'tiny' FOR CLOUD TO PREVENT EXIT CODE 137 (OOM)
        # Use 'base' only if you are running on a machine with 16GB+ RAM
        self.model_name = "tiny" 
        
        try:
            # Load the model and immediately collect garbage to free RAM
            self.model = whisper.load_model(self.model_name, device=device)
            self.loaded_models[self.model_name] = self.model
            gc.collect() 
            print(f"✅ Whisper '{self.model_name}' model loaded. RAM optimized.")
        except Exception as e:
            print(f"⚠️ Load failed: {e}. Attempting emergency load...")
            # Emergency fallback to tiny if not already tried
            self.model = whisper.load_model("tiny", device=device)
            self.loaded_models["tiny"] = self.model
            self.model_name = "tiny"

    def _audio_duration_s(self, audio_path):
        try:
            return float(AudioSegment.from_file(audio_path).duration_seconds)
        except Exception:
            return None

    def _hint_text(self, language=None, dialect_hint=""):
        return f"{language or ''} {dialect_hint or ''}".strip().lower()

    def _is_zero_gpu_available(self):
        return torch.cuda.is_available()

    def _is_underrepresented_language(self, language=None, dialect_hint=""):
        hint = self._hint_text(language, dialect_hint)
        common_tokens = {
            "english", " en", "en ", "korean", " ko", "ko ", "arabic", " ar", "ar ",
            "chinese", "mandarin", "cantonese", " zh", "zh ", "french", " fr", "fr ",
            "spanish", " es", "es ", "german", "deutsch", " de", "de ", "italian",
            " it", "it ", "portuguese", " pt", "pt ", "hindi", " hi", "hi ",
            "japanese", " ja", "ja ", "russian", " ru", "ru "
        }
        if any(token in f" {hint} " for token in common_tokens):
            return False
        return bool(hint)

    def _expected_script(self, language=None, dialect_hint=""):
        lang = str(language or "").strip().lower()
        padded_lang = f" {lang} "
        if lang in {"en", "eng", "fr", "fra", "es", "spa", "de", "deu", "it", "ita", "pt", "por", "pcm"}:
            return "latin"
        if lang in {"ko", "kor"}:
            return "hangul"
        if lang in {"ar", "ara", "ur", "urd", "fa", "fas", "prs", "ps", "pus"}:
            return "arabic"
        if lang in {"zh", "cmn", "yue", "zho"}:
            return "cjk"
        if lang in {"ja", "jpn"}:
            return "japanese"
        if lang in {"hi", "hin", "mr", "mar", "ne", "nep"}:
            return "devanagari"
        if lang in {"th", "tha"}:
            return "thai"

        hint = self._hint_text(language, dialect_hint)
        padded_hint = f" {hint} "
        if any(token in padded_hint for token in ["korean", "hangul", "satoori", "jeju", "gyeongsang", "chungcheong", "jeolla", "busan", " ko "]):
            return "hangul"
        if any(token in padded_hint for token in ["arabic", "urdu", "persian", "farsi", "dari", "pashto", " ar "]):
            return "arabic"
        if any(token in padded_hint for token in ["chinese", "mandarin", "cantonese", "yue", " zh "]):
            return "cjk"
        if any(token in padded_hint for token in ["japanese", " ja "]):
            return "japanese"
        if any(token in padded_hint for token in ["hindi", "marathi", "nepali", " hi "]):
            return "devanagari"
        if "thai" in padded_hint:
            return "thai"
        if any(token in padded_hint for token in ["english", "french", "spanish", "german", "italian", "portuguese", "pidgin", "patois", " en ", " fr ", " es "]):
            return "latin"
        if any(token in padded_lang for token in [" en ", " eng ", " fr ", " fra ", " es ", " spa "]):
            return "latin"
        return None

    def _looks_wrong_script(self, text, language=None, dialect_hint=""):
        text = str(text or "").strip()
        if not text:
            return True
        script = self._expected_script(language, dialect_hint)
        if not script:
            return False
        patterns = {
            "hangul": r"[\uac00-\ud7a3]",
            "arabic": r"[\u0600-\u06ff]",
            "cjk": r"[\u3400-\u9fff]",
            "japanese": r"[\u3040-\u30ff\u3400-\u9fff]",
            "devanagari": r"[\u0900-\u097f]",
            "thai": r"[\u0e00-\u0e7f]",
            "latin": r"[A-Za-z]",
        }
        expected_hits = len(re.findall(patterns[script], text))
        if script == "latin":
            other_hits = len(re.findall(r"[\u0600-\u06ff\u0900-\u097f\u0e00-\u0e7f\u3040-\u30ff\u3400-\u9fff\uac00-\ud7a3]", text))
            return expected_hits == 0 or other_hits > expected_hits
        return expected_hits == 0

    def _mms_language_code(self, language=None, dialect_hint=""):
        hint = self._hint_text(language, dialect_hint)
        mappings = [
            (["nigerian pidgin", "pidgin"], "pcm"),
            (["english", " en "], "eng"),
            (["korean", " ko "], "kor"),
            (["arabic", " ar "], "ara"),
            (["mandarin", "chinese", " zh "], "cmn"),
            (["cantonese", "yue"], "yue"),
            (["french", " fr "], "fra"),
            (["spanish", " es "], "spa"),
            (["german", "deutsch", " de "], "deu"),
            (["italian", " it "], "ita"),
            (["portuguese", " pt "], "por"),
            (["hindi", " hi "], "hin"),
            (["japanese", " ja "], "jpn"),
            (["russian", " ru "], "rus"),
            (["thai"], "tha"),
            (["vietnamese"], "vie"),
            (["swahili"], "swh"),
            (["yoruba"], "yor"),
            (["igbo"], "ibo"),
            (["hausa"], "hau"),
        ]
        padded = f" {hint} "
        for tokens, code in mappings:
            if any(token in padded for token in tokens):
                return code
        return "eng"

    def _normalize_manual_route(self, model_choice):
        choice = str(model_choice or "auto").strip().lower()
        aliases = {
            "tiny": {"engine": "whisper", "model": "tiny", "label": "Whisper tiny.en (Edge fastest)"},
            "whisper tiny": {"engine": "whisper", "model": "tiny", "label": "Whisper tiny.en (Edge fastest)"},
            "base": {"engine": "whisper", "model": "base", "label": "Whisper base (Edge balanced)"},
            "whisper base": {"engine": "whisper", "model": "base", "label": "Whisper base (Edge balanced)"},
            "small": {"engine": "whisper", "model": "small", "label": "Whisper small (Fallback)"},
            "whisper-small": {"engine": "whisper", "model": "small", "label": "Whisper small (Fallback)"},
            "whisper small": {"engine": "whisper", "model": "small", "label": "Whisper small (Fallback)"},
            "turbo": {"engine": "hf-whisper", "model": self.TURBO_MODEL_ID, "label": "Whisper large-v3 turbo (HF Pro)"},
            "large-v3-turbo": {"engine": "hf-whisper", "model": self.TURBO_MODEL_ID, "label": "Whisper large-v3 turbo (HF Pro)"},
            "whisper-large-v3-turbo": {"engine": "hf-whisper", "model": self.TURBO_MODEL_ID, "label": "Whisper large-v3 turbo (HF Pro)"},
            self.TURBO_MODEL_ID: {"engine": "hf-whisper", "model": self.TURBO_MODEL_ID, "label": "Whisper large-v3 turbo (HF Pro)"},
            "distil": {"engine": "hf-whisper", "model": self.DISTIL_MODEL_ID, "label": "Distil-Whisper (English fast)"},
            "distil-large-v3": {"engine": "hf-whisper", "model": self.DISTIL_MODEL_ID, "label": "Distil-Whisper (English fast)"},
            self.DISTIL_MODEL_ID: {"engine": "hf-whisper", "model": self.DISTIL_MODEL_ID, "label": "Distil-Whisper (English fast)"},
            "qwen": {"engine": "hf-whisper", "model": self.QWEN_ASR_MODEL_ID, "label": "Qwen3-ASR 1.7B (Multilingual)"},
            "qwen3-asr": {"engine": "hf-whisper", "model": self.QWEN_ASR_MODEL_ID, "label": "Qwen3-ASR 1.7B (Multilingual)"},
            "qwen/qwen3-asr-1.7b": {"engine": "hf-whisper", "model": self.QWEN_ASR_MODEL_ID, "label": "Qwen3-ASR 1.7B (Multilingual)"},
            self.QWEN_ASR_MODEL_ID.lower(): {"engine": "hf-whisper", "model": self.QWEN_ASR_MODEL_ID, "label": "Qwen3-ASR 1.7B (Multilingual)"},
            "mms": {"engine": "mms", "model": self.MMS_MODEL_ID, "label": "MMS 1B (Low-resource)"},
            self.MMS_MODEL_ID: {"engine": "mms", "model": self.MMS_MODEL_ID, "label": "MMS 1B (Low-resource)"},
        }
        return aliases.get(choice)

    def _auto_routes(self, language=None, dialect_hint="", audio_path=None):
        base = {"engine": "whisper", "model": "base", "label": "Whisper base (Edge balanced)"}
        tiny = {"engine": "whisper", "model": "tiny", "label": "Whisper tiny.en (Edge fastest)"}
        small = {"engine": "whisper", "model": "small", "label": "Whisper small (Fallback)"}
        mms = {"engine": "mms", "model": self.MMS_MODEL_ID, "label": "MMS 1B (Low-resource)"}
        turbo = {"engine": "hf-whisper", "model": self.TURBO_MODEL_ID, "label": "Whisper large-v3 turbo (HF Pro)"}

        duration_s = self._audio_duration_s(audio_path) if audio_path else None
        if self._is_zero_gpu_available() and not self.turbo_disabled and not (duration_s and duration_s > 18):
            return [turbo, base, tiny]
        if self._expected_script(language, dialect_hint) == "latin":
            return [base, tiny]
        if self._is_underrepresented_language(language, dialect_hint):
            return [base, mms, small]
        return [small]

    def resolve_model_name(self, model_choice="auto", language=None, dialect_hint="", audio_path=None):
        manual = self._normalize_manual_route(model_choice)
        if manual:
            return manual["model"]
        return self._auto_routes(language=language, dialect_hint=dialect_hint, audio_path=audio_path)[0]["model"]

    def _get_model(self, model_name):
        model_name = model_name if model_name in {"tiny", "base", "small"} else "small"
        if model_name in self.loaded_models:
            self.model_name = model_name
            self.model = self.loaded_models[model_name]
            return self.model

        try:
            print(f"🔄 Loading Whisper '{model_name}' on demand...")
            self.loaded_models[model_name] = whisper.load_model(model_name, device=self.device)
            self.model_name = model_name
            self.model = self.loaded_models[model_name]
            gc.collect()
            return self.model
        except Exception as e:
            print(f"⚠️ Whisper '{model_name}' unavailable ({e}); falling back to tiny.")
            self.model_name = "tiny"
            self.model = self.loaded_models.get("tiny") or whisper.load_model("tiny", device=self.device)
            self.loaded_models["tiny"] = self.model
            gc.collect()
            return self.model

    def _get_hf_pipeline(self, model_id):
        if model_id in self.loaded_pipelines:
            return self.loaded_pipelines[model_id]
        try:
            from transformers import pipeline
            device_index = 0 if torch.cuda.is_available() else -1
            kwargs = {"model": model_id, "device": device_index}
            if torch.cuda.is_available():
                kwargs["torch_dtype"] = torch.float16
            print(f"🔄 Loading ASR pipeline '{model_id}'...")
            self.loaded_pipelines[model_id] = pipeline("automatic-speech-recognition", **kwargs)
            gc.collect()
            return self.loaded_pipelines[model_id]
        except Exception as e:
            raise RuntimeError(f"ASR pipeline '{model_id}' unavailable: {e}") from e

    def _get_mms_pipeline(self, language=None, dialect_hint=""):
        lang_code = self._mms_language_code(language, dialect_hint)
        cache_key = f"{self.MMS_MODEL_ID}:{lang_code}"
        if cache_key in self.loaded_pipelines:
            return self.loaded_pipelines[cache_key]
        try:
            from transformers import AutoProcessor, Wav2Vec2ForCTC, pipeline
            device_index = 0 if torch.cuda.is_available() else -1
            processor = AutoProcessor.from_pretrained(self.MMS_MODEL_ID)
            model = Wav2Vec2ForCTC.from_pretrained(self.MMS_MODEL_ID)
            try:
                processor.tokenizer.set_target_lang(lang_code)
                model.load_adapter(lang_code)
            except Exception as e:
                print(f"⚠️ MMS adapter '{lang_code}' unavailable ({e}); trying English adapter.")
                lang_code = "eng"
                processor.tokenizer.set_target_lang(lang_code)
                model.load_adapter(lang_code)
                cache_key = f"{self.MMS_MODEL_ID}:{lang_code}"
            if torch.cuda.is_available():
                model = model.to("cuda")
            self.loaded_pipelines[cache_key] = pipeline(
                "automatic-speech-recognition",
                model=model,
                tokenizer=processor.tokenizer,
                feature_extractor=processor.feature_extractor,
                device=device_index,
            )
            gc.collect()
            return self.loaded_pipelines[cache_key]
        except Exception as e:
            raise RuntimeError(f"MMS ASR unavailable: {e}") from e

    def _transcribe_with_route(self, route, clean_path, language=None, dialect_hint=""):
        start = time.perf_counter()
        if route["engine"] == "whisper":
            model = self._get_model(route["model"])
            result = model.transcribe(clean_path, language=language, fp16=False)
        elif route["engine"] == "mms":
            pipe = self._get_mms_pipeline(language, dialect_hint)
            result = pipe(clean_path)
        else:
            pipe = self._get_hf_pipeline(route["model"])
            generate_kwargs = {"task": "transcribe"}
            if language:
                generate_kwargs["language"] = language
            result = pipe(clean_path, generate_kwargs=generate_kwargs)

        elapsed_s = time.perf_counter() - start
        if isinstance(result, dict):
            text = str(result.get("text", "")).strip()
        else:
            text = str(result or "").strip()
        if route["model"] == self.TURBO_MODEL_ID and elapsed_s > self.turbo_slow_threshold_s:
            self.turbo_disabled = True
            print(f"⚠️ Turbo route was slow ({elapsed_s:.2f}s); future Auto calls will use Whisper small.")
        return text, elapsed_s

    # 🟢 NEW HELPER: Sanitizes corrupted browser audio
    def _sanitize_audio(self, audio_path):
        try:
            # Try to load it regardless of format
            audio = AudioSegment.from_file(audio_path)
            # Export it to a clean, standard WAV in a temp file
            temp_path = os.path.join(tempfile.gettempdir(), f"clean_audio_{os.path.basename(audio_path)}.wav")
            audio.export(temp_path, format="wav")
            return temp_path
        except Exception as e:
            print(f"⚠️ Audio Sanitization Warning: {e}")
            return audio_path # Fallback to original if pydub fails
    
    def _should_sanitize_audio(self, sanitize_audio=True):
        if isinstance(sanitize_audio, str):
            return sanitize_audio.strip().lower() not in {"off", "false", "0", "no", "raw", "none"}
        return bool(sanitize_audio)

    def transcribe(self, audio_path, language=None, model_choice="auto", dialect_hint="", sanitize_audio=True):
        if not audio_path: 
            return [{"text": "", "speaker": "SYSTEM"}]
        
        try:
            clean_path = self._sanitize_audio(audio_path) if self._should_sanitize_audio(sanitize_audio) else audio_path
            manual_route = self._normalize_manual_route(model_choice)
            routes = [manual_route] if manual_route else self._auto_routes(language=language, dialect_hint=dialect_hint, audio_path=clean_path)
            if manual_route and manual_route["engine"] != "whisper":
                routes.append({"engine": "whisper", "model": "small", "label": "Whisper small (Fallback)"})

            transcription_text = ""
            used_label = "Unknown"
            tried_mms = any(route["engine"] == "mms" for route in routes)
            for route in routes:
                try:
                    text, elapsed_s = self._transcribe_with_route(route, clean_path, language=language, dialect_hint=dialect_hint)
                    used_label = route["label"]
                    print(f"🎙️ Speech model route: requested={model_choice or 'auto'} resolved={used_label} time={elapsed_s:.2f}s")
                    transcription_text = text
                    if not manual_route and route["model"] == self.TURBO_MODEL_ID and elapsed_s > self.turbo_slow_threshold_s:
                        transcription_text = ""
                        continue
                    if not manual_route and route["engine"] != "mms" and self._looks_wrong_script(text, language, dialect_hint) and not tried_mms:
                        tried_mms = True
                        mms_route = {"engine": "mms", "model": self.MMS_MODEL_ID, "label": "MMS 1B (Low-resource)"}
                        print("🔁 ASR output looked off-script/empty; retrying with MMS 1B.")
                        mms_text, mms_elapsed = self._transcribe_with_route(mms_route, clean_path, language=language, dialect_hint=dialect_hint)
                        if mms_text:
                            transcription_text = mms_text
                            used_label = mms_route["label"]
                            print(f"🎙️ Speech model route: retry resolved={used_label} time={mms_elapsed:.2f}s")
                    if transcription_text:
                        break
                except Exception as route_error:
                    print(f"⚠️ ASR route '{route['label']}' failed: {route_error}")
                    continue

            gc.collect() # Force memory release

            # Clean up temp file
            if clean_path != audio_path and os.path.exists(clean_path):
                os.remove(clean_path)
                
            return [{"text": transcription_text, "speaker": "Speaker 1", "model": used_label}]
        except Exception as e:
            print(f"❌ Transcription Error: {e}")
            return [{"text": "", "speaker": "ERROR"}]