from pathlib import Path
from typing import Any

import numpy as np


SEG_MODEL_CLASSES = {
    "small": "RFDETRSegSmall",
    "medium": "RFDETRSegMedium",
}


def load_rfdetr_seg_model(model_size: str = "small", checkpoint_path: str | None = None) -> Any:
    if model_size not in SEG_MODEL_CLASSES:
        raise ValueError(f"Unsupported model_size={model_size}. Use one of {sorted(SEG_MODEL_CLASSES)}.")
    try:
        import rfdetr
    except ImportError as exc:
        raise ImportError("rfdetr is not installed. Install with `pip install rfdetr`.") from exc

    model_cls = getattr(rfdetr, SEG_MODEL_CLASSES[model_size])
    kwargs = {}
    if checkpoint_path:
        kwargs["pretrain_weights"] = checkpoint_path
    return model_cls(**kwargs)


def detections_to_records(detections: Any) -> list[dict]:
    xyxy = getattr(detections, "xyxy", None)
    confidence = getattr(detections, "confidence", None)
    class_id = getattr(detections, "class_id", None)
    masks = getattr(detections, "mask", None)
    if xyxy is None:
        return []

    records = []
    for idx, box in enumerate(xyxy):
        records.append(
            {
                "bbox": [float(v) for v in box],
                "confidence": float(confidence[idx]) if confidence is not None else 0.0,
                "class_id": int(class_id[idx]) if class_id is not None else 0,
                "mask": np.asarray(masks[idx]).astype(bool) if masks is not None else None,
            }
        )
    return records


def predict_records(
    image_path: str | Path,
    model_size: str = "small",
    checkpoint_path: str | None = None,
    threshold: float = 0.5,
) -> list[dict]:
    model = load_rfdetr_seg_model(model_size=model_size, checkpoint_path=checkpoint_path)
    detections = model.predict(str(image_path), threshold=threshold)
    return detections_to_records(detections)