"""Video overlay generation pipeline.

Reads original video frames, draws skeleton overlays from H5 data,
and writes a composited MP4 at the same resolution/fps.
"""

from __future__ import annotations

import logging
from collections.abc import Callable
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)


def generate_overlay_video(
    video_path: Path,
    h5_path: Path,
    output_path: Path,
    viz_type: str,
    progress_callback: Callable[[int, int], None] | None = None,
) -> Path:
    """Generate a skeleton overlay video from pose H5 data.

    Args:
        video_path: Path to the original video file.
        h5_path: Path to H5 file containing pose detections.
        output_path: Where to write the overlay MP4.
        viz_type: One of "mediapipe", "wilor", "nlf".
        progress_callback: Called with (current_frame, total_frames) during processing.

    Returns:
        Path to the generated overlay video.

    Raises:
        FileNotFoundError: If video or H5 file doesn't exist.
        ValueError: If viz_type is not recognized.
    """
    import cv2
    import h5py

    video_path = Path(video_path)
    h5_path = Path(h5_path)
    output_path = Path(output_path)

    if not video_path.exists():
        raise FileNotFoundError(f"Video not found: {video_path}")
    if not h5_path.exists():
        raise FileNotFoundError(f"H5 file not found: {h5_path}")
    if viz_type not in ("mediapipe", "wilor", "nlf", "teaser", "rtmpose"):
        raise ValueError(
            f"Unknown viz_type: {viz_type!r}. "
            f"Expected 'mediapipe', 'wilor', 'nlf', 'teaser', or 'rtmpose'."
        )

    # Open video
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise RuntimeError(f"Cannot open video: {video_path}")

    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Prepare output
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # type: ignore[attr-defined]
    writer = cv2.VideoWriter(str(output_path), fourcc, fps, (width, height))

    if not writer.isOpened():
        cap.release()
        raise RuntimeError(f"Cannot create video writer for: {output_path}")

    try:
        with h5py.File(h5_path, "r") as f:
            frame_idx = f["frame_idx"][:] if "frame_idx" in f else None

            if viz_type == "mediapipe":
                _render_mediapipe(cap, writer, f, frame_idx, total_frames, progress_callback)
            elif viz_type == "wilor":
                _render_wilor(cap, writer, f, frame_idx, total_frames, progress_callback)
            elif viz_type == "nlf":
                _render_nlf(cap, writer, f, frame_idx, total_frames, progress_callback)
            elif viz_type == "teaser":
                _render_teaser(cap, writer, f, frame_idx, total_frames, progress_callback)
            elif viz_type == "rtmpose":
                _render_rtmpose(cap, writer, f, total_frames, progress_callback)
    finally:
        cap.release()
        writer.release()

    # Re-encode to H.264 for browser compatibility if ffmpeg is available
    _reencode_h264(output_path)

    logger.info("Overlay video written to %s (%d frames)", output_path, total_frames)
    return output_path


def _render_mediapipe(
    cap,
    writer,
    f,
    frame_idx: np.ndarray | None,
    total_frames: int,
    progress_callback: Callable[[int, int], None] | None,
) -> None:
    """Render MediaPipe skeleton overlay."""
    from sltk.io.h5 import get_frame_detections
    from sltk.visualization.skeleton import draw_mediapipe_skeleton

    poses = f["poses"][:] if "poses" in f else None
    confidence = f["confidence"][:] if "confidence" in f else None
    left_hand = f["left_hand"][:] if "left_hand" in f else None
    right_hand = f["right_hand"][:] if "right_hand" in f else None

    if poses is None:
        logger.warning("No 'poses' dataset found in H5 file")
        _copy_video(cap, writer, total_frames, progress_callback)
        return

    # Detect dense vs sparse format
    is_sparse = (
        frame_idx is not None and len(poses) < total_frames and len(frame_idx) >= total_frames
    )

    for frame_num in range(total_frames):
        ret, frame = cap.read()
        if not ret:
            break

        if is_sparse:
            assert frame_idx is not None
            start, end = get_frame_detections(f, frame_num, frame_idx)
            if start >= 0 and start < len(poses):
                for det_i in range(start, min(end, len(poses))):
                    kpts = poses[det_i]
                    conf = (
                        confidence[det_i]
                        if confidence is not None and det_i < len(confidence)
                        else None
                    )
                    lh = (
                        left_hand[det_i]
                        if left_hand is not None and det_i < len(left_hand)
                        else None
                    )
                    rh = (
                        right_hand[det_i]
                        if right_hand is not None and det_i < len(right_hand)
                        else None
                    )
                    draw_mediapipe_skeleton(frame, kpts, conf, lh, rh)
        elif frame_num < len(poses):
            kpts = poses[frame_num]
            conf = (
                confidence[frame_num]
                if confidence is not None and frame_num < len(confidence)
                else None
            )
            lh = (
                left_hand[frame_num]
                if left_hand is not None and frame_num < len(left_hand)
                else None
            )
            rh = (
                right_hand[frame_num]
                if right_hand is not None and frame_num < len(right_hand)
                else None
            )
            draw_mediapipe_skeleton(frame, kpts, conf, lh, rh)

        writer.write(frame)
        if progress_callback:
            progress_callback(frame_num + 1, total_frames)


def _render_wilor(
    cap,
    writer,
    f,
    frame_idx: np.ndarray | None,
    total_frames: int,
    progress_callback: Callable[[int, int], None] | None,
) -> None:
    """Render WiLoR hand overlay.

    Supports two H5 storage formats:
    - Sparse: ``frame_idx (T, 2)`` with ``[start_idx, count]`` per frame.
    - Flat: ``img_idx (N,)`` with per-detection absolute frame number.
    """
    from sltk.visualization.skeleton import draw_wilor_hands

    kpts_2d = f["kpts_2d"][:] if "kpts_2d" in f else None
    right_flags = f["right"][:] if "right" in f else None

    if kpts_2d is None:
        logger.warning("No 'kpts_2d' in H5 file for WiLoR")
        _copy_video(cap, writer, total_frames, progress_callback)
        return

    # Determine storage format
    use_sparse = frame_idx is not None
    img_idx = f["img_idx"][:] if (not use_sparse and "img_idx" in f) else None

    if not use_sparse and img_idx is None:
        logger.warning("No 'frame_idx' or 'img_idx' in H5 file for WiLoR")
        _copy_video(cap, writer, total_frames, progress_callback)
        return

    if use_sparse:
        from sltk.io.h5 import get_frame_detections

    for frame_num in range(total_frames):
        ret, frame = cap.read()
        if not ret:
            break

        if use_sparse:
            assert frame_idx is not None
            start, end = get_frame_detections(f, frame_num, frame_idx)
        else:
            mask = np.where(img_idx == frame_num)[0]
            if len(mask) > 0:
                start, end = int(mask[0]), int(mask[-1]) + 1
            else:
                start, end = -1, -1

        if start >= 0:
            det_kpts = np.array(kpts_2d[start:end], dtype=np.float32)
            det_right = right_flags[start:end] if right_flags is not None else None
            draw_wilor_hands(frame, det_kpts, det_right)

        writer.write(frame)
        if progress_callback:
            progress_callback(frame_num + 1, total_frames)


def _render_nlf(
    cap,
    writer,
    f,
    frame_idx: np.ndarray | None,
    total_frames: int,
    progress_callback: Callable[[int, int], None] | None,
) -> None:
    """Render NLF/SMPL-X body overlay.

    Supports two H5 storage formats:
    - Sparse: ``frame_idx (T, 2)`` with ``[start_idx, count]`` per frame.
    - Flat: ``img_idx (N,)`` with per-detection absolute frame number.
    """
    from sltk.visualization.skeleton import draw_nlf_body

    joints2d = f["joints2d"][:] if "joints2d" in f else None

    if joints2d is None:
        logger.warning("No 'joints2d' in H5 file for NLF")
        _copy_video(cap, writer, total_frames, progress_callback)
        return

    # Determine storage format
    use_sparse = frame_idx is not None
    img_idx = f["img_idx"][:] if (not use_sparse and "img_idx" in f) else None

    if not use_sparse and img_idx is None:
        logger.warning("No 'frame_idx' or 'img_idx' in H5 file for NLF")
        _copy_video(cap, writer, total_frames, progress_callback)
        return

    if use_sparse:
        from sltk.io.h5 import get_frame_detections

    for frame_num in range(total_frames):
        ret, frame = cap.read()
        if not ret:
            break

        if use_sparse:
            assert frame_idx is not None
            start, end = get_frame_detections(f, frame_num, frame_idx)
        else:
            mask = np.where(img_idx == frame_num)[0]
            if len(mask) > 0:
                start, end = int(mask[0]), int(mask[-1]) + 1
            else:
                start, end = -1, -1

        if start >= 0:
            det_joints = joints2d[start:end]  # (N, 55, 2)
            draw_nlf_body(frame, det_joints)

        writer.write(frame)
        if progress_callback:
            progress_callback(frame_num + 1, total_frames)


def _render_teaser(
    cap,
    writer,
    f,
    frame_idx: np.ndarray | None,
    total_frames: int,
    progress_callback: Callable[[int, int], None] | None,
) -> None:
    """Render TEASER face detection + FLAME parameter overlay.

    H5 layout expected:
    - ``boxes (M, 4)`` xyxy bounding boxes
    - ``boxes_scores (M,)`` detection confidences
    - ``flame/jaw_pose (M, 3)``
    - ``flame/expression (M, 50)``
    - ``flame/eyelid_params (M, 2)``
    - ``frame_idx (T, 2)`` sparse index
    """
    from sltk.visualization.skeleton import draw_teaser_face, draw_teaser_landmarks

    boxes = f["boxes"][:] if "boxes" in f else None
    scores = f["boxes_scores"][:] if "boxes_scores" in f else None

    # FLAME params (in group)
    jaw_pose = f["flame/jaw_pose"][:] if "flame/jaw_pose" in f else None
    expression = f["flame/expression"][:] if "flame/expression" in f else None
    eyelid = f["flame/eyelid_params"][:] if "flame/eyelid_params" in f else None

    # Optional FLAME landmarks (computed by FLAMEProcessor)
    landmarks_2d = f["landmarks_2d"][:] if "landmarks_2d" in f else None

    if boxes is None:
        logger.warning("No 'boxes' dataset in H5 file for TEASER")
        _copy_video(cap, writer, total_frames, progress_callback)
        return

    use_sparse = frame_idx is not None
    if use_sparse:
        from sltk.io.h5 import get_frame_detections

    for frame_num in range(total_frames):
        ret, frame = cap.read()
        if not ret:
            break

        if use_sparse:
            assert frame_idx is not None
            start, end = get_frame_detections(f, frame_num, frame_idx)
        else:
            start, end = -1, -1

        if start >= 0:
            det_boxes = np.array(boxes[start:end], dtype=np.float32)
            det_scores = (
                np.array(scores[start:end], dtype=np.float32) if scores is not None else None
            )
            det_jaw = (
                np.array(jaw_pose[start:end], dtype=np.float32) if jaw_pose is not None else None
            )
            det_expr = (
                np.array(expression[start:end], dtype=np.float32)
                if expression is not None
                else None
            )
            det_eyelid = (
                np.array(eyelid[start:end], dtype=np.float32) if eyelid is not None else None
            )
            draw_teaser_face(frame, det_boxes, det_scores, det_jaw, det_expr, det_eyelid)
            if landmarks_2d is not None:
                det_lmk = np.array(landmarks_2d[start:end], dtype=np.float32)
                draw_teaser_landmarks(frame, det_lmk)

        writer.write(frame)
        if progress_callback:
            progress_callback(frame_num + 1, total_frames)


def _render_rtmpose(
    cap,
    writer,
    f,
    total_frames: int,
    progress_callback: Callable[[int, int], None] | None,
) -> None:
    """Render RTMPose whole-body skeleton overlay.

    RTMPose H5 stores dense aligned arrays (filtered mode):
    ``keypoints (T, N, 133, 3)`` where N = number of people.
    """
    from sltk.visualization.skeleton import draw_rtmpose_wholebody

    keypoints = f["keypoints"][:] if "keypoints" in f else None

    if keypoints is None:
        logger.warning("No 'keypoints' dataset in H5 file for RTMPose")
        _copy_video(cap, writer, total_frames, progress_callback)
        return

    for frame_num in range(total_frames):
        ret, frame = cap.read()
        if not ret:
            break

        if frame_num < len(keypoints):
            frame_kpts = np.array(keypoints[frame_num], dtype=np.float32)
            # Filter out all-NaN persons
            valid = ~np.isnan(frame_kpts[:, :, 0]).all(axis=1)
            if valid.any():
                draw_rtmpose_wholebody(frame, frame_kpts[valid])

        writer.write(frame)
        if progress_callback:
            progress_callback(frame_num + 1, total_frames)


def _copy_video(
    cap,
    writer,
    total_frames: int,
    progress_callback: Callable[[int, int], None] | None,
) -> None:
    """Copy video frames unchanged (fallback when no pose data found)."""
    for frame_num in range(total_frames):
        ret, frame = cap.read()
        if not ret:
            break
        writer.write(frame)
        if progress_callback:
            progress_callback(frame_num + 1, total_frames)


def _reencode_h264(path: Path) -> None:
    """Re-encode from mp4v to H.264 for browser playback compatibility."""
    import shutil
    import subprocess

    if shutil.which("ffmpeg") is None:
        logger.debug("ffmpeg not found, skipping H.264 re-encode")
        return

    tmp = path.with_suffix(".tmp.mp4")
    try:
        subprocess.run(
            [
                "ffmpeg",
                "-y",
                "-i",
                str(path),
                "-c:v",
                "libx264",
                "-preset",
                "fast",
                "-crf",
                "23",
                "-movflags",
                "+faststart",
                "-an",  # No audio needed for overlay
                str(tmp),
            ],
            capture_output=True,
            timeout=600,
            check=True,
        )
        tmp.replace(path)
    except Exception as exc:
        logger.warning("H.264 re-encode failed (%s), keeping mp4v version", exc)
        tmp.unlink(missing_ok=True)
        if not isinstance(
            exc,
            (subprocess.CalledProcessError, subprocess.TimeoutExpired, OSError),
        ):
            raise
