# acmenra_cv/inference/backends/yolo_backend.py

import logging
import numpy as np
from enum import EnumType
from ultralytics import YOLO
from datetime import datetime
from typing import List, Dict
from abc import abstractmethod

from acmenra_cv.instance.enums import TaskType
from acmenra_cv.inference import Result, Timing
from acmenra_cv.inference.backend import Backend
from acmenra_cv.instance import Point, Box, Obb, Polygon, Instance, DeviceType


logger = logging.getLogger(__name__)


class YOLOBackend(Backend):
    """
    YOLO-specific inference backend using Ultralytics YOLO models.

    Implements the Backend contract for YOLO models, handling detection, segmentation,
    and OBB tasks. Supports built-in tracking via YOLO's native tracker (BoT-SORT, etc.).

    Key features:
    - Automatic task type detection from YOLO model (detect/segment/obb)
    - Built-in tracking support via YOLO's native tracker
    - Optimized for YOLO's output format (boxes, masks, obb)
    - Full timing metrics extraction from YOLO's speed dictionary
    - High-quality mask generation for segmentation tasks

    Attributes:
        _model (YOLO): Ultralytics YOLO model instance.
        _iou (float): IoU threshold for NMS [0.0, 1.0].
        _imgsz (int): Inference image size in pixels.
        _half (bool): FP16 inference flag.
        _refined (bool): High-quality mask generation flag.

    Note:
        - Task type is automatically determined from model.task
        - Built-in tracking is only available for detect/segment tasks
        - OBB task does not support built-in tracking (frame-based IDs assigned)
        - refined=True improves mask quality but reduces inference speed by ~20-30%
    """

    def __init__(self,
                 model: YOLO,
                 device: DeviceType,
                 category: EnumType,
                 task_type: TaskType,
                 threshold: float = 0.35,
                 iou: float = 0.7,
                 imgsz: int = 640,
                 half: bool = False,
                 refined: bool = False) -> None: # tested (good)
        """
        Initializes the YOLO backend with model and configuration parameters.

        Args:
            model (YOLO): Ultralytics YOLO instance.
            device (DeviceType): Hardware backend for inference.
            category (EnumType): Semantic class enumeration for label mapping.
            task_type (TaskType): Type of task the model performs (detect/segment/obb).
            threshold (float, optional): Confidence threshold [0.0, 1.0]. Defaults to 0.35.
            iou (float, optional): IoU threshold for NMS [0.0, 1.0]. Defaults to 0.7.
            imgsz (int, optional): Inference image size in pixels. Defaults to 640.
            half (bool, optional): FP16 inference flag. Defaults to False.
            refined (bool, optional): High-quality mask generation for segmentation.
                When True, generates masks at full model resolution for smoother boundaries.
                When False, uses lower resolution masks with upsampling (faster but less detailed).
                Defaults to False. Only affects SEGMENT task type.

        Raises:
            TypeError: If any parameter has incorrect type.
            ValueError: If thresholds are out of range.

        Note:
            - task_type must match the model's actual capabilities
            - imgsz should be a multiple of 32 for optimal YOLO performance
            - half=True only works on CUDA devices with compute capability >= 7.0
        """
        super().__init__(device=device, category=category, task_type=task_type, threshold=threshold)
        self.model = model
        self.iou = iou
        self.imgsz = imgsz
        self.half = half
        self.refined = refined

    def __str__(self) -> str: # tested (good)
        """
        Returns a human-readable string representation of the YOLO backend.

        Returns:
            str: Concise summary with device, task type, and key parameters.
                 Format: "YOLOBackend(device=mps, task_type=detect, threshold=0.35, iou=0.70, refined=False)"

        Note:
            - Optimized for console logging and operator dashboards
            - Includes YOLO-specific parameters (iou, imgsz, half, refined)
        """
        return (f"YOLOBackend(device={self.device.value}, "
                f"task_type={self.task_type.value}, "
                f"threshold={self.threshold:.2f}, "
                f"iou={self.iou:.2f}, "
                f"imgsz={self.imgsz}, "
                f"half={self.half}, "
                f"refined={self.refined})")

    def __repr__(self) -> str: # tested (good)
        """
        Returns a detailed string representation suitable for debugging and reproduction.

        Returns:
            str: Exact representation including all canonical parameters.
                 Uses repr-safe formatting (!r) for unambiguous parsing.

        Note:
            - Includes all YOLO-specific parameters with full precision
            - Output can theoretically be used to recreate the object in interactive sessions
            - Essential for debugging pipeline configuration
        """
        return (f"YOLOBackend(model={self.model!r}, "
                f"category={self.category!r}, "
                f"device={self.device!r}, "
                f"task_type={self.task_type!r}, "
                f"threshold={self.threshold!r}, "
                f"iou={self.iou!r}, "
                f"imgsz={self.imgsz!r}, "
                f"half={self.half!r}, "
                f"refined={self.refined!r})")

    @property
    def model(self) -> YOLO: # tested (good)
        """
        YOLO: Gets the object detection instance used for inference.

        Returns:
            YOLO: Ultralytics YOLO instance instance.

        Note:
            - Changing the instance automatically invalidates previous tracking state
            - Via the `_is_update` flag to ensure consistent behavior
            - Model must support the configured task type (`segment`/`detect`/`obb`)
        """
        return self._model

    @model.setter
    def model(self, model: YOLO) -> None: # tested (good)
        """
        Sets the object detection model with type validation.

        Args:
            model (YOLO): New YOLO model instance for detection.

        Raises:
            TypeError: If model is not a YOLO instance.

        Note:
            - Setting a new model automatically triggers state refresh
            - Via `_is_update = True` to ensure subsequent processing uses updated configuration
            - Existing tracks are NOT automatically migrated; consider resetting `_objects` if needed
        """
        if not isinstance(model, YOLO):
            raise TypeError(f"'instance' must be '<YOLO>', but got {type(model).__name__}")
        self._model = model
        self._is_update = True

    @property
    def iou(self) -> float: # tested (good)
        """
        float: Gets the IoU (Intersection over Union) threshold for NMS.

        Returns:
            float: IoU threshold in range [0.0, 1.0].

        Note:
            - Used for Non-Maximum Suppression (NMS) during detection post-processing
            - Higher values = more overlapping boxes retained (less aggressive suppression)
            - Lower values = more aggressive suppression of duplicate detections
            - Typical values: 0.5-0.7 for most applications
        """
        return self._iou

    @iou.setter
    def iou(self, iou: float) -> None: # tested (good)
        """
        Sets the IoU threshold for NMS with validation.

        Args:
            iou (float): IoU threshold to set [0.0, 1.0].

        Raises:
            TypeError: If iou is not a float.
            ValueError: If iou is outside range [0.0, 1.0].

        Note:
            - IoU threshold controls duplicate detection suppression aggressiveness
            - Higher values allow more overlapping detections (useful for crowded scenes)
            - Lower values are more aggressive at removing duplicate boxes
        """
        if not isinstance(iou, float):
            raise TypeError(f"'iou' must be '<float>', but got {type(iou).__name__}")
        if not 0.0 <= iou <= 1.0:
            raise ValueError(f"'iou' must be in range [0.0, 1.0], got {iou}")
        self._iou = iou

    @property
    def imgsz(self) -> int: # tested (good)
        """
        int: Gets the inference image size.

        Returns:
            int: Image size in pixels (must be > 0).

        Note:
            - Should be a multiple of 32 for optimal YOLO performance
            - Typical values: 640, 1280, 1920
            - Larger values = better accuracy but slower inference
        """
        return self._imgsz

    @imgsz.setter
    def imgsz(self, imgsz: int) -> None: # tested (good)
        """
        Sets the inference image size with validation.

        Args:
            imgsz (int): Image size in pixels.

        Raises:
            TypeError: If imgsz is not an integer.
            ValueError: If imgsz is <= 0.

        Note:
            - Must be strictly > 0
            - Should be a multiple of 32 for YOLO compatibility
            - Changes affect subsequent predict() calls
        """
        if not isinstance(imgsz, int) or isinstance(imgsz, bool):
            raise TypeError(f"'imgsz' must be '<int>', but got {type(imgsz).__name__}")
        if imgsz <= 0:
            raise ValueError(f"'imgsz' must be > 0, got {imgsz}")
        self._imgsz = imgsz

    @property
    def half(self) -> bool: # tested (good)
        """
        bool: Gets the FP16 inference flag.

        Returns:
            bool: True if FP16 inference is enabled, False otherwise.

        Note:
            - FP16 inference is faster but may reduce accuracy slightly
            - Only supported on CUDA devices with compute capability >= 7.0
            - Not supported on MPS (Apple Silicon) or CPU
        """
        return self._half

    @half.setter
    def half(self, half: bool) -> None: # tested (good)
        """
        Sets the FP16 inference flag with validation.

        Args:
            half (bool): Enable FP16 inference.

        Raises:
            TypeError: If half is not a boolean.

        Note:
            - Boolean values are explicitly validated (bool is subclass of int)
            - Changes affect subsequent predict() calls
            - May cause errors on unsupported devices
        """
        if not isinstance(half, bool):
            raise TypeError(f"'half' must be '<bool>', but got {type(half).__name__}")
        self._half = half

    @property
    def refined(self) -> bool: # tested (good)
        """
        bool: Gets the high-quality mask generation flag.

        Returns:
            bool: True if refined masks are enabled, False otherwise.

        Note:
            - When True: generates masks at full model resolution (e.g., 640x640)
            - When False: generates masks at lower resolution (e.g., 160x160) with upsampling
            - Only affects SEGMENT task type
            - Improves mask boundary smoothness and detail
            - Reduces inference speed by ~20-30%
        """
        return self._refined

    @refined.setter
    def refined(self, refined: bool) -> None: # tested (good)
        """
        Sets the high-quality mask generation flag with validation.

        Args:
            refined (bool): Enable refined mask generation.

        Raises:
            TypeError: If refined is not a boolean.

        Note:
            - Boolean values are explicitly validated (bool is subclass of int)
            - Changes affect subsequent predict()/track() calls
            - Only relevant for SEGMENT task type
        """
        if not isinstance(refined, bool):
            raise TypeError(f"'refined' must be '<bool>', but got {type(refined).__name__}")
        self._refined = refined

    def predict(self,
                frame: np.ndarray,
                max_det: int = 640,
                track: bool = False,
                verbose: bool = False) -> Result: # TODO needs tests
        """
        Performs YOLO inference and converts results to normalized `Instance` objects.

        This function handles:
        - Model inference on the provided image
        - Processing segmentation masks, bounding boxes, or OBB results
        - Conversion of normalized coordinates to `Instance` objects
        - Semantic class mapping via category enum

        Args:
            frame (np.ndarray): Input image in BGR format with shape `(H, W, 3)`.
            track (bool): Whether to use tracking mode (persistent IDs).
            verbose (bool): Enable verbose logging. Defaults to `False`.
            max_det (int): Maximum detections per frame. Defaults to `640`.

        Returns:
            Result: Universal result container with instances and timing.

        Raises:
            TypeError: If any parameter has incorrect type.
            ValueError: If numeric parameters have invalid values.

        Note:
            - Supports three detection modes: segmentation, bbox, OBB
            - Polygon coordinates are normalized to `[0.0, 1.0]`
            - Requires Ultralytics YOLO instance with appropriate task type
            - Use `track=True` for persistent object IDs across frames
            - Refined masks are automatically applied based on `self.refined` setting
        """
        if frame is None or not isinstance(frame, np.ndarray):
            raise TypeError(f"'image' must be a valid numpy array, got {type(frame)}!")

        if not isinstance(track, bool):
            raise TypeError(f"'track' must be a valid numpy array, got {type(frame)}!")

        if not isinstance(verbose, bool):
            raise TypeError(f"'verbose' must be a valid numpy array, got {type(frame)}!")

        if not isinstance(max_det, int):
            raise TypeError(f"'max_det' must be '<int>', but got {type(max_det).__name__}")
        if max_det < 0:
            raise ValueError(f"'max_det' must be >= 0, got {max_det}")


        if track:
            results = self.model.track(source=frame,
                                       device=self.device.value,
                                       conf=self.threshold,
                                       imgsz=self.imgsz,
                                       verbose=verbose,
                                       max_det=max_det,
                                       half=self.half,
                                       iou=self.iou,
                                       retina_masks=self.refined,
                                       persist=True)
        else:
            results = self.model.predict(source=frame,
                                         device=self.device.value,
                                         conf=self.threshold,
                                         imgsz=self.imgsz,
                                         verbose=verbose,
                                         max_det=max_det,
                                         half=self.half,
                                         retina_masks=self.refined,
                                         iou=self.iou)

        instances = []
        for result in results:
            if result.obb is not None:
                instances += self._process_obb_results(frame=frame, result=result)
            if result.masks is not None and result.obb is None:
                instances += self._process_seg_results(frame=frame, result=result)
            elif result.boxes is not None and result.masks is None and result.obb is None:
                instances += self._process_bbox_results(frame=frame, result=result)

        return Result(device=self.device,
                      task_type=self.task_type,
                      instances=instances,
                      timing=Timing(preprocess=results[0].speed['preprocess'],
                                    prediction=results[0].speed['inference'],
                                    postprocess=results[0].speed['postprocess']),
                      width=results[0].orig_shape[0],
                      height=results[0].orig_shape[1],
                      depth=1,
                      category=self.category
                      )

    def _process_bbox_results(self, frame: np.ndarray, result) -> List[Instance]: # TODO needs tests
        """
        Processes YOLO bounding box results into `Instance` objects.

        Args:
            frame (np.ndarray): Input frame (used for shape reference if needed).
            result: YOLO result object with boxes.

        Returns:
            List[Instance]: List of processed instances.

        Note:
            - Converts normalized box coordinates to `Box`.
            - Used for detection-only models (`task='detect'`).
        """
        instances = []
        if result.boxes is not None:
            for box in result.boxes:
                instances.append(Instance(id=box.id.int().tolist()[0] if box.id is not None else None,
                                          class_id=box.cls.int().tolist()[0],
                                          label=self.category(box.cls.int().tolist()[0]),
                                          conf=box.conf.cpu().numpy().flatten().tolist()[0],
                                          box=Box.from_xyxyn(xyxyn=box.xyxyn.cpu().numpy().flatten().tolist())
                                          )
                                 )
        return instances

    def _process_obb_results(self, frame: np.ndarray, result) -> List[Instance]: # TODO needs tests
        """
        Processes YOLO OBB results into `Instance` objects.

        Args:
            frame (np.ndarray): Input frame (used for shape reference).
            result: YOLO result object with OBB detections.

        Returns:
            List[Instance]: List of processed instances with OBB data.

        Note:
            - Uses true rotated bounding boxes from YOLO OBB models
            - Coordinates are normalized `[0.0, 1.0]`
            - Used for OBB detection models (`task='obb'`)
            - YOLOv8/v11 OBB models only
        """
        instances = []
        h, w = frame.shape[:2]
        if result.obb is not None:
            for id, obb in enumerate(result.obb):
                obb_instance = Obb.from_xywhrn(self.normalize_xywhr(xywhr=obb.xywhr.cpu().numpy().flatten().tolist(),
                                                                    img_w=w, img_h=h))
                logger.warning(
                    f"Using OBB task. Built-in tracking is not supported. "
                    f"A temporary numeric frame-based ID ({id}) has been assigned."
                )
                instances.append(Instance(id=id,
                                          class_id=obb.cls.int().tolist()[0],
                                          label=self.category(obb.cls.int().tolist()[0]),
                                          conf=float(obb.conf.cpu().item()),
                                          obb=obb_instance,
                                          box=Box.safe(left=obb_instance.left, right=obb_instance.right,
                                                       top=obb_instance.top, bottom=obb_instance.bottom,
                                                       front=0.0, back=0.0)
                                          )
                                 )
        return instances

    def _process_seg_results(self, frame: np.ndarray, result) -> List[Instance]: # TODO needs tests
        """
        Processes YOLO segmentation results into `Instance` objects.

        Args:
            frame (np.ndarray): Input frame (used for shape reference).
            result: YOLO result object with segmentation masks.

        Returns:
            List[Instance]: List of processed instances with polygon masks.

        Note:
            - Uses segmentation masks from YOLO segmentation models
            - Polygon coordinates are normalized `[0.0, 1.0]`
            - Creates `Box` from polygon bounds for unified interface
            - Used for segmentation models (`task='segment'`)
        """
        instances = []
        if result.masks is not None and result.boxes is not None:
            for box, mask in zip(result.boxes, result.masks.xyn):
                polygon = Polygon.from_xyn(mask)
                instances.append(Instance(id=box.id.int().tolist()[0] if box.id is not None else None,
                                          class_id=box.cls.int().tolist()[0],
                                          label=self.category(box.cls.int().tolist()[0]),
                                          conf=box.conf.cpu().numpy().flatten().tolist()[0],
                                          polygon=polygon,
                                          box=Box(left=polygon.left, right=polygon.right,
                                                  top=polygon.top, bottom=polygon.bottom,
                                                  front=0.0, back=0.0)
                                          )
                                 )
        return instances

    @staticmethod
    def normalize_xywhr(xywhr: list[float], img_w: float, img_h: float) -> list[float]:
        """
        Явно нормализует [cx, cy, w, h, angle_rad] из пикселей в [0.0, 1.0].
        Угол остаётся в радианах (Obb.from_xywhr сам конвертирует в градусы).
        """
        cx, cy, w, h, angle = xywhr
        return [cx / img_w, cy / img_h, w / img_w, h / img_h, angle]