"""TEASER FLAME face parameter extraction wrapper.

TEASER (ICLR 2025) extracts FLAME 3D face parameters — jaw pose, expression,
shape, eyelid, head pose — from video via YOLOv8 face detection and a
multi-branch MobileNetV3 encoder.

Usage::

    from sltk.extraction.teaser import TeaserExtractor, TeaserConfig

    config = TeaserConfig(device="cuda:0")
    with TeaserExtractor(config) as ext:
        ext.load_model()
        result = ext.extract_from_video("video.mp4", "face.h5")
"""

from __future__ import annotations

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

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,
)

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


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


@dataclass
class TeaserConfig(ExtractionConfig):
    """TEASER-specific configuration."""

    # Model paths (auto-resolved from bundled weights if None)
    checkpoint_path: str | None = None
    face_detector_path: str | None = None

    # Processing
    input_image_size: int = 224
    crop_scale: float = 1.4
    img_batch_size: int = 32
    use_amp: bool = True

    def __post_init__(self) -> None:
        """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("teaser_checkpoint")
            if bundled:
                self.checkpoint_path = str(bundled)

        if self.face_detector_path is None:
            bundled = get_weight_path("teaser_face_detector")
            if bundled:
                self.face_detector_path = str(bundled)


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


class TeaserExtractor(PoseExtractor):
    """TEASER face parameter extractor.

    Extracts FLAME parameters (jaw, expression, shape, eyelid) plus
    head pose and camera from video frames.
    """

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

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

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

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

    def load_model(self, **kwargs: Any) -> None:
        """Load YOLOv8 face detector and TeaserEncoder."""
        checkpoint_path = kwargs.get("checkpoint_path", self.config.checkpoint_path)
        face_detector_path = kwargs.get("face_detector_path", self.config.face_detector_path)

        if checkpoint_path is None:
            raise ConfigurationError(
                "TEASER checkpoint_path required. "
                "Set via TeaserConfig or load_model(checkpoint_path=...)"
            )
        if face_detector_path is None:
            raise ConfigurationError(
                "TEASER face_detector_path required. "
                "Set via TeaserConfig or load_model(face_detector_path=...)"
            )

        try:
            from sltk.extraction.teaser_deps import (
                TeaserProcessorConfig,
                TeaserProcessorOptimized,
            )
        except ImportError as e:
            raise ImportError(
                f"TEASER processor not found. This is an internal error. "
                f"Please reinstall SLTK: pip install -e '.[all]'. Error: {e}"
            )

        import torch

        proc_config = TeaserProcessorConfig(
            checkpoint_path=checkpoint_path,
            face_detector_path=face_detector_path,
            input_image_size=self.config.input_image_size,
            crop_scale=self.config.crop_scale,
            detection_confidence=self.config.confidence_threshold,
            use_amp=self.config.use_amp,
        )
        device = torch.device(self.config.device)
        self._processor = TeaserProcessorOptimized(proc_config, device)
        self._initialized = True

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

    def extract_from_video(
        self,
        video_path: str | Path,
        output_path: str | Path | None = None,
    ) -> ExtractionResult:
        """Extract FLAME face parameters from video.

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

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

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

        try:
            import cv2

            cap = cv2.VideoCapture(str(video_path))
            if not cap.isOpened():
                return ExtractionResult(
                    errors=[f"Cannot open video: {video_path}"],
                    fps=info.get("fps", 25.0),
                )

            total_frames = info["num_frames"]
            batch_size = self.config.img_batch_size

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

            for batch_start in range(0, total_frames, batch_size):
                batch_frames: list[np.ndarray] = []
                batch_indices: list[int] = []

                for idx in range(batch_start, min(batch_start + batch_size, total_frames)):
                    ret, frame = cap.read()
                    if not ret:
                        break
                    batch_frames.append(frame)  # BGR uint8
                    batch_indices.append(idx)

                if not batch_frames:
                    break

                batch_result = self._processor.process_frames(batch_frames, batch_indices)
                all_results.append(batch_result)

            cap.release()

            # Merge batches
            merged = self._merge_batch_results(all_results, total_frames)
            poses = self._results_to_poses(merged, info["fps"], total_frames)

            result = ExtractionResult(
                poses=poses,
                num_frames=total_frames,
                num_detections=merged.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_teaser_results(merged, output_path, info)

            return result

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

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

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

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

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

        import cv2

        # Convert RGB → BGR for YOLO
        bgr_frames = [cv2.cvtColor(f, cv2.COLOR_RGB2BGR) for f in frames]
        indices = list(range(len(frames)))

        results = self._processor.process_frames(bgr_frames, indices)
        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,
        )

    # -------------------------------------------------------------------
    # Merge / convert / save helpers
    # -------------------------------------------------------------------

    @staticmethod
    def _merge_batch_results(
        batch_results: list[dict[str, Any]],
        total_frames: int,
    ) -> dict[str, Any]:
        """Merge results from multiple batches into one dict."""
        if len(batch_results) == 0:
            return {
                "flame": {
                    "jaw_pose": np.zeros((0, 3), dtype=np.float32),
                    "expression": np.zeros((0, 50), dtype=np.float32),
                    "eyelid_params": np.zeros((0, 2), dtype=np.float32),
                    "shape": np.zeros((0, 300), dtype=np.float32),
                },
                "pose_params": np.zeros((0, 3), dtype=np.float16),
                "cam": np.zeros((0, 3), dtype=np.float16),
                "boxes": np.zeros((0, 4), dtype=np.float16),
                "boxes_scores": np.zeros((0,), dtype=np.float16),
                "successful_frames": np.array([], dtype=np.int32),
                "frame_idx": np.zeros((total_frames, 2), dtype=np.int32),
                "num_detections": 0,
            }

        if len(batch_results) == 1:
            return batch_results[0]

        # Concatenate detection arrays
        flame_keys = ["jaw_pose", "expression", "eyelid_params", "shape"]
        flat_keys = ["pose_params", "cam", "boxes", "boxes_scores", "successful_frames"]

        merged_flame: dict[str, np.ndarray] = {}
        for k in flame_keys:
            merged_flame[k] = np.concatenate([r["flame"][k] for r in batch_results], axis=0)

        merged: dict[str, Any] = {"flame": merged_flame}
        for k in flat_keys:
            merged[k] = np.concatenate([r[k] for r in batch_results], axis=0)

        # Rebuild frame_idx from scratch using merged successful_frames
        frame_idx = np.zeros((total_frames, 2), dtype=np.int32)
        for det_idx, fid in enumerate(merged["successful_frames"]):
            if 0 <= fid < total_frames:
                frame_idx[fid] = [det_idx, 1]
        merged["frame_idx"] = frame_idx
        merged["num_detections"] = len(merged["successful_frames"])

        return merged

    def _results_to_poses(
        self,
        results: dict[str, Any],
        fps: float,
        num_frames: int,
    ) -> PoseSequence:
        """Convert TEASER results to a PoseSequence for API compatibility.

        We pack the FLAME parameters into a ``(T, K, 1)`` array where ``K``
        is the total parameter dimension (3+50+2+300+3+3 = 361).
        """
        flame = results["flame"]
        M = len(results["successful_frames"])

        if M == 0:
            data = np.zeros((num_frames, 361, 1), dtype=np.float32)
            confidence = np.zeros((num_frames, 361), dtype=np.float32)
        else:
            # Concatenate all params per detection
            params = np.concatenate(
                [
                    flame["jaw_pose"],  # (M, 3)
                    flame["expression"],  # (M, 50)
                    flame["eyelid_params"],  # (M, 2)
                    flame["shape"],  # (M, 300)
                    results["pose_params"].astype(np.float32),  # (M, 3)
                    results["cam"].astype(np.float32),  # (M, 3)
                ],
                axis=1,
            )  # (M, 361)

            # Dense array (T, 361, 1)
            data = np.zeros((num_frames, 361, 1), dtype=np.float32)
            confidence = np.zeros((num_frames, 361), dtype=np.float32)

            frame_idx = results["frame_idx"]
            for frame in range(num_frames):
                if frame < len(frame_idx):
                    start, count = frame_idx[frame]
                    if count > 0 and start < M:
                        data[frame, :, 0] = params[start]
                        confidence[frame] = 1.0

        return PoseSequence(
            data=data,
            fps=fps,
            format="flame",
            confidence=confidence,
            metadata={
                "num_params": 361,
                "param_layout": {
                    "jaw_pose": (0, 3),
                    "expression": (3, 53),
                    "eyelid_params": (53, 55),
                    "shape": (55, 355),
                    "pose_params": (355, 358),
                    "cam": (358, 361),
                },
                "num_detections": M,
            },
        )

    def _get_compression_kwargs(self, compress: bool) -> dict[str, Any]:
        """Get HDF5 compression kwargs."""
        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):
            return {"compression": "gzip", "compression_opts": 4}

    def _save_teaser_results(
        self,
        results: dict[str, Any],
        output_path: str | Path,
        video_info: dict[str, Any],
    ) -> None:
        """Save TEASER results in H5 format."""
        import h5py

        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        ckw = self._get_compression_kwargs(self.config.compress)

        with h5py.File(output_path, "w") as f:
            # FLAME parameter group
            flame_grp = f.create_group("flame")
            for key, value in results["flame"].items():
                flame_grp.create_dataset(key, data=value, **ckw)

            # Top-level arrays
            for key in [
                "pose_params",
                "cam",
                "boxes",
                "boxes_scores",
                "successful_frames",
                "frame_idx",
            ]:
                if key in results and results[key] is not None:
                    f.create_dataset(key, data=results[key], **ckw)

            # Optional FLAME landmarks (computed post-hoc by FLAMEProcessor)
            if "landmarks_2d" in results and results["landmarks_2d"] is not None:
                f.create_dataset("landmarks_2d", data=results["landmarks_2d"], **ckw)

            # Metadata
            f.attrs["fps"] = video_info["fps"]
            f.attrs["num_frames"] = video_info["num_frames"]
            f.attrs["extractor"] = self.name

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

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

                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

                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                    torch.cuda.empty_cache()
            except ImportError:
                del self._processor
                self._processor = None
            except Exception:
                self._processor = None
        self._initialized = False

    def __enter__(self) -> TeaserExtractor:
        """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_teaser(
    video_path: str | Path,
    output_path: str | Path | None = None,
    device: str = "cuda:0",
) -> ExtractionResult:
    """Quick extraction using TEASER.

    Args:
        video_path: Input video path.
        output_path: Optional output H5 path.
        device: Torch device string.

    Returns:
        ExtractionResult with FLAME face parameters.
    """
    config = TeaserConfig(device=device)
    with TeaserExtractor(config) as extractor:
        extractor.load_model()
        return extractor.extract_from_video(video_path, output_path)
