"""RTMPose whole-body keypoint extraction wrapper.

Extracts 133 COCO-WholeBody keypoints (17 body + 6 foot + 68 face + 42 hand)
from video using MMPose's ``rtmw-x_8xb320-270e_cocktail14-384x288`` model.
Supports multi-person detection with optional person tracking and filtering.

Usage::

    from sltk.extraction.rtmpose import RTMPoseExtractor, RTMPoseConfig

    config = RTMPoseConfig(device="cuda:0", filter=True)
    with RTMPoseExtractor(config) as ext:
        ext.load_model()
        result = ext.extract_from_video("video.mp4", "wholebody.h5")
"""

from __future__ import annotations

import gc
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

from sltk.data.core import PoseSequence
from sltk.extraction.base import (
    ExtractionConfig,
    ExtractionResult,
    PoseExtractor,
)

logger = logging.getLogger(__name__)


def _to_numpy_batch(batch: Any) -> np.ndarray:
    """Convert a decord/torch batch object to a numpy array."""
    if hasattr(batch, "asnumpy"):
        return batch.asnumpy()
    if hasattr(batch, "detach"):
        return batch.detach().cpu().numpy()
    if hasattr(batch, "numpy"):
        return batch.numpy()
    return np.asarray(batch)


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------


@dataclass
class RTMPoseConfig(ExtractionConfig):
    """RTMPose-specific configuration.

    Attributes:
        model_name: MMPose model alias.  Weights are auto-downloaded by
            ``openmim`` on first use — no local bundling needed.
        kpt_thr: Keypoint confidence threshold for MMPose inference.
        bbox_thr: Bounding-box confidence threshold.
        filter: Enable person tracking / bad-detection removal.
        img_batch_size: Frames per decord read batch.
        max_tracking_distance: Max pixel distance for person correspondence.
        kp_dtype_bits: 16 for float16, 32 for float32 storage.
    """

    model_name: str = "rtmw-x_8xb320-270e_cocktail14-384x288"
    kpt_thr: float = 0.5
    bbox_thr: float = 0.5
    filter: bool = True
    img_batch_size: int = 32
    max_tracking_distance: float = 150
    kp_dtype_bits: int = 16

    @property
    def kp_dtype(self) -> type:
        """Numpy dtype for keypoint storage."""
        return np.float16 if self.kp_dtype_bits == 16 else np.float32


# ---------------------------------------------------------------------------
# Extractor
# ---------------------------------------------------------------------------


class RTMPoseExtractor(PoseExtractor):
    """RTMPose whole-body keypoint extractor.

    Extracts 133 COCO-WholeBody keypoints per person using MMPose's RTMPose
    model.  Optionally tracks persons across frames and filters noisy
    detections.
    """

    def __init__(self, config: RTMPoseConfig | None = None) -> None:
        super().__init__(config or RTMPoseConfig())
        self.config: RTMPoseConfig = self.config  # type: ignore[assignment]
        self._processor: Any = None

    @property
    def name(self) -> str:
        return "rtmpose"

    @property
    def output_format(self) -> str:
        return "coco_wholebody"

    # -------------------------------------------------------------------
    # Model loading
    # -------------------------------------------------------------------

    def load_model(self, **kwargs: Any) -> None:
        """Lazy-load the MMPose inferencer.

        MMPose auto-downloads ``rtmw-x`` weights on first use via openmim.
        """
        try:
            from sltk.extraction.rtmpose_deps import (
                RTMPoseProcessor,
                RTMPoseProcessorConfig,
            )
        except ImportError as e:
            raise ImportError(
                f"RTMPose dependencies not found.  Ensure mmpose, mmdet, and "
                f"mmengine are installed (pip install mmpose mmdet mmengine "
                f"openmim).  Error: {e}"
            ) from e

        # CUDA init before decord import to avoid segfault
        # https://github.com/dmlc/decord/issues/324
        import torch

        if torch.cuda.is_available():
            torch.tensor([0], device=self.config.device)

        proc_config = RTMPoseProcessorConfig(
            model_name=kwargs.get("model_name", self.config.model_name),
            kpt_thr=self.config.kpt_thr,
            bbox_thr=self.config.bbox_thr,
            kp_dtype=self.config.kp_dtype,
        )
        self._processor = RTMPoseProcessor(proc_config, device=self.config.device)
        self._initialized = True

    # -------------------------------------------------------------------
    # Video extraction
    # -------------------------------------------------------------------

    def extract_from_video(
        self,
        video_path: str | Path,
        output_path: str | Path | None = None,
    ) -> ExtractionResult:
        """Extract whole-body keypoints from video.

        Args:
            video_path: Path to input video.
            output_path: Optional H5 output path.

        Returns:
            ExtractionResult with keypoints.
        """
        if not self._initialized:
            self.load_model()

        video_path = Path(video_path)

        try:
            import decord

            vr = decord.VideoReader(str(video_path), ctx=decord.cpu(0))
        except Exception as e:
            logger.error("Cannot open video %s: %s", video_path, e)
            return ExtractionResult(errors=[f"Cannot open video: {e}"])

        total_frames = len(vr)
        fps = round(vr.get_avg_fps())
        frame_shape: tuple[int, int] = tuple(vr[0].shape[:2])  # type: ignore[assignment]
        batch_size = self.config.img_batch_size

        # Accumulate raw per-frame detections across batches
        acc_kp: list[np.ndarray | list] = []
        acc_bboxes: list[np.ndarray | list] = []
        acc_bboxes_scores: list[list[float] | list] = []
        acc_sf: dict[int, list[int]] = {}
        total_detections = 0

        try:
            frame_ids = list(range(total_frames))
            for batch_start in range(0, total_frames, batch_size):
                batch_indices = frame_ids[batch_start : batch_start + batch_size]
                batch = vr.get_batch(batch_indices)
                frames_np = list(_to_numpy_batch(batch))

                batch_result = self._processor.process_frames(
                    frames_np,
                    batch_indices,
                )

                acc_kp.extend(batch_result["all_kp"])
                acc_bboxes.extend(batch_result["all_bboxes"])
                acc_bboxes_scores.extend(batch_result["all_bboxes_scores"])
                total_detections += batch_result["num_total_detections"]

                # Merge successful_frames_per_person
                for pid, fids in batch_result["successful_frames_per_person"].items():
                    acc_sf.setdefault(pid, []).extend(fids)

                logger.debug(
                    "Batch %d-%d: %d detections",
                    batch_start,
                    batch_start + len(batch_indices),
                    batch_result["num_total_detections"],
                )

            # Post-process
            if self.config.filter:
                results = self._run_filtering(
                    acc_kp,
                    acc_bboxes,
                    frame_shape,
                    total_frames,
                )
            else:
                results = self._collect_per_person(
                    acc_kp,
                    acc_bboxes,
                    acc_bboxes_scores,
                    acc_sf,
                )

            results["fps"] = fps
            results["num_frames"] = total_frames
            results["resolution"] = frame_shape
            results["num_total_detections"] = total_detections

            poses = self._results_to_poses(results, fps, total_frames)

            result = ExtractionResult(
                poses=poses,
                num_frames=total_frames,
                num_detections=total_detections,
                fps=float(fps),
                metadata={
                    "source_video": str(video_path),
                    "extractor": self.name,
                    "filter": self.config.filter,
                },
            )

            if output_path is not None:
                self._save_rtmpose_results(results, output_path)

            return result

        except Exception as e:
            logger.exception("RTMPose extraction failed")
            return ExtractionResult(
                errors=[str(e)],
                fps=float(fps),
            )

    # -------------------------------------------------------------------
    # Frame extraction
    # -------------------------------------------------------------------

    def extract_from_frames(
        self,
        frames: np.ndarray,
        fps: float = 25.0,
    ) -> ExtractionResult:
        """Extract keypoints from numpy frames.

        Args:
            frames: ``(T, H, W, 3)`` uint8 array.
            fps: Frame rate.

        Returns:
            ExtractionResult with keypoints.
        """
        if not self._initialized:
            self.load_model()

        frames_list = list(frames)
        indices = list(range(len(frames)))
        frame_shape = (frames.shape[1], frames.shape[2])

        batch_result = self._processor.process_frames(frames_list, indices)

        if self.config.filter:
            results = self._run_filtering(
                batch_result["all_kp"],
                batch_result["all_bboxes"],
                frame_shape,
                len(frames),
            )
        else:
            results = self._collect_per_person(
                batch_result["all_kp"],
                batch_result["all_bboxes"],
                batch_result["all_bboxes_scores"],
                batch_result["successful_frames_per_person"],
            )

        results["fps"] = fps
        results["num_frames"] = len(frames)
        results["resolution"] = frame_shape

        poses = self._results_to_poses(results, fps, len(frames))

        return ExtractionResult(
            poses=poses,
            num_frames=len(frames),
            num_detections=batch_result["num_total_detections"],
            fps=fps,
        )

    # -------------------------------------------------------------------
    # Internal helpers
    # -------------------------------------------------------------------

    def _run_filtering(
        self,
        all_kp: list[np.ndarray | list],
        all_bboxes: list[np.ndarray | list],
        frame_shape: tuple[int, int],
        total_frames: int,
    ) -> dict[str, Any]:
        """Apply person tracking and bad-detection removal.

        Returns:
            Dict with ``filtered=True``, ``keypoints``, ``bboxes``,
            ``person_masks``.
        """
        from sltk.extraction.rtmpose_deps.filtering import (
            compute_person_correspondences,
            remove_bad_detections,
        )

        aligned_kp, aligned_bboxes, person_masks = compute_person_correspondences(
            all_kp,
            all_bboxes,
            max_distance=self.config.max_tracking_distance,
        )

        if aligned_kp.size == 0:
            return {
                "filtered": True,
                "keypoints": np.full(
                    (total_frames, 1, 133, 3),
                    np.nan,
                    dtype=self.config.kp_dtype,
                ),
                "bboxes": np.full(
                    (total_frames, 1, 2, 3),
                    np.nan,
                    dtype=self.config.kp_dtype,
                ),
                "person_masks": np.zeros(
                    (total_frames, 1),
                    dtype=bool,
                ),
            }

        clean_kp, clean_bboxes, clean_masks = remove_bad_detections(
            aligned_kp,
            aligned_bboxes,
            person_masks,
            frame_shape,
        )

        assert len(clean_kp) == total_frames, f"Expected {total_frames} frames, got {len(clean_kp)}"

        return {
            "filtered": True,
            "keypoints": clean_kp,
            "bboxes": clean_bboxes,
            "person_masks": clean_masks,
        }

    @staticmethod
    def _collect_per_person(
        all_kp: list[np.ndarray | list],
        all_bboxes: list[np.ndarray | list],
        all_bboxes_scores: list[list[float] | list],
        successful_frames_per_person: dict[int, list[int]],
    ) -> dict[str, Any]:
        """Collect detections into per-person arrays (unfiltered mode).

        Returns:
            Dict with ``filtered=False`` and per-person data.
        """
        per_person_kp: dict[int, list[np.ndarray]] = {}
        per_person_bboxes: dict[int, list[np.ndarray]] = {}
        per_person_scores: dict[int, list[float]] = {}

        for frame_i in range(len(all_kp)):
            for person_i in range(len(all_kp[frame_i])):
                per_person_kp.setdefault(person_i, []).append(
                    all_kp[frame_i][person_i],
                )
                per_person_bboxes.setdefault(person_i, []).append(
                    all_bboxes[frame_i][person_i],
                )
                if all_bboxes_scores[frame_i]:
                    per_person_scores.setdefault(person_i, []).append(
                        all_bboxes_scores[frame_i][person_i],
                    )

        return {
            "filtered": False,
            "per_person_kp": per_person_kp,
            "per_person_bboxes": per_person_bboxes,
            "per_person_scores": per_person_scores,
            "successful_frames_per_person": successful_frames_per_person,
        }

    def _results_to_poses(
        self,
        results: dict[str, Any],
        fps: float,
        num_frames: int,
    ) -> PoseSequence:
        """Convert results to PoseSequence (primary signer only).

        Picks the person with the most detections and packs as dense
        ``(T, 133, 3)`` with NaN for missing frames.
        """
        if results.get("filtered"):
            kp = results["keypoints"]  # (T, N, 133, 3)
            masks = results["person_masks"]  # (T, N)

            # Primary signer = person with most valid frames
            person_counts = masks.sum(axis=0)
            if person_counts.max() == 0:
                data = np.full(
                    (num_frames, 133, 3),
                    np.nan,
                    dtype=np.float32,
                )
                confidence = np.zeros((num_frames, 133), dtype=np.float32)
            else:
                primary = int(np.argmax(person_counts))
                data = kp[:, primary, :, :].astype(np.float32)
                confidence = np.where(
                    masks[:, primary, np.newaxis],
                    data[:, :, 2],
                    0.0,
                ).astype(np.float32)
        else:
            # Unfiltered: pick person_id with most frames
            sf = results.get("successful_frames_per_person", {})
            if not sf:
                data = np.full(
                    (num_frames, 133, 3),
                    np.nan,
                    dtype=np.float32,
                )
                confidence = np.zeros((num_frames, 133), dtype=np.float32)
            else:
                primary_id = max(sf, key=lambda pid: len(sf[pid]))
                person_kp = np.array(
                    results["per_person_kp"][primary_id],
                    dtype=np.float32,
                )
                person_frames = sf[primary_id]

                data = np.full(
                    (num_frames, 133, 3),
                    np.nan,
                    dtype=np.float32,
                )
                confidence = np.zeros((num_frames, 133), dtype=np.float32)
                for i, fid in enumerate(person_frames):
                    if 0 <= fid < num_frames:
                        data[fid] = person_kp[i]
                        confidence[fid] = person_kp[i, :, 2]

        return PoseSequence(
            data=data,
            fps=fps,
            format="coco_wholebody",
            confidence=confidence,
            metadata={
                "num_keypoints": 133,
                "keypoint_layout": {
                    "body": (0, 17),
                    "foot": (17, 23),
                    "face": (23, 91),
                    "left_hand": (91, 112),
                    "right_hand": (112, 133),
                },
            },
        )

    def _save_rtmpose_results(
        self,
        results: dict[str, Any],
        output_path: str | Path,
    ) -> None:
        """Save results to H5 file.

        Filtered mode: top-level ``keypoints (T, N, 133, 3)`` and
        ``bboxes (T, N, 2, 3)``.

        Unfiltered mode: per-person groups with ``keypoints``, ``bboxes``,
        ``successful_frames``, ``bboxes_scores``.
        """
        import h5py

        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        ckw = {"compression": "gzip", "compression_opts": 4}

        with h5py.File(output_path, "w") as hf:
            hf.attrs["fps"] = results.get("fps", 25)
            hf.attrs["num_frames"] = results.get("num_frames", 0)
            hf.attrs["resolution"] = results.get("resolution", (0, 0))
            hf.attrs["kpt_thr"] = self.config.kpt_thr
            hf.attrs["bbox_thr"] = self.config.bbox_thr
            hf.attrs["filter"] = self.config.filter
            hf.attrs["extractor"] = self.name

            if results.get("filtered"):
                hf.create_dataset(
                    "keypoints",
                    data=results["keypoints"],
                    **ckw,
                )
                hf.create_dataset(
                    "bboxes",
                    data=results["bboxes"],
                    **ckw,
                )
            else:
                kp_dtype = self.config.kp_dtype
                for pid in results["per_person_kp"]:
                    grp = hf.create_group(f"person_{pid}")
                    grp.create_dataset(
                        "keypoints",
                        data=np.array(
                            results["per_person_kp"][pid],
                            dtype=kp_dtype,
                        ),
                        **ckw,
                    )
                    grp.create_dataset(
                        "bboxes",
                        data=np.array(
                            results["per_person_bboxes"][pid],
                            dtype=kp_dtype,
                        ),
                        **ckw,
                    )
                    sf = results["successful_frames_per_person"]
                    grp.create_dataset(
                        "successful_frames",
                        data=np.array(sf[pid], dtype=np.int32),
                        **ckw,
                    )
                    if pid in results.get("per_person_scores", {}):
                        grp.create_dataset(
                            "bboxes_scores",
                            data=np.array(
                                results["per_person_scores"][pid],
                                dtype=np.float32,
                            ),
                            **ckw,
                        )

        logger.info("Saved RTMPose results to %s", output_path)

    # -------------------------------------------------------------------
    # Cleanup
    # -------------------------------------------------------------------

    def close(self) -> None:
        """Release resources."""
        if hasattr(self, "_processor") and self._processor is not None:
            del self._processor
            self._processor = None
            gc.collect()
            try:
                import torch

                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                    torch.cuda.empty_cache()
            except ImportError:
                pass
        self._initialized = False

    def __enter__(self) -> RTMPoseExtractor:
        """Context manager entry."""
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Context manager exit — release resources."""
        self.close()


# ---------------------------------------------------------------------------
# Convenience function
# ---------------------------------------------------------------------------


def extract_rtmpose(
    video_path: str | Path,
    output_path: str | Path | None = None,
    device: str = "cuda:0",
    filter: bool = True,
) -> ExtractionResult:
    """Quick extraction using RTMPose.

    Args:
        video_path: Input video path.
        output_path: Optional output H5 path.
        device: Torch device string.
        filter: Enable person tracking / filtering.

    Returns:
        ExtractionResult with whole-body keypoints.
    """
    config = RTMPoseConfig(device=device, filter=filter)
    with RTMPoseExtractor(config) as extractor:
        extractor.load_model()
        return extractor.extract_from_video(video_path, output_path)
