import os
from pathlib import Path

import cv2
import numpy as np
import requests

# ultralytics builds a module-level `SettingsManager()` at import time that calls
# `dist.barrier(device_ids=[RANK])`, passing the GLOBAL rank as a CUDA device ordinal
# (ultralytics/utils/__init__.py -> `torch_distributed_zero_first(RANK)`; note it should have
# used LOCAL_RANK). On multi-node runs the global rank can exceed the per-node GPU count
# (e.g. rank 15 on an 8-GPU node), which crashes the import with "CUDA error: invalid device
# ordinal". Present the node-local rank to ultralytics just for this import so its barrier
# targets a valid device, then restore the true RANK so distributed training is unaffected.
#
# The import itself is wrapped in try/except as well: ultralytics is an optional dependency
# that is only needed when a detector is actually instantiated, so a missing/broken install
# must not make this module unimportable.
def _import_yolo():
    """Import ultralytics' YOLO behind the node-local RANK guard; return None on failure."""
    real_rank = os.environ.get("RANK")
    if real_rank is not None:
        os.environ["RANK"] = os.environ.get("LOCAL_RANK", "0")
    try:
        from ultralytics import YOLO
        return YOLO
    except Exception:
        return None
    finally:
        if real_rank is not None:
            os.environ["RANK"] = real_rank


YOLO = _import_yolo()


def _get_yolo():
    """Return the ultralytics YOLO class, re-importing it lazily if the eager import failed."""
    if YOLO is not None:
        return YOLO
    yolo_cls = _import_yolo()
    if yolo_cls is None:
        raise ImportError(
            "ultralytics is required by the YOLO detectors in this module. "
            "Install it with `pip install ultralytics`."
        )
    return yolo_cls


MODEL_DIR = "./models"
instance_url = "https://zhoumo-bj.oss-cn-beijing.aliyuncs.com/Code/model_data/Stable_Diffusion/annotator/yolov8x-seg.pt"
yolo26_url = "https://github.com/ultralytics/assets/releases/download/v8.4.0/yolo26x-seg.pt"

def urldownload(url, filename):
    """
    Download a file to the specified directory.
    :param url: URL of the file to download.
    :param filename: Destination path, e.g. ./test.xls.
    """
    down_res = requests.get(url)
    with open(filename,'wb') as file:
        file.write(down_res.content)

class ObjectInstanceDetector:
    def __init__(self, model_dir=MODEL_DIR, combine_overlap_boxes=False, device=None):
        model_path = os.path.join(model_dir, "yolov8x-seg.pt")
        if not os.path.exists(model_path):
            urldownload(instance_url, model_path)

        self.model  = _get_yolo()(model_path).to(device)
        self.cob    = combine_overlap_boxes
        self.overlap_threhold = 0.05
    
    def cal_miou(self, mask1, mask2):
        tp = mask1 * mask2
        if np.shape(mask2)[0] != np.shape(mask2)[1]:
            tile_mask1 = np.tile(mask1, [np.shape(mask2)[0], 1, 1])
        tp_fp_fn = np.clip(tile_mask1 + mask2, 0, 1)
        miou = np.sum(np.sum(tp, 1), 1) / np.sum(np.sum(tp_fp_fn, 1), 1) 
        return miou

    def combine_overlap_boxes(self, boxes, masks):       
        outputs         = []
        outputs_masks   = []
        while np.shape(boxes)[0]:
            # Get the mask dimensions of the smallest box
            h, w = np.shape(masks[0])

            # Store the smallest box coordinates and mask
            outputs.append(boxes[0:1, :])
            outputs_masks.append(masks[0:1])
            if len(boxes) == 1:
                break
            
            # Compute mIoU between the smallest box and remaining boxes
            outputs_mask    = np.zeros([1, h, w])
            outputs_mask[0, int(outputs[-1][0][1]):int(outputs[-1][0][3]), int(outputs[-1][0][0]):int(outputs[-1][0][2])] = 1
            miou            = self.cal_miou(outputs_mask, masks[1:])

            # Boxes with mIoU above threshold are considered significantly overlapping
            overlap_boxes   = boxes[1:][miou > self.overlap_threhold]
            overlap_masks   = masks[1:][miou > self.overlap_threhold]
            if len(overlap_boxes) > 0:
                overlap_boxes   = np.concatenate([outputs[-1], overlap_boxes], 0)
                overlap_masks   = np.concatenate([outputs_masks[-1], overlap_masks], 0)

                outputs[-1]         = np.array([[np.min(overlap_boxes[:, 0]), np.min(overlap_boxes[:, 1]), np.max(overlap_boxes[:, 2]), np.max(overlap_boxes[:, 3])]])
                outputs_masks[-1]   = np.max(overlap_masks, 0, keepdims=True)
                
                cv2.imwrite("test.jpg", overlap_masks[0] * 255)

            # Keep non-overlapping masks and boxes
            masks           = masks[1:][miou <= self.overlap_threhold]
            boxes           = boxes[1:][miou <= self.overlap_threhold]

            if len(boxes) == 0:
                break

            # Re-compute mIoU against remaining boxes
            outputs_mask    = np.zeros([1, h, w])
            outputs_mask[0, int(outputs[-1][0][1]):int(outputs[-1][0][3]), int(outputs[-1][0][0]):int(outputs[-1][0][2])] = 1
            miou            = self.cal_miou(outputs_mask, masks)

            # If any remaining boxes overlap significantly, merge them
            overlap_boxes   = boxes[miou > self.overlap_threhold]
            if len(overlap_boxes) > 0:
                boxes       = np.concatenate([outputs[-1], boxes], 0)
                masks       = np.concatenate([outputs_masks[-1], masks], 0)
                del outputs[-1]
                del outputs_masks[-1]
        # Stack results
        if len(outputs) != 0:
            outputs = np.concatenate(outputs, 0)
            outputs_masks = np.concatenate(outputs_masks, 0)
        return outputs, outputs_masks

    def __call__(self, img):
        h, w, c     = np.shape(img)
        results     = self.model.predict(source=img, imgsz=1024, conf=0.40)

        try:
            categories  = np.array([result.boxes.cls.cpu().numpy() for result in results])[0]
            boxes       = np.array([result.boxes.xyxy.cpu().numpy() for result in results])[0]
            masks       = np.array([result.masks.data.cpu().numpy() for result in results])[0]
        except Exception:
            return [], [], []
        
        if len(boxes) != 0:
            if self.cob:
                masks       = np.array([cv2.resize(mask, (w, h), cv2.INTER_NEAREST) for mask in masks])

                areas       = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
                args        = np.argsort(areas)
                boxes       = boxes[args]
                masks       = masks[args]
                
                outputs, outputs_masks = self.combine_overlap_boxes(boxes, masks)
                outputs, outputs_masks = self.combine_overlap_boxes(outputs, outputs_masks)

                areas       = (outputs[:, 3] - outputs[:, 1]) * (outputs[:, 2] - outputs[:, 0])
                args        = np.argsort(areas)[::-1]
                outputs       = outputs[args]
                outputs_masks = outputs_masks[args]

                outputs_masks = np.uint8(outputs_masks) * 255 
                outputs_masks = [cv2.resize(mask, (w, h), cv2.INTER_NEAREST) for box, mask in zip(outputs, outputs_masks)]
                return outputs, boxes, outputs_masks
            else:
                areas       = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
                args        = np.argsort(areas)[::-1]

                boxes       = boxes[args]
                masks       = masks[args]

                masks       = np.uint8(masks) * 255 
                masks       = [cv2.resize(mask, (w, h), cv2.INTER_NEAREST) for box, mask in zip(boxes, masks)]
                return boxes, boxes, masks
        else:
            return [], boxes, []


class ObjectDetector:
    """
    COCO-pretrained object detection and segmentation based on YOLO26-seg.
    Supports 80 COCO categories. Prediction results include class names.
    """

    def __init__(self, model_dir=MODEL_DIR, model_size="x", device=None):
        """
        Initialize YOLO26 COCO detector.

        Args:
            model_dir: Directory to store model weights.
            model_size: Model scale, one of "n", "s", "m", "l", "x".
            device: Inference device, e.g. "cuda:0" or "cpu". Auto-selects if None.
        """
        model_name = f"yolo26{model_size}-seg.pt"
        model_path = os.path.join(model_dir, model_name)
        if not os.path.exists(model_path):
            os.makedirs(model_dir, exist_ok=True)
            urldownload(yolo26_url.replace("yolo26x-seg.pt", model_name), model_path)

        self.model = _get_yolo()(model_path)
        if device is not None:
            self.model.to(device)

    def __call__(self, img, conf=0.25, imgsz=1024):
        """
        Run COCO 80-class object detection on the input image.

        Args:
            img: Input image (numpy array, BGR format) or image file path.
            conf: Confidence threshold, default 0.25.
            imgsz: Inference image size, default 1024.

        Returns:
            list[dict]: Each detection result contains:
                - "class_name" (str): COCO category name
                - "class_id" (int): Category ID
                - "confidence" (float): Confidence score
                - "bbox" (np.ndarray): Bounding box [x1, y1, x2, y2]
                - "mask" (np.ndarray | None): Segmentation mask (H, W), uint8 format
        """
        # Get original image dimensions
        if isinstance(img, str):
            img_bgr = cv2.imread(img)
        else:
            img_bgr = img
        h, w = img_bgr.shape[:2]

        results = self.model.predict(source=img_bgr, imgsz=imgsz, conf=conf)

        detections = []
        for result in results:
            if result.boxes is None or len(result.boxes) == 0:
                continue

            boxes_xyxy = result.boxes.xyxy.cpu().numpy()
            confs = result.boxes.conf.cpu().numpy()
            cls_ids = result.boxes.cls.cpu().numpy().astype(int)

            # Get class name mapping (COCO 80 classes)
            id_to_name = {int(k): v for k, v in result.names.items()} if result.names else {}

            # Extract segmentation masks if available
            has_masks = result.masks is not None and result.masks.data is not None
            if has_masks:
                mask_data = result.masks.data.cpu().numpy()

            for i in range(len(boxes_xyxy)):
                class_id = int(cls_ids[i])
                class_name = id_to_name.get(class_id, str(class_id))
                confidence = float(confs[i])
                bbox = boxes_xyxy[i]

                # Resize mask to original image dimensions
                if has_masks and i < len(mask_data):
                    mask = mask_data[i]
                    mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
                    mask = (mask > 0.5).astype(np.uint8) * 255
                else:
                    mask = None

                detections.append({
                    "class_name": class_name,
                    "class_id": class_id,
                    "confidence": confidence,
                    "bbox": bbox,
                    "mask": mask,
                })

        # Sort by confidence in descending order
        detections.sort(key=lambda x: x["confidence"], reverse=True)
        return detections

    def detect_and_draw(self, img, conf=0.25, imgsz=1024):
        """
        Detect objects and draw results on the image (for debugging/visualization).

        Args:
            img: Input image (numpy array, BGR format) or image file path.
            conf: Confidence threshold.
            imgsz: Inference image size.

        Returns:
            tuple: (annotated_image, detections)
        """
        if isinstance(img, str):
            img_bgr = cv2.imread(img)
        else:
            img_bgr = img.copy()

        detections = self(img_bgr, conf=conf, imgsz=imgsz)

        for det in detections:
            x1, y1, x2, y2 = det["bbox"].astype(int)
            label = f'{det["class_name"]} {det["confidence"]:.2f}'
            cv2.rectangle(img_bgr, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(img_bgr, label, (x1, y1 - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

            if det["mask"] is not None:
                colored_mask = np.zeros_like(img_bgr)
                colored_mask[:, :, 1] = det["mask"]
                img_bgr = cv2.addWeighted(img_bgr, 1.0, colored_mask, 0.4, 0)

        return img_bgr, detections