"""Prediction enrichment callbacks for BADAS inference.

Callbacks are callables that receive each prediction dict and can enrich it
in-place (add new keys).  Pass them to ``predict_video(callbacks=[...])``.

The prediction dict is the single source of truth: downstream visualisation
functions (``generate_overlay_video``, ``visualize_predictions``) simply read
whatever keys are present — they do not need to know which callbacks ran.

Example — attach TTS synthesis so the overlay video has audio without
passing ``tts_provider`` to the visualiser::

    from badas.inference.callbacks import TTSCallback
    from badas.analysis.tts import EdgeTTSProvider

    with TTSCallback(EdgeTTSProvider()) as tts_cb:
        predictions = predictor.predict_video(video, callbacks=[tts_cb])

    generate_overlay_video(video, predictions)   # audio already in predictions

Custom callback — log every high-risk window to a file::

    class LogCallback(PredictionCallback):
        def __call__(self, prediction):
            if prediction['probability'] >= 0.8:
                with open('alerts.jsonl', 'a') as f:
                    json.dump({'ts': prediction['timestamp'],
                               'prob': prediction['probability']}, f)
                    f.write('\\n')
"""

from __future__ import annotations

import os
import shutil
import tempfile
from abc import ABC, abstractmethod
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from badas.analysis.tts import TTSProvider


class PredictionCallback(ABC):
    """Abstract base for prediction enrichment callbacks.

    A callback receives each prediction dict and may enrich it in-place by
    adding new keys, or perform side-effects (logging, alerting, etc.).

    Minimal example::

        class DangerAlertCallback(PredictionCallback):
            def __call__(self, prediction):
                if prediction['risk_level'] == 'danger':
                    send_alert(prediction['timestamp'])
    """

    @abstractmethod
    def __call__(self, prediction: Dict[str, Any]) -> None:
        """Enrich or act on *prediction* in-place.  No return value is used."""


class TTSCallback(PredictionCallback):
    """Synthesise TTS audio for every prediction that carries a ``description``.

    Adds two keys to each matching prediction dict:

    * ``tts_audio_path``  — absolute path to the synthesised WAV file.
    * ``tts_word_events`` — list of ``{offset_s, duration_s, text}`` dicts
      for karaoke word-highlight synchronisation in the overlay video.

    Files are kept in *audio_dir* (a managed temp directory by default) until
    :meth:`cleanup` is called.  ``generate_overlay_video`` reads the paths
    directly from the prediction dicts — no ``tts_provider`` argument needed.

    Usage — context manager (recommended; cleans up automatically)::

        with TTSCallback(EdgeTTSProvider()) as tts_cb:
            predictions = predictor.predict_video(video, callbacks=[tts_cb])

        generate_overlay_video(video, predictions)   # audio in predictions

    Usage — explicit cleanup::

        tts_cb = TTSCallback(provider=EdgeTTSProvider())
        predictions = predictor.predict_video(video, callbacks=[tts_cb])
        generate_overlay_video(video, predictions)
        tts_cb.cleanup()

    Args:
        provider:  A :class:`~badas.analysis.tts.TTSProvider` instance.
        audio_dir: Directory for WAV files.  Defaults to a managed temp dir.
        condition: ``(prediction) -> bool`` selecting which predictions to
                   synthesise.  Default: any prediction with a ``description``.
        quiet:     Suppress progress output.
    """

    def __init__(
        self,
        provider: "TTSProvider",
        audio_dir: Optional[str] = None,
        condition: Optional[Callable[[Dict[str, Any]], bool]] = None,
        quiet: bool = False,
    ) -> None:
        self.provider   = provider
        self._audio_dir = audio_dir
        self._condition = condition
        self.quiet      = quiet
        self._owns_dir  = audio_dir is None   # True → we created it, we clean it
        self._dir: Optional[str] = None       # resolved lazily on first synthesis

    # ── Condition ──────────────────────────────────────────────────────────────

    def _should_run(self, prediction: Dict[str, Any]) -> bool:
        if self._condition is not None:
            return bool(self._condition(prediction))
        return bool(prediction.get('description'))

    # ── Core ───────────────────────────────────────────────────────────────────

    def __call__(self, prediction: Dict[str, Any]) -> None:
        if not self._should_run(prediction):
            return

        # Resolve / create the audio directory on first use
        if self._dir is None:
            self._dir = (
                tempfile.mkdtemp(prefix='badas_tts_')
                if self._owns_dir
                else self._audio_dir
            )
            os.makedirs(self._dir, exist_ok=True)

        ts   = float(prediction.get('timestamp', 0.0))
        desc = _format_description(prediction['description'])
        if not desc:
            return

        wav_path = os.path.join(self._dir, f'tts_{ts:.3f}.wav')
        if not self.quiet:
            print(f"[TTS] Synthesising t={ts:.2f}s: {desc[:60]!r}", flush=True)

        ok, word_events = self.provider.synthesize_with_timing(desc, wav_path)
        if not ok:
            if not self.quiet:
                print(f"[TTS] Synthesis failed at t={ts:.2f}s — skipping.", flush=True)
            return

        prediction['tts_audio_path']  = wav_path
        prediction['tts_word_events'] = word_events or []

        if not self.quiet and word_events:
            n_sents = sum(desc.count(c) for c in '.!?')
            print(
                f"[TTS] {len(word_events)} words captured "
                f"({max(1, n_sents)} sentence(s))",
                flush=True,
            )

    # ── Lifecycle ──────────────────────────────────────────────────────────────

    def cleanup(self) -> None:
        """Remove synthesised WAV files.  Safe to call multiple times."""
        if self._owns_dir and self._dir and os.path.isdir(self._dir):
            shutil.rmtree(self._dir, ignore_errors=True)
            self._dir = None

    def __enter__(self) -> "TTSCallback":
        return self

    def __exit__(self, *_: Any) -> None:
        self.cleanup()


# ── Helpers ────────────────────────────────────────────────────────────────────

def _format_description(desc: str) -> str:
    """Join multi-line descriptions (Action / Reasoning) with natural pauses."""
    lines = [line.strip() for line in desc.splitlines() if line.strip()]
    return "... ".join(lines)
