"""
NLF (Neural Lifting Field) SMPLX pose extraction wrapper.

NLF provides full-body SMPLX parameters including body, hands, and face.
The NLF processor is bundled with SLTK in the nlf_deps subpackage.

Usage:
    from sltk.extraction.nlf import NLFExtractor, NLFConfig

    config = NLFConfig(
        model_path="path/to/nlf_model.torchscript",
        smplx_path="path/to/SMPLX_NEUTRAL.npz",
    )
    extractor = NLFExtractor(config)
    extractor.load_model()
    result = extractor.extract_from_video("video.mp4")
"""

from __future__ import annotations

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 NLFConfig(ExtractionConfig):
    """NLF-specific configuration."""

    # Model paths
    model_path: str | None = None
    smplx_path: str | None = None
    detector_path: str | None = None

    # Processing parameters
    detection_confidence: float = 0.05
    img_batch_size: int = 64
    det_batch_size: int = 64

    # SMPLX parameters
    num_betas: int = 10
    fov: float = 60.0

    # GPU acceleration
    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

        # Resolve model path
        if self.model_path is None:
            bundled = get_weight_path("nlf_model")
            if bundled:
                self.model_path = str(bundled)

        # Resolve SMPLX path
        if self.smplx_path is None:
            bundled = get_weight_path("nlf_smplx")
            if bundled:
                self.smplx_path = str(bundled)

        # Resolve detector path
        if self.detector_path is None:
            bundled = get_weight_path("nlf_detector")
            if bundled:
                self.detector_path = str(bundled)


class NLFExtractor(PoseExtractor):
    """
    NLF full-body pose extractor.

    Extracts 55 SMPLX keypoints (body + hands + face) with full
    SMPLX parameters including shape and pose.
    """

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

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

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

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

        Requires the signmunculusisambard NLF processor.
        """
        model_path = kwargs.get("model_path", self.config.model_path)
        smplx_path = kwargs.get("smplx_path", self.config.smplx_path)
        detector_path = kwargs.get("detector_path", self.config.detector_path)

        if model_path is None:
            raise ConfigurationError(
                "NLF model_path required. Set via NLFConfig or load_model(model_path=...)"
            )

        # Import from bundled nlf_deps (no external dependencies)
        try:
            from sltk.extraction.nlf_deps import (
                NLFProcessorConfig,
                NLFProcessorOptimized,
            )
        except ImportError as e:
            raise ImportError(
                f"NLF processor not found. This is an internal error. "
                f"Please reinstall SLTK: pip install -e '.[all]'. Error: {e}"
            )

        # Create processor config (use bundled field names)
        nlf_config = NLFProcessorConfig(
            nlf_model_path=model_path,
            smplx_model_path=smplx_path if smplx_path else "models/SMPLX_NEUTRAL.npz",
            yolo_model_path=detector_path if detector_path else "pretrained_models/yolov8x.pt",
            detection_confidence=self.config.detection_confidence,
            img_batch_size=self.config.img_batch_size,
            det_batch_size=self.config.det_batch_size,
            use_amp=self.config.use_amp,
            fov=self.config.fov,
            num_betas=self.config.num_betas,
        )

        # Initialize processor
        import torch

        device = torch.device(self.config.device)
        self._processor = NLFProcessorOptimized(nlf_config, device)
        self._initialized = True

    def extract_from_video(
        self,
        video_path: str | Path,
        output_path: str | Path | None = None,
    ) -> ExtractionResult:
        """Extract full-body poses from video."""
        if not self._initialized:
            self.load_model()

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

        try:
            import torch

            # Use GPU video decoding if available
            if self.config.use_gpu_decode:
                frames = self._decode_video_gpu(video_path)
            else:
                from sltk.extraction.base import read_video_frames

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

            # Process frames
            results = self._processor.process_frames(frames, base_frame_idx=0)

            # 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_nlf_results(results, output_path, info)

            return result

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

    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 NLF results to PoseSequence."""
        # NLF outputs sparse detections with frame_idx
        joints3d = results.get("joints3d", np.array([]))
        frame_idx = results.get("frame_idx", np.array([]))

        # Initialize dense arrays (T, 55, 3) for SMPLX joints
        data = np.zeros((num_frames, 55, 3), dtype=np.float32)
        confidence = np.zeros((num_frames, 55), dtype=np.float32)

        # Fill in detected poses (take first detection per frame)
        for frame in range(num_frames):
            if frame < len(frame_idx):
                start_idx, count = frame_idx[frame]
                if count > 0 and start_idx < len(joints3d):
                    data[frame] = joints3d[start_idx]
                    conf = results.get("joint_uncertainties", None)
                    if conf is not None and start_idx < len(conf):
                        confidence[frame] = 1.0 - np.clip(conf[start_idx], 0, 1)
                    else:
                        confidence[frame] = 1.0

        return PoseSequence(
            data=data,
            fps=fps,
            format="smplx",
            confidence=confidence,
            metadata={
                "num_keypoints": 55,
                "keypoint_format": "smplx_55",
                "sparse_results": {
                    "num_detections": len(joints3d),
                    "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_nlf_results(
        self,
        results: dict[str, Any],
        output_path: str | Path,
        video_info: dict[str, Any],
    ) -> None:
        """Save NLF 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 main arrays
            for key in [
                "joints3d",
                "joints2d",
                "joint_uncertainties",
                "img_idx",
                "frame_idx",
                "boxes",
            ]:
                if key in results and results[key] is not None:
                    f.create_dataset(key, data=results[key], **compression_kwargs)

            # Save SMPLX parameters
            if "smplx" in results:
                grp = f.create_group("smplx")
                for key, value in results["smplx"].items():
                    if isinstance(value, np.ndarray):
                        grp.create_dataset(key, data=value, **compression_kwargs)

            # Metadata
            f.attrs["fps"] = video_info["fps"]
            f.attrs["num_frames"] = video_info["num_frames"]
            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 get_smplx_params(
        self,
        video_path: str | Path,
    ) -> dict[str, np.ndarray]:
        """
        Extract full SMPLX parameters (not just keypoints).

        Returns:
            Dictionary with SMPLX parameters:
            - global_orient: (M, 3) root rotation
            - body_pose: (M, 63) body joint rotations
            - left_hand_pose: (M, 45) left hand rotations
            - right_hand_pose: (M, 45) right hand rotations
            - betas: (M, 10) shape parameters
            - trans: (M, 3) translation
        """
        if not self._initialized:
            self.load_model()

        import torch

        video_path = Path(video_path)
        if self.config.use_gpu_decode:
            frames = self._decode_video_gpu(video_path)
        else:
            from sltk.extraction.base import read_video_frames

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

        results = self._processor.process_frames(frames, base_frame_idx=0)

        return results.get("smplx", {})

    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_nlf(
    video_path: str | Path,
    model_path: str,
    output_path: str | Path | None = None,
    smplx_path: str | None = None,
) -> ExtractionResult:
    """
    Quick extraction using NLF.

    Args:
        video_path: Input video path
        model_path: Path to NLF TorchScript model
        output_path: Optional output H5 path
        smplx_path: Optional path to SMPLX neutral model

    Returns:
        ExtractionResult with full-body poses
    """
    config = NLFConfig(
        model_path=model_path,
        smplx_path=smplx_path,
    )
    with NLFExtractor(config) as extractor:
        return extractor.extract_from_video(video_path, output_path)
