"""
Visualization Utilities for BADAS Collision Anticipation.

Provides:
- Progress bars (auto-detects notebook vs terminal)
- Beautiful prediction visualizations
- Video overlay generation
- Video display in notebooks (base64 encoding)
"""

from pathlib import Path
from typing import List, Dict, Any, Optional

from badas.visualization.charts import (
    VisualizationConfig,
    visualize_predictions,
    plot_attribution_events,
)
from badas.visualization.display_utils import (
    is_notebook,
    get_progress_bar,
    _DummyPbar,
    get_risk_color,
    get_risk_level,
    format_risk_level,
    get_risk_emoji,
    print_prediction,
)


def generate_heatmap_video(
    video_path: str,
    predictions: List[Dict[str, Any]],
    output_path: Optional[str] = None,
    heatmap_alpha: float = 0.5,
    show_bbox: bool = False,
    show_risk: bool = True,
    quiet: bool = False,
) -> str:
    """Generate a video with per-window heatmap overlays.

    Each frame shows the original video with the spatial attribution heatmap
    from the nearest inference window blended on top.  Predictions must have
    been collected with ``attach_attribution(always_compute=True)`` so that
    every prediction dict contains a ``hazard_map`` field.

    Args:
        video_path:    Path to the source video.
        predictions:   Output of ``BADAS.predict_video()``.
        output_path:   Destination path.  Defaults to
                       ``<stem>_heatmap.mp4`` next to the source video.
        heatmap_alpha: Blend weight of the heatmap layer [0–1].  Default 0.5.
        show_bbox:     Draw bounding boxes if present.
        show_risk:     Render the standard risk panel at the bottom.
        quiet:         Suppress progress messages.

    Returns:
        Path to the written video file.
    """
    from badas.visualization.heatmap import HeatmapVideoProcessor

    video_path_obj = Path(video_path)
    if output_path is None:
        output_path = str(video_path_obj.parent / f"{video_path_obj.stem}_heatmap.mp4")

    processor = HeatmapVideoProcessor()
    return processor.create_heatmap_video(
        video_path=str(video_path),
        predictions=predictions,
        output_path=str(output_path),
        heatmap_alpha=heatmap_alpha,
        show_bbox=show_bbox,
        show_risk=show_risk,
        quiet=quiet,
    )


def generate_overlay_video(
    video_path: str,
    predictions: List[Dict[str, Any]],
    output_path: Optional[str] = None,
    quiet: bool = False,
    config: VisualizationConfig = None,
    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,
) -> str:
    """
    Generate video with collision risk overlay.

    TTS audio is read from prediction dicts (keys ``tts_audio_path`` and
    ``tts_word_events``) when present.  Populate these by attaching a
    :class:`~badas.inference.callbacks.TTSCallback` to ``predict_video``.

    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 status messages.
        config:          Unused; kept for API compatibility.
        pause_on_hazard: Freeze on hazard predictions.
        pause_duration:  Freeze length when no TTS audio is present (seconds).
        post_tts_buffer: Extra seconds to hold after TTS audio ends.
        use_tts:         Use pre-synthesised TTS data from predictions when
                         present.  Set False to render without audio.
        show_heatmap:    Blend spatial heatmap onto every frame.
        heatmap_alpha:   Heatmap blend weight (0 = invisible, 1 = opaque).

    Returns:
        Path to output video.
    """
    if config is not None:
        import warnings
        warnings.warn(
            "The 'config' parameter of generate_overlay_video() is unused and will be "
            "removed in a future version. Remove it from your call.",
            DeprecationWarning,
            stacklevel=2,
        )

    from badas.visualization.overlay import VideoOverlayProcessor

    video_path_obj = Path(video_path)
    if output_path is None:
        output_path = video_path_obj.parent / f"{video_path_obj.stem}_overlay.mp4"

    if not quiet:
        print("Creating overlay video...")

    processor = VideoOverlayProcessor()
    result_path = processor.create_overlay_video(
        video_path=str(video_path),
        predictions=predictions,
        output_path=str(output_path),
        quiet=True,  # suppress processor's own "Saved:" line; wrapper owns output
        pause_on_hazard=pause_on_hazard,
        pause_duration=pause_duration,
        post_tts_buffer=post_tts_buffer,
        use_tts=use_tts,
        show_heatmap=show_heatmap,
        heatmap_alpha=heatmap_alpha,
    )

    if not quiet:
        print(f"Saved to: {result_path}")

    return result_path


def display_video(video_path: str, width: int = 640, height: int = 480):
    """
    Display video in Jupyter notebook.

    Uses IPython.display.Video for reliable playback.

    Args:
        video_path: Path to video file
        width: Display width
        height: Display height
    """
    from IPython.display import display, Video

    video_path = Path(video_path).resolve()
    if not video_path.exists():
        print(f"Video not found: {video_path}")
        return

    display(Video(str(video_path), width=width, height=height, embed=True, html_attributes='controls'))


def display_video_comparison(
    original_path: str,
    overlay_path: str,
    width: int = 480,
    height: int = 360
):
    """
    Display original and overlay videos side by side in notebook.

    Uses IPython.display.Video with embed=True for reliable browser playback.

    Args:
        original_path: Path to original video
        overlay_path: Path to overlay video
        width: Width of each video
        height: Height of each video
    """
    from IPython.display import display, HTML, Video

    original_path = Path(original_path).resolve()
    overlay_path = Path(overlay_path).resolve()

    # Check files exist
    if not original_path.exists():
        print(f"Original video not found: {original_path}")
        return
    if not overlay_path.exists():
        print(f"Overlay video not found: {overlay_path}")
        return

    # Create Video objects with embedding (reads file once, stores in notebook)
    original_video = Video(str(original_path), width=width, height=height, embed=True)
    overlay_video = Video(str(overlay_path), width=width, height=height, embed=True)

    # Extract the HTML from Video objects and combine in a table
    original_html = original_video._repr_html_()
    overlay_html = overlay_video._repr_html_()

    html = f'''
    <table style="margin: 0 auto; border-collapse: separate; border-spacing: 20px 0;">
        <tr>
            <th style="color: #2c3e50; font-size: 16px; padding-bottom: 10px;">Original Video</th>
            <th style="color: #2c3e50; font-size: 16px; padding-bottom: 10px;">With Risk Overlay</th>
        </tr>
        <tr>
            <td style="vertical-align: top;">{original_html}</td>
            <td style="vertical-align: top;">{overlay_html}</td>
        </tr>
    </table>
    '''

    display(HTML(html))
