"""Video overlay processor for BADAS — creates overlay videos from prediction dicts."""

import cv2
import numpy as np
import os
import subprocess
import shutil
import tempfile
import threading
from pathlib import Path
from collections import deque
from concurrent.futures import ThreadPoolExecutor
from typing import List, Dict, Any, Optional, Tuple
from fractions import Fraction

from badas.visualization.renderer import (
    OverlayRenderer,
    _check_ffmpeg,
    _check_decord,
    _check_nvenc,
    _is_hazard_event,
    _tts_duration,
    _find_active_word_idx,
    _sanitize_text,
)
from badas.visualization.video_base import BaseVideoProcessor


# ── VideoOverlayProcessor ─────────────────────────────────────────────────────

class VideoOverlayProcessor(BaseVideoProcessor):
    """
    Processor for creating video overlays with collision predictions.

    Uses FFmpeg pipe for fast H.264 encoding compatible with browser playback.
    Falls back to OpenCV if FFmpeg is not available.
    """

    INIT_PERIOD = 2.0

    def __init__(self):
        super().__init__()
        self.renderer    = OverlayRenderer()
        self._has_decord = _check_decord()
        # nvenc check is expensive (~4-6s) on compute GPUs (A100/H100) that lack
        # NVENC hardware.  Disabled by default; set to True explicitly on consumer /
        # workstation GPUs (RTX, Quadro) where nvenc is available.
        self._has_nvenc  = False
        self._show_heatmap  = False
        self._heatmap_alpha = 0.5

    def create_overlay_video(
        self,
        video_path: str,
        predictions: List[Dict[str, Any]],
        output_path: Optional[str] = None,
        quiet: bool = False,
        use_ffmpeg: bool = True,
        pause_on_hazard: bool = True,
        pause_duration: float = 4.5,
        post_tts_buffer: float = 1.5,
        use_tts: bool = True,
        show_heatmap: bool = False,
        heatmap_alpha: float = 0.5,
        **kwargs
    ) -> str:
        """
        Create a video with prediction overlays burned in.

        TTS audio and word-timing data are read directly from the prediction
        dicts (keys ``tts_audio_path`` and ``tts_word_events``).  Attach a
        :class:`~badas.inference.callbacks.TTSCallback` to ``predict_video``
        to populate those keys before calling this function.

        Args:
            video_path:      Path to input video.
            predictions:     List of prediction dicts from ``predict_video``.
            output_path:     Destination path (default: ``<stem>_overlay.mp4``).
            quiet:           Suppress progress output.
            use_ffmpeg:      Use FFmpeg for encoding (faster, browser-compatible).
            pause_on_hazard: Freeze on hazard predictions (those with a
                             ``description`` or ``bboxes``).
            pause_duration:  Freeze length when no TTS audio is present (seconds).
            post_tts_buffer: Extra seconds to hold the freeze after TTS audio ends.
            use_tts:         Use pre-synthesised TTS data from predictions when
                             present.  Set to False to render without audio even
                             if predictions carry ``tts_audio_path`` data.
            show_heatmap:    Blend the spatial heatmap onto every frame.
            heatmap_alpha:   Heatmap blend weight (0 = invisible, 1 = opaque).

        Returns:
            Path to the output video file.
        """
        if output_path is None:
            video_path_obj = Path(video_path)
            output_path = str(
                video_path_obj.parent / f"{video_path_obj.stem}_overlay.mp4"
            )

        self._show_heatmap  = show_heatmap
        self._heatmap_alpha = heatmap_alpha

        # Collect pre-synthesised TTS from predictions (populated by TTSCallback).
        _pretts = (
            {p['timestamp']: p
             for p in predictions
             if 'tts_audio_path' in p and 'tts_word_events' in p}
            if use_tts else {}
        )

        if not _pretts:
            # ── No TTS: direct render ─────────────────────────────────────────
            if use_ffmpeg and self._has_nvenc:
                return self._create_with_gpu(
                    video_path, predictions, output_path, quiet,
                    pause_on_hazard=pause_on_hazard, pause_duration=pause_duration,
                )
            elif use_ffmpeg and self._has_ffmpeg:
                return self._create_with_ffmpeg(
                    video_path, predictions, output_path, quiet,
                    pause_on_hazard=pause_on_hazard, pause_duration=pause_duration,
                )
            else:
                if not quiet and use_ffmpeg:
                    print("Warning: FFmpeg not found. Using OpenCV (slower, may not play in browser).")
                return self._create_with_opencv(
                    video_path, predictions, output_path, quiet,
                    pause_on_hazard=pause_on_hazard, pause_duration=pause_duration,
                )

        # ── TTS flow ──────────────────────────────────────────────────────────
        with tempfile.TemporaryDirectory() as tmp:
            # 1. Build maps from pre-synthesised data in predictions
            tts_wav_map    = {ts: p['tts_audio_path']  for ts, p in _pretts.items()}
            tts_timing_map = {ts: p['tts_word_events'] for ts, p in _pretts.items()}

            if not quiet:
                print(f"[TTS] {len(tts_wav_map)} audio clip(s) ready.")

            # 2. Get video properties for output-time calculation
            cap = cv2.VideoCapture(video_path)
            fps_v        = cap.get(cv2.CAP_PROP_FPS)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            cap.release()

            # 3. Map hazard timestamps → output timestamps (TTS-duration aware)
            tts_events = self._collect_tts_events(
                predictions, fps_v, total_frames,
                tts_wav_map, tts_timing_map,
                pause_on_hazard, pause_duration, post_tts_buffer)
            if not quiet and tts_wav_map and not tts_events:
                print("[TTS] Warning: clips were synthesised but no output timestamps matched.")

            # 4. Render visual track with karaoke word highlighting
            visual_path = os.path.join(tmp, "visual.mp4")
            if use_ffmpeg and self._has_nvenc:
                self._create_with_gpu(
                    video_path, predictions, visual_path, quiet,
                    pause_on_hazard=pause_on_hazard, pause_duration=pause_duration,
                    post_tts_buffer=post_tts_buffer, tts_timing_map=tts_timing_map,
                )
            elif use_ffmpeg and self._has_ffmpeg:
                self._create_with_ffmpeg(
                    video_path, predictions, visual_path, quiet,
                    pause_on_hazard=pause_on_hazard, pause_duration=pause_duration,
                    post_tts_buffer=post_tts_buffer, tts_timing_map=tts_timing_map,
                )
            else:
                if not quiet and use_ffmpeg:
                    print("Warning: FFmpeg not found. Using OpenCV.")
                self._create_with_opencv(
                    video_path, predictions, visual_path, quiet,
                    pause_on_hazard=pause_on_hazard, pause_duration=pause_duration,
                    post_tts_buffer=post_tts_buffer, tts_timing_map=tts_timing_map,
                )

            if not tts_events:
                # Nothing to mix — just move visual to final destination
                shutil.move(visual_path, output_path)
                return output_path

            # 5. Determine output video duration from the rendered file
            cap_out = cv2.VideoCapture(visual_path)
            out_frames = int(cap_out.get(cv2.CAP_PROP_FRAME_COUNT))
            out_fps    = cap_out.get(cv2.CAP_PROP_FPS) or fps_v
            cap_out.release()
            total_output_duration = out_frames / out_fps if out_fps else 0

            # 6. Mix TTS audio into the video
            from badas.analysis.tts import mix_tts_into_video
            success = mix_tts_into_video(
                visual_path, output_path, tts_events,
                total_output_duration, quiet=quiet,
            )

            if not success:
                if not quiet:
                    print("[TTS] Audio mixing failed — saving visual-only output.")
                shutil.move(visual_path, output_path)

        if not quiet:
            print(f"Saved: {output_path}")
        return output_path

    def _collect_tts_events(
        self,
        predictions: List[Dict[str, Any]],
        fps: float,
        total_frames: int,
        tts_wav_map: Dict[float, str],
        tts_timing_map: Dict[float, List[Dict]],
        pause_on_hazard: bool,
        pause_duration: float,
        post_tts_buffer: float,
    ) -> List[Tuple[float, str]]:
        """Compute (output_time_s, wav_path) for each hazard, mirroring the render loop.

        The output time accounts for accumulated freeze-frame offsets from earlier events.
        When TTS timing is available the freeze length is ``tts_duration + post_tts_buffer``;
        otherwise it falls back to the fixed ``pause_duration``.
        """
        frame_lookup = self._create_frame_lookup(predictions, fps, total_frames)
        seen_ts: set = set()
        accumulated_offset = 0.0
        events: List[Tuple[float, str]] = []

        for frame_idx in range(total_frames):
            pred = frame_lookup.get(frame_idx)
            if pred is None:
                continue
            ts = pred["timestamp"]
            if ts in seen_ts:
                continue
            if pause_on_hazard and _is_hazard_event(pred):
                seen_ts.add(ts)
                wav_path = tts_wav_map.get(ts)
                word_events = tts_timing_map.get(ts) or []
                if wav_path:
                    output_time = frame_idx / fps + accumulated_offset
                    events.append((output_time, wav_path))
                freeze_s = (
                    _tts_duration(word_events) + post_tts_buffer
                    if word_events else pause_duration
                )
                accumulated_offset += freeze_s

        return events

    # ── internal: FFmpeg path ─────────────────────────────────────────────────

    def _create_with_ffmpeg(
        self,
        video_path: str,
        predictions: List[Dict[str, Any]],
        output_path: str,
        quiet: bool,
        pause_on_hazard: bool = True,
        pause_duration: float = 4.5,
        post_tts_buffer: float = 1.5,
        tts_timing_map: Optional[Dict[float, List[Dict]]] = None,
    ) -> str:
        cap, fps, width, height, total_frames = self._open_video(video_path)

        if not quiet:
            print(f"Processing {total_frames} frames @ {fps:.1f}fps with FFmpeg...")

        frame_predictions = self._create_frame_lookup(predictions, fps, total_frames)
        timeline_strip    = self._build_timeline_strip(predictions, total_frames, width, fps)

        first_pred  = min(predictions, key=lambda p: p['timestamp']) if predictions else None
        init_period = first_pred['timestamp'] if first_pred else self.INIT_PERIOD
        init_frames = first_pred['window_end'] if first_pred else 16

        proc = self._open_ffmpeg_pipe(width, height, fps, output_path)

        # Parallel rendering: thread pool renders frames while main thread
        # reads and writes in order.  Hazard-freeze bursts are also parallelised
        # since each freeze frame is a pure function of its index.
        NWORKERS    = min(os.cpu_count() or 4, 8)
        MAX_INFLIGHT = NWORKERS * 4   # max frames queued at once

        def _render_normal(fidx: int, raw: np.ndarray) -> bytes:
            return self._draw_frame_overlay(
                raw, fidx, fps, frame_predictions,
                timeline_strip, total_frames,
                init_period=init_period, init_frames=init_frames,
            ).tobytes()

        def _render_freeze(raw: np.ndarray, pred, chars, awid, fidx: int) -> bytes:
            fz = self.renderer.draw_hazard_freeze(
                raw.copy(), pred, chars_visible=chars, active_word_idx=awid)
            return self._blit_timeline(fz, timeline_strip, fidx, total_frames).tobytes()

        seen_hazard_timestamps: set = set()
        frame_idx  = 0
        pending: deque = deque()   # deque of Future[bytes]

        def _drain_to(target_len: int) -> None:
            """Write completed futures until len(pending) <= target_len."""
            while len(pending) > target_len:
                proc.stdin.write(pending.popleft().result())

        try:
            with ThreadPoolExecutor(max_workers=NWORKERS) as pool:
                while True:
                    ret, frame = cap.read()
                    if not ret:
                        break

                    pred = frame_predictions.get(frame_idx)
                    ts   = pred['timestamp'] if pred else None

                    if (
                        pause_on_hazard
                        and pred is not None
                        and _is_hazard_event(pred)
                        and ts not in seen_hazard_timestamps
                    ):
                        # Flush all in-flight frames before the hazard burst
                        _drain_to(0)

                        seen_hazard_timestamps.add(ts)
                        word_events = (tts_timing_map or {}).get(ts) or []
                        freeze_s = (_tts_duration(word_events) + post_tts_buffer
                                    if word_events else pause_duration)
                        n_freeze    = int(fps * freeze_s)
                        raw_desc    = _sanitize_text(pred.get('description') or '')
                        total_chars = len(raw_desc)
                        type_frames = int(n_freeze * 0.33) if (total_chars > 0 and not word_events) else 0

                        frozen = frame  # read-only base; each thread copies
                        for fi in range(n_freeze):
                            elapsed_s = fi / fps
                            if word_events:
                                awid  = _find_active_word_idx(word_events, elapsed_s)
                                chars = None
                            else:
                                awid  = None
                                chars = (
                                    (max(1, int((fi + 1) / type_frames * total_chars))
                                     if fi < type_frames else total_chars)
                                    if total_chars > 0 else None
                                )
                            pending.append(
                                pool.submit(_render_freeze, frozen, pred, chars, awid, frame_idx))
                            _drain_to(MAX_INFLIGHT)

                        # Flush freeze frames before the normal frame
                        _drain_to(0)

                    # Normal frame
                    pending.append(
                        pool.submit(_render_normal, frame_idx, frame))
                    frame_idx += 1
                    _drain_to(MAX_INFLIGHT)

                # Drain remaining
                _drain_to(0)

        except BrokenPipeError:
            pass
        finally:
            cap.release()
            proc.stdin.close()

        stderr_bytes = proc.stderr.read()
        proc.wait()
        if proc.returncode != 0:
            raise RuntimeError(
                f"FFmpeg failed (exit {proc.returncode}):\n"
                + stderr_bytes.decode(errors='replace')
            )

        if not quiet:
            print(f"Saved: {output_path}")
        return output_path

    # ── internal: GPU pipeline (decord NVDEC + PyAV NVENC) ───────────────────

    def _create_with_gpu(
        self,
        video_path: str,
        predictions: List[Dict[str, Any]],
        output_path: str,
        quiet: bool,
        pause_on_hazard: bool = True,
        pause_duration: float = 4.5,
        post_tts_buffer: float = 1.5,
        tts_timing_map: Optional[Dict[float, List[Dict]]] = None,
    ) -> str:
        import av

        # ── Probe metadata (fast header read) ──────────────────────────────
        _probe       = cv2.VideoCapture(video_path)
        fps          = _probe.get(cv2.CAP_PROP_FPS)
        width        = int(_probe.get(cv2.CAP_PROP_FRAME_WIDTH))
        height       = int(_probe.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(_probe.get(cv2.CAP_PROP_FRAME_COUNT))
        _probe.release()

        if fps == 0 or total_frames == 0:
            raise RuntimeError(f"Failed to open video: {video_path!r}")

        # ── GPU / CPU video reader ──────────────────────────────────────────
        import decord
        try:
            _vr = decord.VideoReader(video_path, ctx=decord.gpu(0))
            _use_gpu_read = True
        except Exception:
            _vr = decord.VideoReader(video_path, ctx=decord.cpu(0))
            _use_gpu_read = False

        def _read_frame(idx: int) -> np.ndarray:
            """Return BGR uint8 HWC array for frame *idx*."""
            return _vr[idx].asnumpy()[:, :, ::-1]  # decord returns RGB

        if not quiet:
            codec_label = 'h264_nvenc' if self._has_nvenc else 'libx264'
            read_label  = 'NVDEC' if _use_gpu_read else 'CPU'
            print(f"Processing {total_frames} frames @ {fps:.1f}fps "
                  f"({read_label} decode + {codec_label} encode)...")

        frame_predictions = self._create_frame_lookup(predictions, fps, total_frames)
        timeline_strip    = self._build_timeline_strip(predictions, total_frames, width, fps)

        first_pred  = min(predictions, key=lambda p: p['timestamp']) if predictions else None
        init_period = first_pred['timestamp'] if first_pred else self.INIT_PERIOD
        init_frames = first_pred['window_end'] if first_pred else 16

        # ── PyAV output (h264_nvenc → libx264 fallback) ─────────────────────
        codec_name = 'h264_nvenc' if self._has_nvenc else 'libx264'
        container  = av.open(output_path, 'w', options={'movflags': 'faststart'})
        stream     = container.add_stream(codec_name,
                                          rate=Fraction(fps).limit_denominator(1001),
                                          options={'preset': 'fast'})
        stream.width  = width
        stream.height = height
        stream.pix_fmt = 'yuv420p'

        # ── Parallel rendering (same pattern as _create_with_ffmpeg) ────────
        NWORKERS     = min(os.cpu_count() or 4, 8)
        MAX_INFLIGHT = NWORKERS * 4

        def _render_normal(fidx: int, raw: np.ndarray) -> bytes:
            return self._draw_frame_overlay(
                raw, fidx, fps, frame_predictions,
                timeline_strip, total_frames,
                init_period=init_period, init_frames=init_frames,
            ).tobytes()

        def _render_freeze(raw: np.ndarray, pred, chars, awid, fidx: int) -> bytes:
            fz = self.renderer.draw_hazard_freeze(
                raw.copy(), pred, chars_visible=chars, active_word_idx=awid)
            return self._blit_timeline(fz, timeline_strip, fidx, total_frames).tobytes()

        def _write_bytes(frame_bytes: bytes) -> None:
            frame_bgr = np.frombuffer(frame_bytes, dtype=np.uint8).reshape(height, width, 3)
            frame_rgb = np.ascontiguousarray(frame_bgr[:, :, ::-1])
            av_frame  = av.VideoFrame.from_ndarray(frame_rgb, format='rgb24')
            for packet in stream.encode(av_frame):
                container.mux(packet)

        seen_hazard_timestamps: set = set()
        pending: deque = deque()

        def _drain_to(target_len: int) -> None:
            while len(pending) > target_len:
                _write_bytes(pending.popleft().result())

        try:
            with ThreadPoolExecutor(max_workers=NWORKERS) as pool:
                for frame_idx in range(total_frames):
                    frame = _read_frame(frame_idx)
                    pred  = frame_predictions.get(frame_idx)
                    ts    = pred['timestamp'] if pred else None

                    if (
                        pause_on_hazard
                        and pred is not None
                        and _is_hazard_event(pred)
                        and ts not in seen_hazard_timestamps
                    ):
                        _drain_to(0)
                        seen_hazard_timestamps.add(ts)
                        word_events = (tts_timing_map or {}).get(ts) or []
                        freeze_s = (_tts_duration(word_events) + post_tts_buffer
                                    if word_events else pause_duration)
                        n_freeze    = int(fps * freeze_s)
                        raw_desc    = _sanitize_text(pred.get('description') or '')
                        total_chars = len(raw_desc)
                        type_frames = int(n_freeze * 0.33) if (total_chars > 0 and not word_events) else 0

                        frozen = frame
                        for fi in range(n_freeze):
                            elapsed_s = fi / fps
                            if word_events:
                                awid  = _find_active_word_idx(word_events, elapsed_s)
                                chars = None
                            else:
                                awid  = None
                                chars = (
                                    (max(1, int((fi + 1) / type_frames * total_chars))
                                     if fi < type_frames else total_chars)
                                    if total_chars > 0 else None
                                )
                            pending.append(
                                pool.submit(_render_freeze, frozen, pred, chars, awid, frame_idx))
                            _drain_to(MAX_INFLIGHT)

                        _drain_to(0)

                    pending.append(pool.submit(_render_normal, frame_idx, frame))
                    _drain_to(MAX_INFLIGHT)

                _drain_to(0)

        finally:
            try:
                for packet in stream.encode(None):
                    container.mux(packet)
            except Exception:
                pass
            container.close()

        if not quiet:
            print(f"Saved: {output_path}")
        return output_path

    # ── internal: OpenCV fallback ─────────────────────────────────────────────

    def _create_with_opencv(
        self,
        video_path: str,
        predictions: List[Dict[str, Any]],
        output_path: str,
        quiet: bool,
        pause_on_hazard: bool = True,
        pause_duration: float = 4.5,
        post_tts_buffer: float = 1.5,
        tts_timing_map: Optional[Dict[float, List[Dict]]] = None,
    ) -> str:
        cap, fps, width, height, total_frames = self._open_video(video_path)

        if not quiet:
            print(f"Processing {total_frames} frames @ {fps:.1f}fps with OpenCV...")

        frame_predictions = self._create_frame_lookup(predictions, fps, total_frames)
        timeline_strip    = self._build_timeline_strip(predictions, total_frames, width, fps)

        first_pred  = min(predictions, key=lambda p: p['timestamp']) if predictions else None
        init_period = first_pred['timestamp'] if first_pred else self.INIT_PERIOD
        init_frames = first_pred['window_end'] if first_pred else 16

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out    = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        if not out.isOpened():
            out = cv2.VideoWriter(output_path,
                                  cv2.VideoWriter_fourcc(*'avc1'), fps, (width, height))

        seen_hazard_timestamps: set = set()
        frame_idx = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break

            pred = frame_predictions.get(frame_idx)
            ts   = pred['timestamp'] if pred else None

            if (
                pause_on_hazard
                and pred is not None
                and _is_hazard_event(pred)
                and ts not in seen_hazard_timestamps
            ):
                seen_hazard_timestamps.add(ts)
                word_events = (tts_timing_map or {}).get(ts) or []
                if word_events:
                    freeze_s = _tts_duration(word_events) + post_tts_buffer
                else:
                    freeze_s = pause_duration
                n_freeze    = int(fps * freeze_s)

                raw_desc    = _sanitize_text(pred.get('description') or '')
                total_chars = len(raw_desc)
                type_frames = int(n_freeze * 0.33) if (total_chars > 0 and not word_events) else 0

                for fi in range(n_freeze):
                    elapsed_s = fi / fps
                    if word_events:
                        active_word_idx = _find_active_word_idx(word_events, elapsed_s)
                        chars = None
                    else:
                        active_word_idx = None
                        if total_chars > 0:
                            chars = (max(1, int((fi + 1) / type_frames * total_chars))
                                     if fi < type_frames else total_chars)
                        else:
                            chars = None
                    fz = self.renderer.draw_hazard_freeze(
                        frame.copy(), pred,
                        chars_visible=chars, active_word_idx=active_word_idx)
                    fz = self._blit_timeline(fz, timeline_strip, frame_idx, total_frames)
                    out.write(fz)

            frame = self._draw_frame_overlay(
                frame, frame_idx, fps, frame_predictions,
                timeline_strip, total_frames,
                init_period=init_period, init_frames=init_frames)
            out.write(frame)
            frame_idx += 1

        cap.release()
        out.release()

        if not quiet:
            print(f"Saved: {output_path}")
        return output_path

    # ── internal: per-frame helpers ───────────────────────────────────────────

    def _draw_frame_overlay(
        self,
        frame: np.ndarray,
        frame_idx: int,
        fps: float,
        frame_predictions: Dict[int, Dict[str, Any]],
        timeline_strip: Optional[np.ndarray] = None,
        total_frames: int = 0,
        init_period: float = None,
        init_frames: int = 16,
    ) -> np.ndarray:
        """Draw appropriate overlay on a single frame."""
        current_time = frame_idx / fps
        pred = frame_predictions.get(frame_idx)

        if pred is None:
            if init_period and init_period > 0:
                buffer_fill = min(init_frames, int(current_time / init_period * init_frames))
            else:
                buffer_fill = 0
            frame = self.renderer.draw_initialization_overlay(
                frame, current_time,
                buffer_fill=buffer_fill, buffer_size=init_frames,
                init_period=init_period)
        else:
            heatmap_data = pred.get('attention_map') if pred.get('attention_map') is not None else pred.get('hazard_map')
            if self._show_heatmap and heatmap_data is not None:
                h, w = frame.shape[:2]
                # Resize to frame dimensions and convert to uint8 in one pass
                heat_u8 = (cv2.resize(heatmap_data, (w, h),
                                      interpolation=cv2.INTER_LINEAR) * 255
                           ).astype(np.uint8)
                # cv2.applyColorMap is ~6× faster than numpy fancy indexing
                heat_bgr = cv2.applyColorMap(heat_u8, cv2.COLORMAP_INFERNO)
                frame = cv2.addWeighted(frame, 1.0 - self._heatmap_alpha,
                                        heat_bgr, self._heatmap_alpha, 0)

            prob     = pred.get('probability', 0)
            raw_prob = pred.get('raw_probability', prob)
            latency  = pred.get('latency_ms', 0)
            led_state = ('red'    if prob >= 0.8 else
                         'orange' if prob >= 0.5 else 'green')
            frame = self.renderer.draw_prediction_overlay(
                frame,
                probability=prob, raw_probability=raw_prob,
                current_time=current_time, latency_ms=latency,
                led_state=led_state, model_info="BADAS")

        if timeline_strip is not None:
            frame = self._blit_timeline(frame, timeline_strip, frame_idx, total_frames)
        return frame

    def _build_timeline_strip(
        self,
        predictions: List[Dict[str, Any]],
        total_frames: int,
        width: int,
        fps: float,
        strip_h: int = 12,
    ) -> np.ndarray:
        """Build a (strip_h, width, 3) BGR strip coloured green/orange/red by risk."""
        strip = np.full((strip_h, width, 3), 40, dtype=np.uint8)
        if not predictions or total_frames == 0 or fps == 0 or width <= 1:
            return strip

        sorted_preds   = sorted(predictions, key=lambda x: x['timestamp'])
        ts_array       = np.array([p['timestamp'] for p in sorted_preds])
        prob_array     = np.array([p.get('probability', 0.0) for p in sorted_preds])
        total_duration = total_frames / fps

        x_coords = np.arange(width)
        t_coords = x_coords / (width - 1) * total_duration
        idxs     = np.searchsorted(ts_array, t_coords, side='right') - 1
        idxs     = np.clip(idxs, 0, len(sorted_preds) - 1)
        probs    = prob_array[idxs]

        first_pred_x = int(ts_array[0] / total_duration * (width - 1)) if total_duration > 0 else 0
        valid  = x_coords >= first_pred_x
        colors = np.full((width, 3), 40, dtype=np.uint8)
        colors[valid & (probs >= 0.8)] = [55,  55,  225]
        colors[valid & (probs >= 0.5) & (probs < 0.8)] = [30, 145, 255]
        colors[valid & (probs < 0.5)] = [80, 200, 100]

        strip[:] = colors[np.newaxis, :]
        return strip

    def _blit_timeline(
        self,
        frame: np.ndarray,
        timeline_strip: np.ndarray,
        frame_idx: int,
        total_frames: int,
    ) -> np.ndarray:
        """Blit the timeline strip onto the bottom of *frame* with a white time cursor."""
        height, width = frame.shape[:2]
        strip_h = min(timeline_strip.shape[0], height)
        frame[height - strip_h:height, :] = timeline_strip[:strip_h, :]
        cursor_x = (int(frame_idx / (total_frames - 1) * (width - 1))
                    if total_frames > 1 else 0)
        cursor_x = max(0, min(cursor_x, width - 1))
        cv2.line(frame,
                 (cursor_x, height - strip_h),
                 (cursor_x, height - 1),
                 (255, 255, 255), 2)
        return frame


# ── Convenience function ──────────────────────────────────────────────────────

def create_overlay_video(
    video_path: str,
    predictions: List[Dict[str, Any]],
    output_path: Optional[str] = None,
    quiet: bool = False,
    pause_on_hazard: bool = True,
    pause_duration: float = 4.5,
    **kwargs
) -> str:
    """Convenience wrapper around :class:`VideoOverlayProcessor`."""
    processor = VideoOverlayProcessor()
    return processor.create_overlay_video(
        video_path, predictions, output_path, quiet=quiet,
        pause_on_hazard=pause_on_hazard, pause_duration=pause_duration,
        **kwargs,
    )
