"""
Segment Anything Model 3 (SAM3) inference capability.
"""

import logging
import os
import time
import uuid
from typing import Any, Dict, List, NoReturn, Optional, Tuple, cast

import numpy as np
from PIL import Image
from pydantic import BaseModel, ConfigDict, Field
from shapely.geometry import MultiPolygon, Point, Polygon

from aiko_services import StreamEvent
from highlighter.agent.capabilities.base_capability import Capability
from highlighter.client import Annotation, Entity, Observation
from highlighter.core.data_models import DataSample
from highlighter.core.enums import ContentTypeEnum


class SAM3InitialisationError(Exception):
    """
    Raised when the SAM3 capability fails to provision model weights, or to
    initialise those weights.
    """

    pass


try:
    from sam3.model.sam3_image_processor import Sam3Processor  # type: ignore
    from sam3.model_builder import build_sam3_image_model  # type: ignore
except ImportError as e:
    raise SAM3InitialisationError("SAM3 libraries missing.") from e

logger = logging.getLogger(__name__)


class SAM3TextPrompt(BaseModel):
    """
    Static text prompt configuration.
    """

    prompt: str = Field(..., description="Text concept to segment (e.g., 'dog').")
    object_class_id: uuid.UUID = Field(..., description="Object class UUID associated with concept.")


class SAM3Prompts(BaseModel):
    """
    Validation schema for SAM3 static inference prompts.

    Note: static geometric prompts (points and boxes) are forbidden in pipeline
    definition to prevent blind, frame-agnostic prompting. Geometric prompts
    must be dynamic, and are provided via entities emitted from an upstream
    capability.
    """

    model_config = ConfigDict(extra="forbid")

    text: List[SAM3TextPrompt] = Field(
        default_factory=list,
        description="List of text concepts (strings) or objects specifying concepts and associated object class IDs.",
    )


class SAM3(Capability):
    """
    Standalone inference capability for Meta's Segment Anything Model 3.

    Supports multimodal prompting (text, points, boxes) for segmentation.

    Usage notes:
    - Text concepts defined statically via `StreamParameters.prompts.text`.
    - Geometric prompts (points, boxes) provided dynamically via upstream entities.
    - All prompts sharing same object class ID batched into single inference.
    - Points internally converted to 2x2 boxes for grounding compatibility.

    Upstream entity requirements:
    - Object class `Observation`.
    - `Annotation` with box or point geometry.
    """

    # multi-use string constants
    PROMPT_TYPE_TEXT = "text"
    PROMPT_TYPE_POINTS = "points"
    PROMPT_TYPE_BOXES = "boxes"

    RESULT_KEY_POLYGON = "polygon"
    RESULT_KEY_SCORE = "score"
    RESULT_KEY_LABEL = "label"
    RESULT_KEY_OBJECT_CLASS_ID = "object_class_id"

    PROMPT_DATA_KEY_TEXT = "text"
    PROMPT_DATA_KEY_BOXES = "boxes"
    PROMPT_DATA_KEY_LABELS = "labels"

    # multi-use numeric constants
    LABEL_POSITIVE = 1
    LABEL_NEGATIVE = 0
    PROCESSOR_RESOLUTION = 1008

    class InitParameters(Capability.InitParameters):
        """
        Initialisation parameters for SAM3.
        """

        model_config = ConfigDict(extra="ignore")

        model_path: Optional[str] = None
        training_run_artefact_id: Optional[uuid.UUID] = None

    class StreamParameters(Capability.StreamParameters):
        """
        Runtime stream parameters for SAM3.
        """

        model_config = ConfigDict(extra="ignore")

        confidence_threshold: float = Field(
            default=0.4, description="Minimum confidence score for detections."
        )
        nms_iou_threshold: float = Field(
            default=0.7,
            description="Mask IoU threshold for non-maximum suppression (discards lower-confidence mask in any pair with IoU above this value).",
        )
        max_detections: int = Field(
            default=300,
            description="Maximum number of masks to return per frame.",
        )
        max_image_size: int = Field(
            default=1024,
            description="Input images with longest side exceeding this value are downscaled before inference to limit GPU memory usage.",
        )
        prompts: SAM3Prompts = Field(
            default_factory=SAM3Prompts,
            description="Static text prompts applied to every frame (text).",
        )

    def __init__(self, context: Any) -> None:
        """
        Initialises SAM3 capability and its predictor.
        """
        super().__init__(context)

        self._model, self._processor = self._setup_predictor()
        self.device = self._model.device

    def _setup_predictor(self) -> Tuple[Any, Any]:
        """
        Initialises and warms up SAM3 model and processor.
        """
        params = cast(SAM3.InitParameters, self.init_parameters)
        if params.model_path and params.training_run_artefact_id:
            msg = "Cannot specify both 'model_path' and 'training_run_artefact_id' simultaneously."
            logger.error(msg)
            raise SAM3InitialisationError(msg)

        # select between HL artefact and explicit path for weights
        if params.training_run_artefact_id:
            from highlighter.licenses.sam3_license import require_acceptance

            require_acceptance()

            from highlighter.client.gql_client import HLClient
            from highlighter.client.training_runs import TrainingRunArtefactType

            training_run_artefact = HLClient.get_client().trainingRunArtefact(
                return_type=TrainingRunArtefactType,
                id=params.training_run_artefact_id,
            )
            path = training_run_artefact.file_url
        elif params.model_path is not None:
            path = params.model_path
        else:
            msg = "Path to model weights could not be resolved from parameters."
            logger.error(msg)
            raise SAM3InitialisationError(msg)

        try:
            model = build_sam3_image_model(checkpoint_path=path)
            processor = self._create_processor(model)

            logger.info(f"Initialising model weights '{os.path.basename(path)}' on {model.device}.")

            # Warm up: use processor resolution so CUDA kernels are compiled
            # at the same tensor shapes used during real inference.
            res = SAM3.PROCESSOR_RESOLUTION
            warmup_img = Image.fromarray(np.zeros((res, res, 3), dtype=np.uint8))
            processor.set_image(warmup_img)

            return model, processor
        except Exception as e:
            self._handle_fatal_error(e, "model load")

    @staticmethod
    def _create_processor(model: Any, resolution: Optional[int] = None) -> Any:
        """
        Creates a Sam3Processor configured to run on the model's device.
        """
        if resolution is None:
            resolution = SAM3.PROCESSOR_RESOLUTION
        return Sam3Processor(model, resolution=resolution, device=str(model.device))

    def _handle_fatal_error(self, error: Exception, context: str = "initialisation") -> NoReturn:
        """
        Logs fatal error and raises SAM3InitialisationError.
        """
        msg = f"Fatal {context} error: {error}."
        logger.error(msg)
        raise SAM3InitialisationError(msg) from error

    @staticmethod
    def extract_prompts_from_entities(
        pos: Dict[str, Entity],
        neg: Dict[str, Entity],
        sample: DataSample,
    ) -> Dict[str, List[Any]]:
        """
        Extracts multimodal prompts from strict positive and negative entity sets.
        """
        prompts: Dict[str, List[Any]] = {
            SAM3.PROMPT_TYPE_TEXT: [],
            SAM3.PROMPT_TYPE_POINTS: [],
            SAM3.PROMPT_TYPE_BOXES: [],
        }

        def _extract(entities: Dict[str, Entity], label: int):
            for e in entities.values():
                # resolve class id
                cid = None
                for o in e.all_observations():
                    try:
                        # o.value could be a LabeledUUID or a string uuid
                        val = o.value
                        new_cid = val if isinstance(val, uuid.UUID) else uuid.UUID(str(val))
                    except (ValueError, TypeError):
                        logger.warning(
                            f"Entity {e.id} has invalid class ID '{o.value}'. "
                            "Object class observation value must be a UUID."
                        )
                        continue

                    if cid and cid != new_cid:
                        logger.warning(
                            f"Entity {e.id} has multiple class observations. "
                            f"Using '{cid}', ignoring '{new_cid}'."
                        )
                        continue
                    cid = new_cid

                if not cid:
                    continue

                # visual exemplars (bounding boxes, points)
                for a in e.annotations:
                    if a.location and a.data_sample == sample:
                        if isinstance(a.location, Point):
                            prompts[SAM3.PROMPT_TYPE_POINTS].append(
                                ([a.location.x, a.location.y, label], str(cid))
                            )
                        elif isinstance(a.location, (Polygon, MultiPolygon)):
                            b = list(a.location.bounds)
                            prompts[SAM3.PROMPT_TYPE_BOXES].append(
                                ([b[0], b[1], b[2], b[3], label], str(cid))
                            )

        _extract(pos, SAM3.LABEL_POSITIVE)
        _extract(neg, SAM3.LABEL_NEGATIVE)
        return prompts

    @staticmethod
    def _mask_to_polygon(mask: np.ndarray, approx_level: float = 0.005) -> Polygon:
        """
        Converts a boolean mask to a Shapely Polygon using the largest contour.

        SAM3 returns one mask per detection, so multi-component masks represent
        disconnected visible regions of a single entity (e.g. partial occlusion).
        We keep only the largest contour because Annotation.from_points requires
        a single Polygon.  Discarded components are logged as a warning.
        """
        import cv2

        binary = mask.astype(np.uint8) * 255
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            raise ValueError("Mask produced no contours")

        if len(contours) > 1:
            areas = [cv2.contourArea(c) for c in contours]
            total = sum(areas)
            largest_idx = int(np.argmax(areas))
            kept_pct = areas[largest_idx] / total * 100 if total > 0 else 100
            logger.warning(
                f"Mask has {len(contours)} disconnected components; "
                f"keeping largest ({kept_pct:.0f}% of mask area), discarding {len(contours) - 1}."
            )
            cnt = contours[largest_idx]
        else:
            cnt = contours[0]

        epsilon = approx_level * cv2.arcLength(cnt, True)
        cnt = cv2.approxPolyDP(cnt, epsilon, True)
        if cnt.shape[0] < 3:
            raise ValueError("Contour has fewer than 3 points after approximation")

        pts = np.squeeze(cnt)
        coords = [(int(x), int(y)) for x, y in pts]
        coords.append(coords[0])
        location = Polygon(coords)
        if not location.is_valid:
            location = location.buffer(0)
        if location.is_empty:
            raise ValueError("Mask produced empty geometry after repair")
        if isinstance(location, MultiPolygon):
            location = max(location.geoms, key=lambda g: g.area)
        return location

    def _create_entity_from_detection(self, result: Dict[str, Any], sample: DataSample) -> Entity:
        """
        Creates entity with annotation and object class observation from detection.
        """
        anno = Annotation.from_points(
            list(result[self.RESULT_KEY_POLYGON].exterior.coords),
            confidence=result[self.RESULT_KEY_SCORE],
            data_sample=sample,
        )

        object_class_id = result[self.RESULT_KEY_OBJECT_CLASS_ID]
        try:
            if not isinstance(object_class_id, uuid.UUID):
                object_class_id = uuid.UUID(str(object_class_id))
        except (ValueError, TypeError) as e:
            self._handle_fatal_error(e, f"invalid object class ID '{object_class_id}'")

        obs = Observation.make_object_class_observation(
            object_class_id,
            result[self.RESULT_KEY_LABEL],
            result[self.RESULT_KEY_SCORE],
            sample.recorded_at,
            frame_id=sample.media_frame_index,
        )
        anno.observations.add(obs)
        return anno.get_or_create_entity()

    def process_frame(
        self,
        stream: Any,
        data_samples: List[DataSample],
        *args: Any,
        **kwargs: Any,
    ) -> Tuple[StreamEvent, Dict[str, Any]]:
        """
        Processes single frame through SAM3 model.
        """
        # 1. strictly extract prompt entities from dedicated pins
        pos = kwargs.get("positive_prompt_entities", {})
        neg = kwargs.get("negative_prompt_entities", {})

        # 2. initialise output state
        new_ent = {}

        p = cast(SAM3.StreamParameters, self.stream_parameters(stream.stream_id))

        for s in data_samples:
            try:
                content = s.content
                if content is None:
                    raise ValueError(f"Data sample {s} has no content.")

                if s.content_type != ContentTypeEnum.IMAGE:
                    continue

                # convert to PIL image
                if isinstance(content, Image.Image):
                    img = content
                else:
                    img = Image.fromarray(content)  # type: ignore

                # extract multimodal prompts aligned on data sample
                all_prompts = self.extract_prompts_from_entities(pos, neg, s)

                # add static text prompts
                for t in p.prompts.text:
                    if isinstance(t, SAM3TextPrompt):
                        all_prompts[self.PROMPT_TYPE_TEXT].append((t.prompt, str(t.object_class_id)))
                    else:
                        raise ValueError(f"Invalid static prompt: {t}. Expected SAM3TextPrompt.")

                if not any(all_prompts.values()):
                    continue

                # run inference
                results = self._predict(img, all_prompts, p)

                # create annotations
                for r in results:
                    try:
                        entity = self._create_entity_from_detection(r, s)
                    except Exception as det_err:
                        logger.warning(
                            f"Skipping detection (score={r.get(self.RESULT_KEY_SCORE):.3f}): {det_err}"
                        )
                        continue
                    new_ent[str(entity.id)] = entity

            except Exception as e:
                return StreamEvent.ERROR, {"diagnostic": str(e)}  # type: ignore

        return StreamEvent.OKAY, {"entities": new_ent}  # type: ignore

    def _group_prompts_by_class(self, prompts: Dict[str, List[Any]]) -> Dict[str, Dict[str, Any]]:
        """
        Groups multimodal prompts by object class UUID.

        Returns dict mapping class UUID to
        {PROMPT_DATA_KEY_TEXT: str|None, PROMPT_DATA_KEY_BOXES: [...], PROMPT_DATA_KEY_LABELS: [...]}.
        """
        by_class: Dict[str, Dict[str, Any]] = {}

        def _ensure_class(cls_uuid: str) -> None:
            if cls_uuid not in by_class:
                by_class[cls_uuid] = {
                    self.PROMPT_DATA_KEY_TEXT: None,
                    self.PROMPT_DATA_KEY_BOXES: [],
                    self.PROMPT_DATA_KEY_LABELS: [],
                }

        # text prompts — collect unique texts per class, then join
        texts_by_class: Dict[str, List[str]] = {}
        for text, cls_uuid in prompts[self.PROMPT_TYPE_TEXT]:
            _ensure_class(cls_uuid)
            if cls_uuid not in texts_by_class:
                texts_by_class[cls_uuid] = []
            if text not in texts_by_class[cls_uuid]:
                texts_by_class[cls_uuid].append(text)
        for cls_uuid, texts in texts_by_class.items():
            if texts:
                by_class[cls_uuid][self.PROMPT_DATA_KEY_TEXT] = ", ".join(texts)

        # point prompts — convert to 2x2 boxes for grounding
        for pt, cls_uuid in prompts[self.PROMPT_TYPE_POINTS]:
            _ensure_class(cls_uuid)
            x, y = pt[:2]
            label_val = int(pt[2]) if len(pt) > 2 else self.LABEL_POSITIVE
            by_class[cls_uuid][self.PROMPT_DATA_KEY_BOXES].append([x - 1, y - 1, x + 1, y + 1])
            by_class[cls_uuid][self.PROMPT_DATA_KEY_LABELS].append(label_val)

        # box prompts
        for box, cls_uuid in prompts[self.PROMPT_TYPE_BOXES]:
            _ensure_class(cls_uuid)
            b = box[:4]
            label_val = int(box[4]) if len(box) > 4 else self.LABEL_POSITIVE
            by_class[cls_uuid][self.PROMPT_DATA_KEY_BOXES].append(b)
            by_class[cls_uuid][self.PROMPT_DATA_KEY_LABELS].append(label_val)

        return by_class

    def _predict(self, img: Image.Image, prompts: Dict[str, List[Any]], p: Any) -> List[Dict[str, Any]]:
        """
        Executes SAM3 inference with multimodal prompt grouping.
        """
        t0 = time.monotonic()
        results = []

        orig_w, orig_h = img.size

        # downscale large images to fit GPU memory
        longest = max(orig_w, orig_h)
        if longest > p.max_image_size:
            scale = p.max_image_size / longest
            img = img.resize((int(orig_w * scale), int(orig_h * scale)), Image.LANCZOS)

        w, h = img.size

        state = self._processor.set_image(img)
        self._processor.set_confidence_threshold(p.confidence_threshold, state)

        try:
            by_class = self._group_prompts_by_class(prompts)

            # Run group inference.  Sam3Processor accumulates prompts in `state`;
            # each call returns the cumulative result, so we only need the final one.
            for cls_uuid, data in by_class.items():
                self._processor.reset_all_prompts(state)

                label_visual = "visual"
                label_key = (
                    data[self.PROMPT_DATA_KEY_TEXT] if data[self.PROMPT_DATA_KEY_TEXT] else label_visual
                )
                result = None

                if data[self.PROMPT_DATA_KEY_TEXT]:
                    result = self._processor.set_text_prompt(
                        prompt=data[self.PROMPT_DATA_KEY_TEXT],
                        state=state,
                    )

                for box_px, label_val in zip(
                    data[self.PROMPT_DATA_KEY_BOXES], data[self.PROMPT_DATA_KEY_LABELS]
                ):
                    normalized_box = self._normalise_box_prompt(box_px, orig_w, orig_h)
                    positive = label_val == self.LABEL_POSITIVE
                    result = self._processor.add_geometric_prompt(
                        box=normalized_box,
                        label=positive,
                        state=state,
                    )

                if result is None:
                    continue

                parsed = self._parse_results(
                    result, p.confidence_threshold, p.nms_iou_threshold, p.max_detections, label_key, cls_uuid
                )
                results.extend(parsed)

            # scale polygon coordinates back to original dimensions if we downscaled
            if (orig_w, orig_h) != (w, h):
                from shapely import affinity

                sx = orig_w / w
                sy = orig_h / h
                for r in results:
                    r[self.RESULT_KEY_POLYGON] = affinity.scale(
                        r[self.RESULT_KEY_POLYGON], xfact=sx, yfact=sy, origin=(0, 0)
                    )
        finally:
            # Release backbone feature tensors cached in state to free GPU memory
            # between frames.  The processor itself remains reusable.
            state.clear()

        elapsed_ms = (time.monotonic() - t0) * 1000
        logger.info(
            f"Produced {len(results)} detections in {elapsed_ms:.0f}ms. "
            f"Image: ({orig_w}x{orig_h}), inference: ({w}x{h})."
        )
        return results

    @staticmethod
    def _normalise_box_prompt(box_xyxy: List[float], image_width: int, image_height: int) -> List[float]:
        """
        Converts pixel-space XYXY coordinates to normalised CXCYWH for sam3.
        """
        w = max(image_width, 1)
        h = max(image_height, 1)
        x1, y1, x2, y2 = box_xyxy
        x1 = min(max(float(x1), 0.0), float(w))
        y1 = min(max(float(y1), 0.0), float(h))
        x2 = min(max(float(x2), 0.0), float(w))
        y2 = min(max(float(y2), 0.0), float(h))
        box_width = max(x2 - x1, 0.0)
        box_height = max(y2 - y1, 0.0)
        center_x = x1 + box_width / 2.0
        center_y = y1 + box_height / 2.0
        return [
            center_x / w,
            center_y / h,
            box_width / w,
            box_height / h,
        ]

    def _parse_results(
        self,
        result: Dict[str, Any],
        confidence_threshold: float,
        nms_iou_threshold: float,
        max_detections: int,
        label: str,
        cls_uuid: str,
    ) -> List[Dict[str, Any]]:
        """
        Converts raw sam3 result dict into polygon detections.

        Masks are converted to polygons early so that NMS uses Shapely polygon
        IoU (fast, ~30 vertices) instead of pixel-wise mask IoU (slow, ~1M bools).
        Results are sorted by descending confidence.
        """
        raw_masks = result.get("masks")
        raw_scores = result.get("scores")

        if raw_masks is None or raw_scores is None:
            return []

        # convert masks to polygons, filtering by confidence
        candidates: List[Tuple[float, Polygon]] = []
        for i, score in enumerate(raw_scores):
            score = float(score)
            if score < confidence_threshold:
                continue
            mask = raw_masks[i]
            if hasattr(mask, "cpu"):
                mask = mask.cpu().numpy()
            mask = np.asarray(mask, dtype=bool).squeeze()
            try:
                poly = self._mask_to_polygon(mask)
            except ValueError:
                continue
            candidates.append((score, poly))

        # sort descending by score so NMS keeps highest-confidence detections
        candidates.sort(key=lambda x: x[0], reverse=True)

        # polygon-IoU NMS: suppress lower-confidence polygons that overlap a kept one
        kept: List[Tuple[float, Polygon]] = []
        for score, poly in candidates:
            if any(self._polygon_iou(poly, k_poly) > nms_iou_threshold for _, k_poly in kept):
                continue
            kept.append((score, poly))
            if len(kept) >= max_detections:
                break

        return [
            {
                self.RESULT_KEY_POLYGON: poly,
                self.RESULT_KEY_LABEL: label,
                self.RESULT_KEY_OBJECT_CLASS_ID: cls_uuid,
                self.RESULT_KEY_SCORE: score,
            }
            for score, poly in kept
        ]

    @staticmethod
    def _polygon_iou(a: Polygon, b: Polygon) -> float:
        """Computes IoU between two Shapely polygons."""
        intersection = a.intersection(b).area
        union = a.union(b).area
        return intersection / union if union > 0 else 0.0
