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

import inspect
import logging
import os
import tempfile
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 ultralytics.models.sam.predict import SAM3SemanticPredictor  # type: ignore
except ImportError as e:
    raise SAM3InitialisationError("Ultralytics SAM3 libraries missing.") from e

logger = logging.getLogger(__name__)


def _apply_ultralytics_patches():
    """
    Applies monkey-patch to bad function call in Ultralytics library (see FIX comment).
    """
    try:
        from ultralytics.models.sam.sam3.decoder import (  # type: ignore
            TransformerDecoder,
        )

        if "_patched" not in TransformerDecoder._get_rpb_matrix.__name__:
            source = inspect.getsource(TransformerDecoder._get_rpb_matrix)
            buggy_code = "self.coord_cache[feat_size] = self._get_coords(H, W, reference_boxes.device)"

            if buggy_code not in source:
                return

            orig_get_rpb_matrix = TransformerDecoder._get_rpb_matrix

            def patched_get_rpb_matrix(self, reference_boxes, feat_size):
                """
                Patched version of _get_rpb_matrix to fix dtype error.
                """
                import numpy as np
                import torch
                from ultralytics.utils.ops import xywh2xyxy

                H, W = feat_size
                boxes_xyxy = xywh2xyxy(reference_boxes).transpose(0, 1)
                bs, num_queries, _ = boxes_xyxy.shape
                if self.compilable_cord_cache is None:
                    self.compilable_cord_cache = self._get_coords(
                        H, W, reference_boxes.device, reference_boxes.dtype
                    )
                    self.compilable_stored_size = (H, W)

                if torch.compiler.is_dynamo_compiling() or self.compilable_stored_size == (H, W):
                    coords_h, coords_w = self.compilable_cord_cache
                else:
                    if feat_size not in self.coord_cache:
                        # FIX: add reference_boxes.dtype to _get_coords call
                        self.coord_cache[feat_size] = self._get_coords(
                            H, W, reference_boxes.device, reference_boxes.dtype
                        )
                    coords_h, coords_w = self.coord_cache[feat_size]

                deltas_y = coords_h.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 1:4:2]
                deltas_y = deltas_y.view(bs, num_queries, -1, 2)
                deltas_x = coords_w.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 0:3:2]
                deltas_x = deltas_x.view(bs, num_queries, -1, 2)

                if self.boxRPB in ["log", "both"]:
                    deltas_x_log = deltas_x * 8
                    deltas_x_log = (
                        torch.sign(deltas_x_log) * torch.log2(torch.abs(deltas_x_log) + 1.0) / np.log2(8)
                    )
                    deltas_y_log = deltas_y * 8
                    deltas_y_log = (
                        torch.sign(deltas_y_log) * torch.log2(torch.abs(deltas_y_log) + 1.0) / np.log2(8)
                    )
                    if self.boxRPB == "log":
                        deltas_x = deltas_x_log
                        deltas_y = deltas_y_log
                    else:
                        deltas_x = torch.cat([deltas_x, deltas_x_log], dim=-1)
                        deltas_y = torch.cat([deltas_y, deltas_y_log], dim=-1)

                deltas_x = self.boxRPB_embed_x(x=deltas_x)
                deltas_y = self.boxRPB_embed_y(x=deltas_y)
                B = deltas_y.unsqueeze(3) + deltas_x.unsqueeze(2)
                B = B.flatten(2, 3)
                B = B.permute(0, 3, 1, 2)
                return B.contiguous()

            patched_get_rpb_matrix.__name__ = f"{orig_get_rpb_matrix.__name__}_patched"
            TransformerDecoder._get_rpb_matrix = patched_get_rpb_matrix

    except Exception as e:
        logger.warning(f"Failed to apply Ultralytics patches: {e}.")


# apply at module level for fast-fail
_apply_ultralytics_patches()


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_MASK = "mask"
    RESULT_KEY_SCORE = "score"
    RESULT_KEY_LABEL = "label"
    RESULT_KEY_OBJECT_CLASS_ID = "object_class_id"

    TIMING_PREPROCESS = "preprocess"
    TIMING_INFERENCE = "inference"
    TIMING_POSTPROCESS = "postprocess"

    SHAPE_UNKNOWN = "unknown"

    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

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

        conf: float = Field(default=0.4, description="Minimum confidence score for detections.")
        iou: float = Field(
            default=0.7,
            description="Intersection over Union (IoU) threshold for non-maximum suppression (NMS; discards lower-confidence mask in any pair of masks with IoU > iou).",
        )
        max_det: int = Field(
            default=300,
            description="Maximum number of masks to return per frame.",
        )
        max_imgsz: int = Field(
            default=1036,
            description="Upper bound for longest dimension of image inputted to model; automatically rounded up to nearest multiple of 14 (SAM3's architectural stride value).",
        )
        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._predictor: SAM3SemanticPredictor = self._setup_predictor()
        self.device = self._predictor.device

    def _setup_predictor(self) -> SAM3SemanticPredictor:
        """
        Initialises and warms up SAM3 predictor.
        """
        # configure logging
        ul_logger = logging.getLogger("ultralytics")
        ul_logger.setLevel(logging.INFO)
        logging.getLogger(__name__)

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

        warmup_imgsz = 644
        cfg = {
            "task": "segment",
            "mode": "predict",
            "model": path,
            "imgsz": warmup_imgsz,
            "agnostic_nms": True,
            "project": os.path.join(tempfile.gettempdir(), "sam3_internal"),
            "name": "inference",
            "save": False,
            "visualize": False,
            "verbose": False,
        }

        try:
            predictor = SAM3SemanticPredictor(overrides=cfg)
            if not getattr(predictor, "model", None):
                predictor.setup_model(None, verbose=False)

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

            # warm up
            warmup_img = np.zeros((warmup_imgsz, warmup_imgsz, 3), dtype=np.uint8)
            predictor.set_image(warmup_img)
            predictor.reset_image()
            return predictor
        except Exception as e:
            self._handle_fatal_error(e, "model load")

    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

    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_mask(
            result[self.RESULT_KEY_MASK],
            confidence=result[self.RESULT_KEY_SCORE],
            data_sample=sample,
            dilate_args=None,
            approx_level=0.005,
        )
        if not anno.location or not anno.location.is_valid:
            raise ValueError(f"Detection produced invalid or missing location: {anno.location}")

        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:
                    entity = self._create_entity_from_detection(r, s)
                    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 _predict(self, img: Image.Image, prompts: Dict[str, List[Any]], p: Any) -> List[Dict[str, Any]]:
        """
        Executes SAM3 inference with multimodal prompt grouping.
        """
        results = []
        timing = {
            self.TIMING_PREPROCESS: 0.0,
            self.TIMING_INFERENCE: 0.0,
            self.TIMING_POSTPROCESS: 0.0,
        }
        img_shape = self.SHAPE_UNKNOWN

        w, h = img.size
        stride = 14
        self._predictor.args.imgsz = ((min(max(w, h), p.max_imgsz) + (stride - 1)) // stride) * stride
        self._predictor.set_image(np.array(img))

        try:
            # group prompts by class id for multimodal inference
            # cid -> {self.PROMPT_DATA_KEY_TEXT: str|None, self.PROMPT_DATA_KEY_BOXES: [], self.PROMPT_DATA_KEY_LABELS: []}
            by_class = {}

            # process text prompts
            for text, cls_uuid in prompts[self.PROMPT_TYPE_TEXT]:
                if cls_uuid not in by_class:
                    by_class[cls_uuid] = {
                        self.PROMPT_DATA_KEY_TEXT: text,
                        self.PROMPT_DATA_KEY_BOXES: [],
                        self.PROMPT_DATA_KEY_LABELS: [],
                    }
                else:
                    # join multiple texts for same class
                    if (
                        by_class[cls_uuid][self.PROMPT_DATA_KEY_TEXT]
                        and by_class[cls_uuid][self.PROMPT_DATA_KEY_TEXT] != text
                    ):
                        by_class[cls_uuid][self.PROMPT_DATA_KEY_TEXT] = (
                            f"{by_class[cls_uuid][self.PROMPT_DATA_KEY_TEXT]}, {text}"
                        )
                    elif not by_class[cls_uuid][self.PROMPT_DATA_KEY_TEXT]:
                        by_class[cls_uuid][self.PROMPT_DATA_KEY_TEXT] = text

            # process point prompts; convert to 2x2 boxes for grounding
            for pt, cls_uuid in prompts[self.PROMPT_TYPE_POINTS]:
                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: [],
                    }
                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)

            # process box prompts
            for box, cls_uuid in prompts[self.PROMPT_TYPE_BOXES]:
                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: [],
                    }
                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)

            # run group inference
            for cls_uuid, data in by_class.items():
                # multimodal: combine text prompt with image exemplars
                text_arg = [data[self.PROMPT_DATA_KEY_TEXT]] if data[self.PROMPT_DATA_KEY_TEXT] else None
                boxes_arg = data[self.PROMPT_DATA_KEY_BOXES] if data[self.PROMPT_DATA_KEY_BOXES] else None
                labels_arg = data[self.PROMPT_DATA_KEY_LABELS] if data[self.PROMPT_DATA_KEY_LABELS] else None

                res = self._predictor(
                    text=text_arg,
                    bboxes=boxes_arg,
                    labels=labels_arg,
                    conf=p.conf,
                    iou=p.iou,
                    max_det=p.max_det,
                    verbose=False,
                )

                # map labels back to class uuid
                label_visual = "visual"
                label_key = (
                    data[self.PROMPT_DATA_KEY_TEXT] if data[self.PROMPT_DATA_KEY_TEXT] else label_visual
                )
                parsed, t, shape = self._parse_results(res, p.conf, {label_key: cls_uuid})

                results.extend(parsed)
                for k in timing:
                    timing[k] += t.get(k, 0.0)
                img_shape = shape

        finally:
            self._predictor.reset_image()

        logger.info(
            f"Produced {len(results)} detections. "
            f"Resolution: {img_shape} -> ({self._predictor.args.imgsz}, {self._predictor.args.imgsz}). "
            f"Speed: {timing[self.TIMING_PREPROCESS]:.1f}ms pre-process, "
            f"{timing[self.TIMING_INFERENCE]:.1f}ms inference, "
            f"{timing[self.TIMING_POSTPROCESS]:.1f}ms post-process."
        )

        return results

    def _parse_results(
        self, results: Any, conf: float, id_map: Dict[str, str]
    ) -> Tuple[List[Dict[str, Any]], Dict[str, float], str]:
        """
        Converts raw results into mask dictionaries.
        """
        parsed = []
        timing = {
            self.TIMING_PREPROCESS: 0.0,
            self.TIMING_INFERENCE: 0.0,
            self.TIMING_POSTPROCESS: 0.0,
        }
        shape = self.SHAPE_UNKNOWN

        for r in results:
            if hasattr(r, "speed"):
                for cat in timing:
                    timing[cat] += r.speed.get(cat, 0.0)

            if hasattr(r, "orig_shape"):
                shape = str(r.orig_shape)

            if r.masks is None:
                continue

            masks = r.masks.data.cpu().numpy().astype(bool)
            for i, box in enumerate(r.boxes):
                score = float(box.conf)
                if score < conf:
                    continue

                label = r.names[int(box.cls)]
                parsed.append(
                    {
                        self.RESULT_KEY_MASK: masks[i],
                        self.RESULT_KEY_LABEL: label,
                        self.RESULT_KEY_OBJECT_CLASS_ID: id_map.get(label) or label,
                        self.RESULT_KEY_SCORE: score,
                    }
                )

        return parsed, timing, shape
