"""Live audio -> transcript -> translated caption pipeline (EN<->AR).

Design goals (fixing the old app's flaws):
  * The microphone stream NEVER closes - capture runs in its own callback
    thread, so no speech is lost while transcription/translation happen.
  * Clause-based segmentation: a caption is produced when David pauses
    (>= silence_ms) or when a segment hits max_segment_s. Whole clauses give
    correct Arabic; fixed 5-second chunks were why the old output was wrong.
  * If processing falls behind, queued segments are merged instead of piling
    up, so captions can lag by at most ~one clause.

Threads:
  audio callback (sounddevice) -> frame queue
  segmenter thread             -> segment queue        (VAD state machine)
  worker thread                -> STT + translation -> ui_callback(caption, source, status)

Direction ("en2ar" / "ar2en") is read live from cfg per segment, so the UI
toggle takes effect immediately. Each direction has its own Whisper model
(cfg["model_size"] for English, cfg["model_size_ar"] multilingual for Arabic);
loaded models are cached so switching back is instant.
"""

import os
import queue
import threading
import time

import numpy as np

import translator
from config import load as load_config

FRAME_MS = 30  # VAD frame size

# Whisper hallucinates these on noise/silence; drop them.
_JUNK_EN = {"thank you", "thanks for watching", "you"}
_JUNK_AR = {"اشتركوا في القناة", "لا تنسوا الاشتراك", "لا تنسوا الاشتراك بالقناة",
            "شكرا للمشاهدة", "شكرا على المشاهدة", "شكرا لكم", "شكرا جزيلا",
            "شكرا", "مع السلامة"}
_AR_DIACRITICS = [chr(c) for c in range(0x064B, 0x0653)] + ["ـ"]


def _is_junk(text, lang):
    t = text.strip().rstrip(".!؟?").strip()
    if lang == "en":
        return t.lower() in _JUNK_EN
    for d in _AR_DIACRITICS:
        t = t.replace(d, "")
    # "ترجمة ..." = subtitle-credit hallucination
    return t in _JUNK_AR or t.startswith("ترجمة ")


class _VAD:
    """webrtcvad if available, otherwise an adaptive energy gate."""

    def __init__(self, sample_rate):
        self.sample_rate = sample_rate
        self._noise = 1e-4
        try:
            import webrtcvad
            self._vad = webrtcvad.Vad(2)
        except Exception:
            self._vad = None

    def is_speech(self, frame_f32):
        if self._vad is not None:
            pcm16 = (np.clip(frame_f32, -1.0, 1.0) * 32767).astype(np.int16)
            try:
                return self._vad.is_speech(pcm16.tobytes(), self.sample_rate)
            except Exception:
                self._vad = None  # bad frame size etc. -> fall back for good
        rms = float(np.sqrt(np.mean(frame_f32 ** 2)) + 1e-9)
        self._noise = 0.995 * self._noise + 0.005 * min(rms, self._noise * 4)
        return rms > max(self._noise * 3.5, 0.0025)


class CaptionPipeline:
    def __init__(self, ui_callback, cfg=None):
        """ui_callback(arabic|None, english, status_str) - called from worker
        threads; the UI layer is responsible for hopping to the main thread."""
        self.cfg = cfg or load_config()
        self.ui = ui_callback
        self.model = None
        self._models = {}        # model name -> loaded WhisperModel
        self._model_locks = {}   # model name -> Lock (per-model, so a big
        self._locks_guard = threading.Lock()  # download never blocks the other)
        self._frames = queue.Queue()
        self._segments = queue.Queue()
        self._running = threading.Event()
        self._threads = []
        self._stream = None
        self._last_caption_at = 0.0
        self._executor = None

    # ---------------- model ----------------

    @staticmethod
    def _direction_langs(direction):
        """-> (spoken language, caption language)."""
        return ("en", "ar") if direction == "en2ar" else ("ar", "en")

    def _model_name_for(self, direction):
        return (self.cfg["model_size"] if direction == "en2ar"
                else self.cfg["model_size_ar"])

    def model_cached(self, direction):
        """True if the model files for `direction` are already on disk (so
        loading it will NOT trigger a big download)."""
        import os
        name = self._model_name_for(direction)
        repo = os.path.expanduser(
            "~/.cache/huggingface/hub/models--Systran--faster-whisper-" + name)
        return name in self._models or os.path.isdir(repo)

    def _lock_for(self, name):
        with self._locks_guard:
            return self._model_locks.setdefault(name, threading.Lock())

    def _load_into_cache(self, name, warm_lang):
        """Load a model into the cache (download on first ever use)."""
        from faster_whisper import WhisperModel
        with self._lock_for(name):
            if name not in self._models:
                # Use every CPU core (default is 4) - meaningful STT speedup
                # on Apple Silicon / multi-core PCs. STT is the latency king.
                model = WhisperModel(name, device="cpu", compute_type="int8",
                                     cpu_threads=max(4, os.cpu_count() or 4))
                # Warm up so the first real caption isn't slow. (Low-level
                # noise, not silence - pure zeros trigger harmless-but-noisy
                # numpy warnings.)
                rng = np.random.default_rng(0)
                warmup = (rng.standard_normal(3200) * 0.001).astype(np.float32)
                model.transcribe(warmup, language=warm_lang)
                self._models[name] = model
            return self._models[name]

    def load_model(self):
        """Blocking; call from a background thread. Downloads on first run."""
        self.ensure_model(self.cfg.get("direction", "en2ar"))

    def ensure_model(self, direction=None):
        """Load (or fetch from cache) the model for `direction` and make it
        the current one. Blocking; call from a background thread."""
        d = direction or self.cfg.get("direction", "en2ar")
        self.model = self._load_into_cache(self._model_name_for(d),
                                           self._direction_langs(d)[0])

    def preload(self, direction):
        """Warm the cache for `direction` WITHOUT switching to it. Blocking;
        call from a background thread (used to fetch the second model while
        the app is already usable in the first direction)."""
        self._load_into_cache(self._model_name_for(direction),
                              self._direction_langs(direction)[0])

    # ---------------- control ----------------

    def start(self, device=None):
        if self._running.is_set():
            return
        import sounddevice as sd
        sr = self.cfg["sample_rate"]
        block = int(sr * FRAME_MS / 1000)
        self._running.set()

        def on_audio(indata, _frames, _time, _status):
            self._frames.put(indata[:, 0].copy())

        self._stream = sd.InputStream(
            samplerate=sr, channels=1, dtype="float32", device=device,
            blocksize=block, callback=on_audio)
        self._stream.start()

        self._threads = [
            threading.Thread(target=self._segment_loop, daemon=True),
            threading.Thread(target=self._worker_loop, daemon=True),
        ]
        for t in self._threads:
            t.start()
        self.ui(None, "", "listening")

    def stop(self):
        self._running.clear()
        if self._stream is not None:
            try:
                self._stream.stop()
                self._stream.close()
            except Exception:
                pass
            self._stream = None
        for q in (self._frames, self._segments):
            try:
                while True:
                    q.get_nowait()
            except queue.Empty:
                pass
        self.ui(None, "", "stopped")

    # ---------------- segmenter ----------------

    def _segment_loop(self):
        cfg = self.cfg
        sr = cfg["sample_rate"]
        vad = _VAD(sr)
        silence_frames_end = max(1, int(cfg["silence_ms"] / FRAME_MS))
        min_speech_frames = max(1, int(cfg["min_speech_ms"] / FRAME_MS))
        pre_roll_frames = max(1, int(cfg["pre_roll_ms"] / FRAME_MS))
        max_frames = int(cfg["max_segment_s"] * 1000 / FRAME_MS)
        linger_frames = int(cfg["caption_linger_s"] * 1000 / FRAME_MS)

        pre_roll, voiced, speech_frames = [], [], 0
        silent_run = 0
        in_speech = False
        silence_since_caption = 0

        while self._running.is_set():
            try:
                frame = self._frames.get(timeout=0.25)
            except queue.Empty:
                continue

            speech = vad.is_speech(frame)

            if not in_speech:
                pre_roll.append(frame)
                if len(pre_roll) > pre_roll_frames:
                    pre_roll.pop(0)
                if speech:
                    in_speech = True
                    voiced = list(pre_roll)
                    voiced.append(frame)
                    speech_frames = 1
                    silent_run = 0
                else:
                    silence_since_caption += 1
                    if silence_since_caption >= linger_frames:
                        self._segments.put(None)  # ask UI to clear caption
                        silence_since_caption = 0  # re-fires if still silent
                continue

            # in speech
            voiced.append(frame)
            if speech:
                speech_frames += 1
                silent_run = 0
            else:
                silent_run += 1

            ended = silent_run >= silence_frames_end
            too_long = len(voiced) >= max_frames
            if ended or too_long:
                if speech_frames >= min_speech_frames:
                    self._segments.put(np.concatenate(voiced))
                    silence_since_caption = 0
                in_speech = False
                pre_roll, voiced, speech_frames, silent_run = [], [], 0, 0

    # ---------------- STT + translation worker ----------------

    def _worker_loop(self):
        while self._running.is_set():
            try:
                seg = self._segments.get(timeout=0.25)
            except queue.Empty:
                continue

            if seg is None:
                # Clear the caption - but not if one was just displayed
                # (a stale clear marker can sit behind segments in the queue).
                if time.time() - self._last_caption_at >= \
                        self.cfg["caption_linger_s"] - 0.5:
                    self.ui("", "", "listening")
                continue

            # Merge any backlog so we never fall more than one clause behind.
            merged = [seg]
            try:
                while True:
                    extra = self._segments.get_nowait()
                    if extra is not None:
                        merged.append(extra)
            except queue.Empty:
                pass
            audio = np.concatenate(merged) if len(merged) > 1 else seg

            src_lang, tgt_lang = self._direction_langs(
                self.cfg.get("direction", "en2ar"))

            t0 = time.time()
            source_text = self._transcribe(audio, src_lang)
            if not source_text:
                continue
            caption = self._translate_with_timeout(
                source_text, src_lang, tgt_lang, timeout=5.0)
            dt = time.time() - t0

            if caption is None:
                self.ui(None, source_text,
                        "translation offline - check internet")
            else:
                self._last_caption_at = time.time()
                self.ui(caption, source_text,
                        "ok %.1fs [%s]" % (dt, translator.last_engine))

    def _translate_with_timeout(self, text, source, target, timeout=5.0):
        """Never let a hung network stall the caption flow on stage."""
        if self._executor is None:
            from concurrent.futures import ThreadPoolExecutor
            self._executor = ThreadPoolExecutor(max_workers=1)
        from concurrent.futures import TimeoutError as FTimeout
        try:
            return self._executor.submit(translator.translate, text,
                                         source, target).result(timeout=timeout)
        except FTimeout:
            return None
        except Exception:
            return None

    def _transcribe(self, audio, lang):
        try:
            segments, _info = self.model.transcribe(
                audio, language=lang, beam_size=1,
                condition_on_previous_text=False,
                no_speech_threshold=0.5)
            text = " ".join(s.text.strip() for s in segments).strip()
            # Whisper hallucinates short phrases on noise; drop obvious junk.
            if _is_junk(text, lang):
                return ""
            return text
        except Exception as e:
            self.ui(None, "", "transcribe error: %s" % e)
            return ""
