# 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

    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.

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

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

        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

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

        Note:
            - Optimized for console logging and operator dashboards
            - Includes YOLO-specific parameters (iou, imgsz, half)
        """
        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})")

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

    @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 instance with type validation.

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

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

        Note:
            - Setting a new instance 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

    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 `YOLOInstance` objects.

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

        Args:
            image (np.ndarray): Input image in BGR format with shape `(H, W, 3)`.
            track (bool): Whether to use tracking mode (persistent IDs).
            imgsz (int): Inference image size. Defaults to `1280`.
            verbose (bool): Enable verbose logging. Defaults to `False`.
            max_det (int): Maximum detections per frame. Defaults to `640`.
            half (bool): Use FP16 inference. Defaults to `False`.
            iou (float): IoU threshold for NMS. Defaults to `0.7`.

        Returns:
            List[Instance]: Detected objects with normalized coordinates.
                Empty list if no valid detections are found.

        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
        """
        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,
                                       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,
                                         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 `YOLOInstance` objects.

        Args:
            result: YOLO result object with boxes.

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

        Note:
            - Converts normalized box coordinates to `Polygon` (4 corners)
            - Creates pseudo-OBB from box corners (not true rotated 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],
                                          category=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 `YOLOInstance` objects.

        Args:
            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],
                                          category=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 `YOLOInstance` objects.

        Args:
            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 pseudo-OBB from polygon bounding box (not true rotated box)
            - 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],
                                          category=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]