"""SAM 2 video tracker: prompt dispatch, mask propagation, drift detection."""

from __future__ import annotations

from typing import Any

import cv2
import numpy as np
from loguru import logger


class SAMTracker:
    """Wraps SAM 2 video predictor for watermark/object tracking."""

    def __init__(
        self,
        model_size: str = "large",
        device: str = "mps",
        dilation_kernel: int = 15,
        dilation_iterations: int = 2,
        drift_area_ratio_threshold: float = 0.5,
        sam2_cfg_dir: str = "configs/sam2.1",
        checkpoint_dir: str = "checkpoints",
    ) -> None:
        self.model_size = model_size
        self.device = device
        self.dilation_kernel = dilation_kernel
        self.dilation_iterations = dilation_iterations
        self.drift_threshold = drift_area_ratio_threshold
        self.sam2_cfg_dir = sam2_cfg_dir
        self.checkpoint_dir = checkpoint_dir

        self._image_predictor: Any = None
        self._video_predictor: Any = None
        self._grounding_model: Any = None
        self._grounding_processor: Any = None
        self._loaded = False

    def load_model(self) -> None:
        """Load SAM 2 predictor onto device."""
        if self._loaded:
            return

        try:
            from sam2.build_sam import build_sam2, build_sam2_video_predictor
            from sam2.sam2_image_predictor import SAM2ImagePredictor
        except ImportError:
            raise ImportError(
                "segment-anything-2 not installed. "
                "Run: pip install git+https://github.com/facebookresearch/segment-anything-2.git"
            )

        config_map = {
            "tiny": "sam2.1_hiera_t.yaml",
            "small": "sam2.1_hiera_s.yaml",
            "base_plus": "sam2.1_hiera_b+.yaml",
            "large": "sam2.1_hiera_l.yaml",
        }
        cfg_file = f"{self.sam2_cfg_dir}/{config_map.get(self.model_size, 'sam2.1_hiera_l.yaml')}"

        ckpt_map = {
            "tiny": "sam2.1_hiera_tiny.pt",
            "small": "sam2.1_hiera_small.pt",
            "base_plus": "sam2.1_hiera_base_plus.pt",
            "large": "sam2.1_hiera_large.pt",
        }
        ckpt = f"{self.checkpoint_dir}/{ckpt_map.get(self.model_size, 'sam2.1_hiera_large.pt')}"

        sam2_model = build_sam2(cfg_file, ckpt, device=self.device)
        self._image_predictor = SAM2ImagePredictor(sam2_model)

        self._video_predictor = build_sam2_video_predictor(
            cfg_file, ckpt, device=self.device,
        )

        self._loaded = True
        logger.info(f"SAM 2 ({self.model_size}) loaded on {self.device}")

    def _load_grounding_dino(self) -> bool:
        """Load Grounding DINO from HuggingFace transformers. Returns False on failure."""
        if self._grounding_model is not None:
            return True

        try:
            from transformers import AutoModelForZeroShotObjectDetection, AutoProcessor

            model_id = "IDEA-Research/grounding-dino-base"
            self._grounding_processor = AutoProcessor.from_pretrained(model_id)
            self._grounding_model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id)
            # Set to inference mode
            self._grounding_model.requires_grad_(False)
            logger.info("Grounding DINO loaded from HuggingFace transformers (CPU)")
            return True
        except Exception as e:
            logger.warning(f"Grounding DINO load failed: {e}. Text prompts unavailable.")
            return False

    def _text_to_box(self, frame: np.ndarray, prompt: str) -> tuple[int, int, int, int] | None:
        """Use Grounding DINO to convert text prompt to bounding box."""
        if not self._load_grounding_dino():
            return None

        import torch
        from PIL import Image

        image = Image.fromarray(frame)
        inputs = self._grounding_processor(images=image, text=prompt, return_tensors="pt")

        with torch.no_grad():
            outputs = self._grounding_model(**inputs)

        results = self._grounding_processor.post_process_grounded_object_detection(
            outputs,
            inputs["input_ids"],
            threshold=0.3,
            text_threshold=0.25,
            target_sizes=[(frame.shape[0], frame.shape[1])],
        )

        if len(results) == 0 or len(results[0]["boxes"]) == 0:
            logger.warning(f"Grounding DINO found no objects for prompt: '{prompt}'")
            return None

        # Take highest confidence box
        boxes = results[0]["boxes"]
        scores = results[0]["scores"]
        best_idx = scores.argmax().item()
        box = boxes[best_idx].int().tolist()

        logger.info(f"Grounding DINO detected box {box} (score={scores[best_idx]:.3f}) for '{prompt}'")
        return tuple(box)

    def segment_first_frame(
        self,
        frame: np.ndarray,
        prompt: str | None = None,
        points: list[tuple[int, int]] | None = None,
        box: tuple[int, int, int, int] | None = None,
    ) -> np.ndarray:
        """Segment first frame using text/box/point prompt. Returns mask (H, W) bool."""
        if not self._loaded:
            raise RuntimeError("Model not loaded. Call load_model() first.")

        # Text prompt -> box via Grounding DINO
        if prompt and box is None and points is None:
            detected_box = self._text_to_box(frame, prompt)
            if detected_box is not None:
                box = detected_box
            else:
                raise ValueError(
                    f"Cannot detect '{prompt}' in frame. "
                    "Provide box or points prompt instead."
                )

        self._image_predictor.set_image(frame)

        sam_kwargs: dict = {}
        if box is not None:
            sam_kwargs["box"] = np.array(box)
        if points is not None:
            sam_kwargs["point_coords"] = np.array(points)
            sam_kwargs["point_labels"] = np.ones(len(points), dtype=np.int32)

        masks, scores, _ = self._image_predictor.predict(**sam_kwargs, multimask_output=True)

        # Pick best mask
        best_idx = scores.argmax()
        mask = masks[best_idx].astype(bool)

        logger.info(
            f"First frame segmented: mask area={mask.sum()} pixels "
            f"({mask.sum() / mask.size * 100:.1f}%)"
        )
        return mask

    def propagate_masks(
        self,
        video_path: str,
        initial_mask: np.ndarray,
        start_frame: int = 0,
        count: int | None = None,
        frames_for_redetect: list[np.ndarray] | None = None,
        prompt: str | None = None,
        box: tuple[int, int, int, int] | None = None,
        points: list[tuple[int, int]] | None = None,
    ) -> list[np.ndarray]:
        """Propagate mask across frames using SAM 2 video predictor.

        SAM 2 requires video_path (MP4 or JPEG dir). Pass the original
        video file directly — SAM 2 handles its own decoding.

        Args:
            video_path: path to MP4 video or JPEG frame directory
            initial_mask: first frame mask (H, W) bool
            start_frame: not used by SAM 2 (always processes full video)
            count: expected number of output frames
            frames_for_redetect: numpy frames for drift re-detection
            prompt/box/points: for re-detection on drift

        Returns list of masks (H, W) bool, one per frame.
        """
        if not self._loaded:
            raise RuntimeError("Model not loaded. Call load_model() first.")

        import torch

        # Initialize video predictor with video path
        state = self._video_predictor.init_state(video_path=video_path)
        n_frames = state["num_frames"]

        if n_frames == 0:
            return []

        # Add initial mask as conditioning on frame 0
        # SAM 2 expects 2D mask (H, W), not 3D
        mask_tensor = torch.from_numpy(initial_mask.astype(np.float32))
        self._video_predictor.add_new_mask(state, frame_idx=0, obj_id=1, mask=mask_tensor)

        # Propagate
        masks: list[np.ndarray] = []
        prev_area = initial_mask.sum()
        drift_count = 0

        for frame_idx, (out_idx, obj_ids, mask_logits) in enumerate(
            self._video_predictor.propagate_in_video(state)
        ):
            mask = (mask_logits[0, 0] > 0.0).cpu().numpy()
            current_area = mask.sum()

            # Drift detection + re-detection (S5-R1-012)
            if prev_area > 0:
                area_ratio = current_area / prev_area
                if area_ratio < self.drift_threshold or area_ratio > (1.0 / self.drift_threshold):
                    drift_count += 1
                    logger.warning(
                        f"Frame {out_idx}: mask drift detected "
                        f"(area ratio={area_ratio:.2f}, count={drift_count})"
                    )

                    # Re-detection using provided frames
                    if drift_count >= 3 and frames_for_redetect and frame_idx < len(frames_for_redetect):
                        try:
                            new_mask = self.segment_first_frame(
                                frames_for_redetect[frame_idx],
                                prompt=prompt, box=box, points=points,
                            )
                            if new_mask.sum() > 0:
                                mask = new_mask
                                current_area = mask.sum()
                                drift_count = 0
                                logger.info(f"Frame {out_idx}: re-detection successful")

                                new_tensor = torch.from_numpy(mask.astype(np.float32))
                                self._video_predictor.add_new_mask(
                                    state, frame_idx=out_idx, obj_id=1, mask=new_tensor,
                                )
                        except Exception as e:
                            logger.warning(f"Frame {out_idx}: re-detection failed: {e}")
                else:
                    drift_count = 0

            prev_area = max(current_area, 1)
            masks.append(mask)

        # Reset state to free memory
        self._video_predictor.reset_state(state)

        return masks

    def dilate_mask(self, mask: np.ndarray, sensitivity: float | None = None) -> np.ndarray:
        """Dilate mask to cover watermark edges.

        If sensitivity is provided (0.1-1.0), maps to kernel_size + iterations.
        Otherwise uses instance defaults.
        """
        if sensitivity is not None:
            s = max(0.1, min(1.0, sensitivity))
            kernel_size = int(7 + (31 - 7) * s)
            iterations = int(1 + 2 * s)
            # Ensure odd kernel
            if kernel_size % 2 == 0:
                kernel_size += 1
        else:
            kernel_size = self.dilation_kernel
            iterations = self.dilation_iterations

        kernel = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)
        )
        dilated = cv2.dilate(
            mask.astype(np.uint8), kernel, iterations=iterations
        )
        return dilated.astype(bool)

    def release(self) -> None:
        """Release models and GPU memory."""
        self._image_predictor = None
        self._video_predictor = None
        self._grounding_model = None
        self._grounding_processor = None
        self._loaded = False

        try:
            import torch
            torch.mps.synchronize()
            torch.mps.empty_cache()
        except Exception:
            pass

        logger.info("SAM tracker released")
