"""
Text-to-speech providers for BADAS.

Used in two modes:
- **Batch** (`synthesize`): write a WAV file for each hazard description, then mix
  the clips into the output video via FFmpeg.  Used by VideoOverlayProcessor.
- **Real-time** (`speak`): play synthesized audio through system speakers immediately.
  Designed for live inference (e.g. predict_stream dashcam loop).  The interface is
  defined now; concrete implementations will be added when needed.

Usage (video overlay)::

    from badas.analysis.tts import FasterQwen3TTSProvider

    tts = FasterQwen3TTSProvider(
        model_id="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
        ref_audio="voice.wav",
        ref_text="Reference transcript here.",
    )
    processor.create_overlay_video(video, predictions, tts_provider=tts)

Usage (future real-time)::

    for pred in badas.predict_stream("dashcam"):
        if pred.get("description") and pred["probability"] > 0.8:
            tts.speak(pred["description"], block=False)
"""

from __future__ import annotations

import os
import subprocess
import shutil
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple


# ── WAV writing helper ────────────────────────────────────────────────────────

def _write_wav(audio, sr: int, path: str) -> None:
    """Write a float32 numpy array (or list of arrays) to a WAV file.

    Tries soundfile → scipy.io.wavfile → stdlib wave (in that order).
    """
    import numpy as np

    if isinstance(audio, (list, tuple)):
        audio = np.concatenate([np.asarray(a) for a in audio])
    audio = np.asarray(audio, dtype=np.float32)

    # Attempt 1: soundfile
    try:
        import soundfile as sf
        sf.write(path, audio, sr)
        return
    except ImportError:
        pass

    # Attempt 2: scipy
    try:
        from scipy.io import wavfile
        audio_i16 = (audio * 32767.0).clip(-32768, 32767).astype(np.int16)
        wavfile.write(path, sr, audio_i16)
        return
    except ImportError:
        pass

    # Attempt 3: stdlib wave (no external deps)
    import wave
    import struct
    audio_i16 = (audio * 32767.0).clip(-32768, 32767).astype(np.int16)
    with wave.open(path, "w") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sr)
        wf.writeframes(audio_i16.tobytes())


# ── Word-timing helpers ───────────────────────────────────────────────────────

def _wav_duration(path: str) -> float:
    """Return duration of a WAV file in seconds, or 0.0 on error."""
    import wave
    try:
        with wave.open(path, "rb") as wf:
            return wf.getnframes() / wf.getframerate()
    except Exception:
        return 0.0


def _estimate_word_timings(text: str, duration_s: float) -> List[Dict]:
    """Uniform word-timing fallback when real boundaries are unavailable.

    Distributes words evenly across *duration_s* with a small lead-in,
    so karaoke highlighting still works even without real boundary events.
    """
    words = text.split()
    if not words or duration_s <= 0:
        return []
    lead = 0.1                          # brief silence before first word
    slot = (duration_s - lead) / len(words)
    return [
        {"offset_s": lead + i * slot, "duration_s": slot * 0.85, "text": w}
        for i, w in enumerate(words)
    ]


def _sentence_events_to_words(sentence_events: List[Dict]) -> List[Dict]:
    """Convert sentence-boundary timing events to per-word timing estimates.

    Words within each sentence are weighted by character count (excluding
    trailing punctuation) as a proxy for spoken duration.  This gives much
    better karaoke sync than equal slots: short words like "Brake" get less
    time than long ones like "immediately".

    The last 15% of each sentence's duration is treated as trailing silence
    and is not allocated to any word.
    """
    word_events: List[Dict] = []
    for ev in sentence_events:
        words = ev["text"].split()
        if not words:
            continue
        # Strip trailing punctuation when counting chars so "immediately."
        # doesn't unfairly outweigh "immediately".
        chars = [max(1, len(w.rstrip(".,!?;:"))) for w in words]
        total_chars = sum(chars)
        speech_dur = ev["duration_s"] * 0.85   # last 15% is trailing silence
        t = ev["offset_s"]
        for word, char_count in zip(words, chars):
            word_dur = speech_dur * char_count / total_chars
            word_events.append({
                "offset_s":   t,
                "duration_s": word_dur,
                "text":       word,
            })
            t += word_dur
    return word_events


# ── Abstract base ─────────────────────────────────────────────────────────────

class TTSProvider(ABC):
    """Abstract base for text-to-speech synthesis.

    Concrete subclasses must implement `synthesize()`.
    `speak()` is a stub for future real-time audio playback.
    """

    @abstractmethod
    def synthesize(self, text: str, output_path: str) -> bool:
        """Synthesize *text* to speech and write a WAV to *output_path*.

        Args:
            text: The text to speak.
            output_path: Absolute path where the WAV should be written.

        Returns:
            True on success, False on failure.
        """

    def speak(self, text: str, block: bool = False) -> None:
        """Play *text* aloud through system audio (real-time mode).

        Override in concrete subclasses.  For faster-qwen3-tts this will
        call ``generate_voice_clone_streaming`` and pipe chunks to
        ``sounddevice`` for minimal latency.

        Args:
            text: The text to speak.
            block: If True, block until audio finishes playing.

        Raises:
            NotImplementedError: Until the subclass implements this.
        """
        raise NotImplementedError(
            f"{type(self).__name__} does not yet implement speak(). "
            "Add a speak() override to enable real-time audio."
        )

    def synthesize_with_timing(
        self, text: str, output_path: str
    ) -> Tuple[bool, List[Dict]]:
        """Synthesize *text* and return word-level timing events for karaoke.

        Returns:
            (success, word_events) where each event is a dict with keys:
                ``offset_s``   – word start time in seconds
                ``duration_s`` – word duration in seconds
                ``text``       – the word string

        The default implementation calls :meth:`synthesize` and returns empty
        timings.  Override in providers that support word boundaries (e.g.
        :class:`EdgeTTSProvider`).
        """
        ok = self.synthesize(text, output_path)
        return ok, []


# ── edge-tts implementation (zero-config, Microsoft Neural TTS) ───────────────

# Reference sentence used when auto-generating a bootstrap voice clip.
_DEFAULT_REF_TEXT = (
    "Attention: a hazard has been detected ahead. "
    "Please slow down and proceed with caution."
)
_DEFAULT_REF_VOICE = "en-US-JennyNeural"


def _make_default_ref_audio(cache_path: str, edge_voice: str = _DEFAULT_REF_VOICE) -> bool:
    """Generate a reference WAV using edge-tts and cache it at *cache_path*.

    Used by :class:`FasterQwen3TTSProvider` when no ``ref_audio`` is supplied.
    Requires ``edge-tts``::

        pip install edge-tts

    Returns True on success, False on failure.
    """
    try:
        import edge_tts  # noqa: F401
    except ImportError:
        return False

    import asyncio
    import subprocess

    os.makedirs(os.path.dirname(cache_path), exist_ok=True)
    tmp_mp3 = cache_path + ".tmp.mp3"

    async def _synth() -> None:
        communicate = edge_tts.Communicate(_DEFAULT_REF_TEXT, edge_voice)
        await communicate.save(tmp_mp3)

    try:
        asyncio.run(_synth())
        result = subprocess.run(
            ["ffmpeg", "-y", "-i", tmp_mp3, cache_path],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        return result.returncode == 0
    except Exception:
        return False
    finally:
        if os.path.exists(tmp_mp3):
            os.unlink(tmp_mp3)


class EdgeTTSProvider(TTSProvider):
    """TTS provider backed by edge-tts (Microsoft Neural TTS, free, no API key).

    Works out of the box with no GPU and no reference audio.  Requires an
    internet connection on first call (results are not cached).

    Install::

        pip install edge-tts

    Args:
        voice: Edge TTS voice name.  Defaults to ``"en-US-JennyNeural"``.
            Run ``edge-tts --list-voices`` for the full list.

    Example::

        from badas.analysis.tts import EdgeTTSProvider
        tts = EdgeTTSProvider()                         # default voice
        tts = EdgeTTSProvider(voice="en-US-AriaNeural") # female voice
    """

    def __init__(self, voice: str = _DEFAULT_REF_VOICE) -> None:
        self.voice = voice

    def synthesize(self, text: str, output_path: str) -> bool:
        """Synthesize *text* via Microsoft Neural TTS and write WAV to *output_path*."""
        try:
            import edge_tts
        except ImportError:
            raise ImportError(
                "edge-tts is not installed.  Run:  pip install edge-tts"
            ) from None

        import asyncio
        import subprocess
        import tempfile

        async def _synth() -> str:
            communicate = edge_tts.Communicate(text, self.voice)
            tmp_mp3 = output_path + ".tmp.mp3"
            await communicate.save(tmp_mp3)
            return tmp_mp3

        try:
            import concurrent.futures
            # asyncio.run() fails when called from a running event loop (e.g.
            # Jupyter).  Running in a fresh thread guarantees a clean loop.
            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                tmp_mp3 = pool.submit(asyncio.run, _synth()).result()
            result = subprocess.run(
                ["ffmpeg", "-y", "-i", tmp_mp3, output_path],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            if os.path.exists(tmp_mp3):
                os.unlink(tmp_mp3)
            return result.returncode == 0
        except Exception as exc:
            print(f"[TTS] EdgeTTS synthesis failed for '{text[:60]}': {exc}")
            return False

    def synthesize_with_timing(
        self, text: str, output_path: str
    ) -> Tuple[bool, List[Dict]]:
        """Synthesize *text* and capture per-word timing via the EdgeTTS stream API.

        Returns (success, word_events).  Each event is::

            {'offset_s': float, 'duration_s': float, 'text': str}

        All data is fetched in a single streaming request — no extra round-trip.
        """
        try:
            import edge_tts
        except ImportError:
            raise ImportError(
                "edge-tts is not installed.  Run:  pip install edge-tts"
            ) from None

        import asyncio

        async def _stream() -> Tuple[bytes, List[Dict], List[Dict]]:
            communicate = edge_tts.Communicate(text, self.voice)
            audio_chunks: List[bytes] = []
            word_events: List[Dict] = []
            sentence_events: List[Dict] = []
            async for chunk in communicate.stream():
                t = chunk.get("type")
                if t == "audio":
                    audio_chunks.append(chunk["data"])
                elif t == "WordBoundary":
                    # Preferred: word-level granularity
                    word_events.append({
                        "offset_s":   chunk["offset"]   / 10_000_000,
                        "duration_s": chunk["duration"] / 10_000_000,
                        "text":       chunk["text"],
                    })
                elif t == "SentenceBoundary":
                    # Fallback: sentence-level (some edge-tts versions)
                    sentence_events.append({
                        "offset_s":   chunk["offset"]   / 10_000_000,
                        "duration_s": chunk["duration"] / 10_000_000,
                        "text":       chunk["text"],
                    })
            return b"".join(audio_chunks), word_events, sentence_events

        try:
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                audio_bytes, word_events, sentence_events = pool.submit(
                    asyncio.run, _stream()
                ).result()

            tmp_mp3 = output_path + ".tmp.mp3"
            with open(tmp_mp3, "wb") as fh:
                fh.write(audio_bytes)
            result = subprocess.run(
                ["ffmpeg", "-y", "-i", tmp_mp3, output_path],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            if os.path.exists(tmp_mp3):
                os.unlink(tmp_mp3)
            if result.returncode != 0:
                return False, []

            if word_events:
                # Best case: word-level events from EdgeTTS
                pass
            elif sentence_events:
                # Good fallback: expand sentence boundaries to per-word estimates
                word_events = _sentence_events_to_words(sentence_events)
                print(f"[TTS] Using sentence-boundary timing "
                      f"({len(sentence_events)} sentences → {len(word_events)} words)")
            else:
                # Last resort: uniform estimate from WAV duration
                dur = _wav_duration(output_path)
                if dur > 0:
                    word_events = _estimate_word_timings(text, dur)
                    print(f"[TTS] No boundary events — using uniform timing estimate "
                          f"({len(word_events)} words, {dur:.1f}s)")

            return True, word_events
        except Exception as exc:
            print(f"[TTS] EdgeTTS synthesis (with timing) failed for '{text[:60]}': {exc}")
            return False, []

    def speak(self, text: str, block: bool = False) -> None:
        """Speak *text* aloud via Microsoft Neural TTS.

        Requires ``sounddevice``::

            pip install sounddevice
        """
        try:
            import sounddevice as sd
            import numpy as np
        except ImportError:
            raise ImportError(
                "sounddevice is required for real-time audio playback. "
                "Run:  pip install sounddevice"
            ) from None

        import asyncio
        import tempfile
        import subprocess

        async def _synth():
            import edge_tts
            communicate = edge_tts.Communicate(text, self.voice)
            tmp_mp3 = tempfile.mktemp(suffix=".mp3")
            await communicate.save(tmp_mp3)
            return tmp_mp3

        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            tmp_mp3 = pool.submit(asyncio.run, _synth()).result()

        # Convert MP3 → raw PCM via ffmpeg, pipe to sounddevice
        import subprocess
        proc = subprocess.run(
            ["ffmpeg", "-y", "-i", tmp_mp3, "-f", "f32le", "-ar", "24000", "-ac", "1", "-"],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
        )
        if os.path.exists(tmp_mp3):
            os.unlink(tmp_mp3)

        audio = np.frombuffer(proc.stdout, dtype=np.float32)

        if block:
            sd.play(audio, 24000)
            sd.wait()
        else:
            import threading
            def _play():
                sd.play(audio, 24000)
                sd.wait()
            threading.Thread(target=_play, daemon=True).start()


# ── faster-qwen3-tts implementation ──────────────────────────────────────────

class FasterQwen3TTSProvider(TTSProvider):
    """TTS provider backed by faster-qwen3-tts (CUDA-graph accelerated Qwen3-TTS).

    Install::

        pip install faster-qwen3-tts

    Args:
        model_id: HuggingFace model ID.  Defaults to the compact 0.6 B base model.
        language: Language hint passed to the model (e.g. ``"English"``).
        ref_audio: Path to a WAV file to clone the voice from.  If *None* (default),
            a professional reference voice is auto-generated using ``edge-tts`` on
            the first synthesis call and cached at ``~/.cache/badas/tts/default_ref.wav``.
        ref_text: Transcript of *ref_audio* (improves clone quality).  When
            *ref_audio* is auto-generated, the transcript is set automatically.

    The model is loaded **lazily** on the first call to :meth:`synthesize` or
    :meth:`speak` so that importing this class does not pay the GPU warm-up cost.
    """

    _DEFAULT_CACHE = os.path.expanduser("~/.cache/badas/tts/default_ref.wav")

    def __init__(
        self,
        model_id: str = "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
        language: str = "English",
        ref_audio: Optional[str] = None,
        ref_text: Optional[str] = None,
    ) -> None:
        self.model_id  = model_id
        self.language  = language
        self.ref_audio = ref_audio
        self.ref_text  = ref_text or ""
        self._model    = None   # lazy init

    # ── private helpers ───────────────────────────────────────────────────────

    def _get_model(self):
        if self._model is None:
            try:
                from faster_qwen3_tts import FasterQwen3TTS
            except ImportError:
                raise ImportError(
                    "faster-qwen3-tts is not installed. "
                    "Run:  pip install faster-qwen3-tts"
                ) from None
            self._model = FasterQwen3TTS.from_pretrained(self.model_id)
        return self._model

    def _resolve_ref_audio(self) -> Tuple[str, str]:
        """Return (ref_audio_path, ref_text) — generating a default if needed."""
        if self.ref_audio:
            return self.ref_audio, self.ref_text

        cache = self._DEFAULT_CACHE
        if not os.path.exists(cache):
            print("[TTS] No ref_audio provided — generating default voice clip via edge-tts …")
            ok = _make_default_ref_audio(cache)
            if not ok:
                raise RuntimeError(
                    "Could not generate a default reference voice.  "
                    "Either provide ref_audio= to FasterQwen3TTSProvider, "
                    "or install edge-tts:  pip install edge-tts"
                )
            print(f"[TTS] Default reference voice cached at {cache}")
        return cache, _DEFAULT_REF_TEXT

    # ── TTSProvider interface ─────────────────────────────────────────────────

    def synthesize(self, text: str, output_path: str) -> bool:
        """Synthesize *text* and write to *output_path* (WAV).

        Returns True on success, False on any error (prints a warning).
        """
        try:
            ref_audio, ref_text = self._resolve_ref_audio()
            model = self._get_model()
            audio_list, sr = model.generate_voice_clone(
                text=text,
                language=self.language,
                ref_audio=ref_audio,
                ref_text=ref_text,
                non_streaming_mode=True,
            )
            _write_wav(audio_list, sr, output_path)
            return True
        except Exception as exc:
            print(f"[TTS] Synthesis failed for '{text[:60]}': {exc}")
            return False

    def speak(self, text: str, block: bool = False) -> None:
        """Speak *text* aloud through system audio using streaming synthesis.

        Requires ``sounddevice``::

            pip install sounddevice

        Args:
            text: The text to speak.
            block: If True, block until audio finishes.  If False, playback
                runs in a background thread (fire-and-forget).
        """
        try:
            import sounddevice as sd
            import numpy as np
        except ImportError:
            raise ImportError(
                "sounddevice is required for real-time audio playback. "
                "Run:  pip install sounddevice"
            ) from None

        ref_audio, ref_text = self._resolve_ref_audio()
        model = self._get_model()

        chunks = []
        for audio_chunk, sr, _ in model.generate_voice_clone_streaming(
            text=text,
            language=self.language,
            ref_audio=ref_audio,
            ref_text=ref_text,
            non_streaming_mode=True,
        ):
            chunks.append(np.asarray(audio_chunk, dtype=np.float32))
        if not chunks:
            return
        audio = np.concatenate(chunks)

        if block:
            sd.play(audio, sr)
            sd.wait()
        else:
            import threading
            def _play():
                sd.play(audio, sr)
                sd.wait()
            threading.Thread(target=_play, daemon=True).start()


# ── FFmpeg audio mixing helper ────────────────────────────────────────────────

def mix_tts_into_video(
    visual_path: str,
    output_path: str,
    tts_events: List[Tuple[float, str]],
    total_output_duration: float,
    quiet: bool = False,
) -> bool:
    """Embed TTS audio clips into a silent overlay video using FFmpeg.

    Creates a silent base audio track spanning the full output duration, then
    mixes each TTS WAV clip at its designated output timestamp.

    Args:
        visual_path: Path to the visual-only overlay video (temp file).
        output_path: Destination path for the final video with audio.
        tts_events: List of ``(output_time_seconds, wav_path)`` pairs, sorted
            by output time.
        total_output_duration: Duration of *visual_path* in seconds.
        quiet: Suppress FFmpeg stderr output.

    Returns:
        True on success, False on failure.  On failure the caller should fall
        back to using *visual_path* as the output.
    """
    if not shutil.which("ffmpeg"):
        return False

    if not tts_events:
        shutil.copy2(visual_path, output_path)
        return True

    # Build filter_complex:
    #   anullsrc → silent track for the full output duration
    #   each WAV  → adelay to its output timestamp, then amix
    parts: List[str] = []

    # Silent base (sample rate 24 kHz = Qwen3-TTS native; ffmpeg resamples for aac)
    parts.append(
        f"anullsrc=r=24000:cl=mono,"
        f"atrim=duration={total_output_duration:.3f},"
        f"asetpts=PTS-STARTPTS[silence]"
    )

    delay_labels: List[str] = []
    audio_inputs: List[str] = []
    for i, (out_time_s, wav_path) in enumerate(tts_events):
        audio_inputs.extend(["-i", wav_path])
        delay_ms = int(out_time_s * 1000)
        label = f"[tts{i}]"
        # aresample normalises any sample-rate differences before delay
        parts.append(
            f"[{i + 1}:a]aresample=24000,"
            f"adelay={delay_ms}|{delay_ms}{label}"
        )
        delay_labels.append(label)

    n_inputs = 1 + len(tts_events)
    mix_inputs = "[silence]" + "".join(delay_labels)
    parts.append(
        f"{mix_inputs}amix=inputs={n_inputs}:duration=first:normalize=0[outa]"
    )

    cmd = (
        ["ffmpeg", "-y", "-i", visual_path]
        + audio_inputs
        + [
            "-filter_complex", ";".join(parts),
            "-map", "0:v",
            "-map", "[outa]",
            "-c:v", "copy",
            "-c:a", "aac",
            output_path,
        ]
    )

    stderr = subprocess.DEVNULL if quiet else None
    result = subprocess.run(
        cmd, stdout=subprocess.DEVNULL, stderr=stderr
    )
    return result.returncode == 0
