"""
WiLoR (Wrist-Led Hand Reconstruction) pose extraction wrapper.

WiLoR provides accurate 3D hand pose estimation with MANO parameters.
This wrapper integrates with the signmunculusisambard pipeline for
GPU-accelerated extraction.

Note: WiLoR requires specific model checkpoints and dependencies.
See: https://github.com/warmshao/WiLoR

Usage:
    from sltk.extraction.wilor import WiLoRExtractor, WiLoRConfig

    config = WiLoRConfig(
        checkpoint_path="path/to/wilor_checkpoint.pth",
        detector_path="path/to/yolo_hand.pt",
    )
    extractor = WiLoRExtractor(config)
    extractor.load_model()
    result = extractor.extract_from_video("video.mp4")
"""

from __future__ import annotations

# Note: MANO model files should be pre-converted to remove chumpy dependencies
# See scripts in sltk/extraction/wilor_deps/ for conversion if needed
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import torch

import numpy as np

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


@dataclass
class WiLoRConfig(ExtractionConfig):
    """WiLoR-specific configuration."""

    # Model paths
    checkpoint_path: str | None = None
    config_path: str | None = None
    detector_path: str | None = None

    # Processing parameters
    rescale_factor: float = 2.0
    detection_confidence: float = 0.3
    img_batch_size: int = 128
    det_batch_size: int = 256

    # GPU acceleration
    use_compile: bool = True
    use_amp: bool = True
    use_gpu_decode: bool = True

    def __post_init__(self):
        """Resolve bundled weights for any unset paths."""
        from sltk.extraction.weights import get_weight_path

        if self.checkpoint_path is None:
            bundled = get_weight_path("wilor_checkpoint")
            if bundled:
                self.checkpoint_path = str(bundled)

        # If checkpoint path is provided, try to derive other paths from the same directory
        # This handles the case where user provides a custom checkpoint location
        checkpoint_dir = None
        if self.checkpoint_path:
            checkpoint_dir = Path(self.checkpoint_path).parent

        if self.detector_path is None:
            bundled = get_weight_path("wilor_detector")
            if bundled:
                self.detector_path = str(bundled)
            elif checkpoint_dir:
                # Try to find detector in the same directory as checkpoint
                candidate = checkpoint_dir / "detector.pt"
                if candidate.exists():
                    self.detector_path = str(candidate)

        if self.config_path is None:
            bundled = get_weight_path("wilor_config")
            if bundled:
                self.config_path = str(bundled)
            elif checkpoint_dir:
                # Try to find config in the same directory as checkpoint
                candidate = checkpoint_dir / "model_config.yaml"
                if candidate.exists():
                    self.config_path = str(candidate)


class WiLoRExtractor(PoseExtractor):
    """
    WiLoR hand pose extractor.

    Extracts 21 3D hand keypoints per hand with MANO parameters.
    Uses YOLO for hand detection and WiLoR for pose estimation.
    """

    def __init__(self, config: WiLoRConfig | None = None):
        super().__init__(config or WiLoRConfig())
        self.config: WiLoRConfig = self.config
        self._processor: Any = None

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

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

    def load_model(self, **kwargs) -> None:
        """
        Load WiLoR model.

        Requires the signmunculusisambard WiLoR processor.
        """
        checkpoint = kwargs.get("checkpoint_path", self.config.checkpoint_path)
        config = kwargs.get("config_path", self.config.config_path)
        detector = kwargs.get("detector_path", self.config.detector_path)

        if checkpoint is None:
            raise ConfigurationError(
                "WiLoR checkpoint_path required. "
                "Set via WiLoRConfig or load_model(checkpoint_path=...)"
            )

        # Import from local wilor_deps
        try:
            import sys

            wilor_deps_path = Path(__file__).parent / "wilor_deps"
            if str(wilor_deps_path) not in sys.path:
                sys.path.insert(0, str(wilor_deps_path))

            from wilor_processor import WiLORConfig, WiLORProcessorOptimized
        except ImportError as e:
            raise ImportError(
                f"WiLoR processor not found: {e}. "
                "Ensure wilor_deps is present in sltk/extraction/"
            )

        # Create processor config, propagating all relevant settings
        wilor_config = WiLORConfig(
            checkpoint_path=checkpoint,
            cfg_path=config,
            detector_path=detector,
            img_batch_size=self.config.img_batch_size,
            det_batch_size=self.config.det_batch_size,
            use_compile=self.config.use_compile,
            use_amp=self.config.use_amp,
        )

        # Initialize processor
        import torch

        device = torch.device(self.config.device)
        self._processor = WiLORProcessorOptimized(wilor_config, device)
        self._initialized = True

    def extract_from_video(
        self,
        video_path: str | Path,
        output_path: str | Path | None = None,
    ) -> ExtractionResult:
        """Extract hand poses from video using GPU pipeline.

        Processes frames in batches of ``img_batch_size`` to stay within GPU
        memory on consumer cards (e.g. RTX 3090 with 24 GB).
        """
        if not self._initialized:
            self.load_model()

        video_path = Path(video_path)
        info = get_video_info(video_path)

        try:
            import torch

            from sltk.extraction.base import read_video_frames

            # Read all frames on CPU (numpy, ~2 GB for a typical video)
            all_frames_np = read_video_frames(video_path)
            total = len(all_frames_np)
            batch_size = self.config.img_batch_size

            # Accumulate results across batches
            all_results: list[dict] = []

            for start in range(0, total, batch_size):
                end = min(start + batch_size, total)
                chunk_np = all_frames_np[start:end]
                chunk_gpu = torch.from_numpy(chunk_np).to(self.config.device)

                batch_result = self._processor.process_frames(
                    chunk_gpu,
                    base_frame_idx=start,
                )

                # Free GPU memory immediately
                del chunk_gpu
                torch.cuda.empty_cache()

                if batch_result is not None:
                    all_results.append(batch_result)

            # Merge batch results
            results = self._merge_batch_results(all_results)

            # Convert to PoseSequence
            poses = self._results_to_poses(results, info["fps"], info["num_frames"])

            result = ExtractionResult(
                poses=poses,
                num_frames=info["num_frames"],
                num_detections=results.get("num_detections", 0),
                fps=info["fps"],
                metadata={
                    "source_video": str(video_path),
                    "video_info": info,
                    "extractor": self.name,
                },
            )

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

            return result

        except Exception as e:
            return ExtractionResult(
                errors=[str(e)],
                fps=info.get("fps", 25.0),
            )

    @staticmethod
    def _merge_batch_results(batch_results: list[dict]) -> dict:
        """Merge results from multiple frame batches into a single dict."""
        if not batch_results:
            return {
                "num_detections": 0,
                "kpts_2d": np.empty((0, 21, 2)),
                "kpts_3d": np.empty((0, 21, 3)),
                "frame_idx": np.empty((0,), dtype=np.int32),
                "right": np.empty((0,), dtype=np.float32),
            }

        merged: dict = {}
        # Concatenate all array-valued keys
        for key in batch_results[0]:
            vals = [r[key] for r in batch_results if key in r]
            if isinstance(vals[0], np.ndarray):
                merged[key] = np.concatenate(vals, axis=0)
            elif isinstance(vals[0], (int, float)):
                merged[key] = sum(vals)
            elif isinstance(vals[0], dict):
                # Recursively merge nested dicts (e.g. mano params)
                sub_merged: dict = {}
                for sub_key in vals[0]:
                    sub_vals = [v[sub_key] for v in vals if sub_key in v]
                    if isinstance(sub_vals[0], np.ndarray):
                        sub_merged[sub_key] = np.concatenate(sub_vals, axis=0)
                    else:
                        sub_merged[sub_key] = sub_vals[-1]
                merged[key] = sub_merged
            else:
                merged[key] = vals[-1]
        return merged

    def _decode_video_gpu(self, video_path: Path) -> torch.Tensor:
        """Decode video using GPU acceleration."""
        try:
            import decord

            decord.bridge.set_bridge("torch")

            device_id = 0
            if ":" in self.config.device:
                try:
                    device_id = int(self.config.device.split(":")[-1])
                except ValueError:
                    device_id = 0

            vr = decord.VideoReader(
                str(video_path),
                ctx=decord.gpu(device_id),
            )
            frames = vr.get_batch(range(len(vr)))
            return frames

        except Exception as e:
            # Fallback to CPU decoding
            # Note: decord raises DECORDError which isn't a standard exception type
            import logging

            import torch

            logger = logging.getLogger(__name__)
            logger.warning(f"GPU video decode failed, using CPU fallback: {e}")

            from sltk.extraction.base import read_video_frames

            frames = read_video_frames(video_path)
            return torch.from_numpy(frames).to(self.config.device)

    def _results_to_poses(
        self,
        results: dict[str, Any],
        fps: float,
        num_frames: int,
    ) -> PoseSequence:
        """Convert WiLoR results to PoseSequence."""
        # WiLoR outputs sparse detections with frame_idx
        # Convert to dense (T, N, C) format for PoseSequence

        kpts_3d = results.get("kpts_3d", np.array([]))
        frame_idx = results.get("frame_idx", np.array([]))
        right = results.get("right", np.array([]))

        # Initialize dense arrays (T, 2*21, 3) for left and right hands
        data = np.zeros((num_frames, 42, 3), dtype=np.float32)
        confidence = np.zeros((num_frames, 42), dtype=np.float32)

        # Fill in detected poses
        for frame in range(num_frames):
            if frame < len(frame_idx):
                start_idx, count = frame_idx[frame]
                for i in range(count):
                    det_idx = start_idx + i
                    if det_idx < len(kpts_3d):
                        is_right = right[det_idx] if det_idx < len(right) else True
                        offset = 21 if is_right else 0
                        data[frame, offset : offset + 21] = kpts_3d[det_idx]
                        confidence[frame, offset : offset + 21] = 1.0

        return PoseSequence(
            data=data,
            fps=fps,
            format="wilor",
            confidence=confidence,
            metadata={
                "num_keypoints": 21,
                "keypoint_order": ["left_hand", "right_hand"],
                "sparse_results": {
                    "num_detections": len(kpts_3d),
                    "frame_idx": frame_idx,
                },
            },
        )

    def _get_compression_kwargs(self, compress: bool) -> dict:
        """Get HDF5 compression kwargs compatible with h5py.create_dataset."""
        if not compress:
            return {}
        try:
            import hdf5plugin

            lz4_filter = hdf5plugin.LZ4()
            return {
                "compression": lz4_filter.filter_id,
                "compression_opts": lz4_filter.filter_options,
            }
        except (ImportError, AttributeError):
            # Fallback to gzip if hdf5plugin not available
            return {"compression": "gzip", "compression_opts": 4}

    def _save_wilor_results(
        self,
        results: dict[str, Any],
        output_path: str | Path,
        video_info: dict[str, Any],
    ) -> None:
        """Save WiLoR results in native H5 format."""
        try:
            import h5py
        except ImportError:
            raise ImportError("h5py required for saving")

        output_path = Path(output_path)
        compression_kwargs = self._get_compression_kwargs(self.config.compress)

        with h5py.File(output_path, "w") as f:
            # Save detection data
            for key, value in results.items():
                if isinstance(value, np.ndarray):
                    f.create_dataset(key, data=value, **compression_kwargs)
                elif isinstance(value, dict):
                    grp = f.create_group(key)
                    for k, v in value.items():
                        if isinstance(v, np.ndarray):
                            grp.create_dataset(k, data=v, **compression_kwargs)

            # Metadata
            f.attrs["fps"] = video_info["fps"]
            f.attrs["num_frames"] = video_info["num_frames"]
            f.attrs["resolution"] = (video_info["width"], video_info["height"])
            f.attrs["extractor"] = self.name

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

        import torch

        frames_tensor = torch.from_numpy(frames).to(self.config.device)

        results = self._processor.process_frames(frames_tensor, base_frame_idx=0)
        poses = self._results_to_poses(results, fps, len(frames))

        return ExtractionResult(
            poses=poses,
            num_frames=len(frames),
            num_detections=results.get("num_detections", 0),
            fps=fps,
        )

    def close(self) -> None:
        """Release resources and GPU memory."""
        if hasattr(self, "_processor") and self._processor is not None:
            # Move model to CPU before deletion to free GPU memory
            try:
                import gc

                import torch

                # Try to move model components to CPU first
                if hasattr(self._processor, "model") and self._processor.model is not None:
                    self._processor.model.cpu()
                if hasattr(self._processor, "detector") and self._processor.detector is not None:
                    self._processor.detector.cpu()

                del self._processor
                self._processor = None

                # Force garbage collection
                gc.collect()

                # Synchronize and clear CUDA cache
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                    torch.cuda.empty_cache()
            except ImportError:
                del self._processor
                self._processor = None
            except Exception:
                # If cleanup fails, still try to delete processor
                self._processor = None
        self._initialized = False

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

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - release resources."""
        self.close()
        return False


# Convenience function
def extract_wilor(
    video_path: str | Path,
    checkpoint_path: str,
    output_path: str | Path | None = None,
    detector_path: str | None = None,
) -> ExtractionResult:
    """
    Quick extraction using WiLoR.

    Args:
        video_path: Input video path
        checkpoint_path: Path to WiLoR checkpoint
        output_path: Optional output H5 path
        detector_path: Optional path to YOLO hand detector

    Returns:
        ExtractionResult with hand poses
    """
    config = WiLoRConfig(
        checkpoint_path=checkpoint_path,
        detector_path=detector_path,
    )
    with WiLoRExtractor(config) as extractor:
        return extractor.extract_from_video(video_path, output_path)
