from typing import Dict, List
import os
import numpy as np

from app.config import settings
from app.utils import read_classes, crop_with_padding, draw_results, encode_image_to_base64
from app.inference.yolo_onnx import YOLOv5ONNXDetector
from app.inference.convnext_onnx import ConvNeXtONNXClassifier


class BirdCascadePipeline:
    def __init__(self):
        self.det_classes = read_classes(settings.det_classes_path, keep_empty=False)

        # 分类类别文件必须保留空项，因为 ;; 表示该类别不显示
        self.cls_classes = read_classes(settings.cls_classes_path, keep_empty=True)

        if not os.path.exists(settings.det_model_path):
            raise FileNotFoundError(f"未找到 YOLOv5 模型: {settings.det_model_path}")

        if not os.path.exists(settings.cls_model_path):
            raise FileNotFoundError(f"未找到 ConvNeXt 模型: {settings.cls_model_path}")

        self.detector = YOLOv5ONNXDetector(
            model_path=settings.det_model_path,
            class_names=self.det_classes,
            img_size=settings.det_img_size,
            device=settings.device,
        )

        self.classifier = ConvNeXtONNXClassifier(
            model_path=settings.cls_model_path,
            class_names=self.cls_classes,
            img_size=settings.cls_img_size,
            device=settings.device,
        )

        self.target_det_class_ids = settings.target_det_class_ids

    def _filter_target_detections(self, detections: List[Dict]) -> List[Dict]:
        """
        只保留 YOLO 检测类别为 0 的目标，也就是“鸟”。
        """
        if self.target_det_class_ids is None or len(self.target_det_class_ids) == 0:
            return detections

        filtered = []
        for det in detections:
            det_class_id = int(det.get("det_class_id", -1))
            if det_class_id in self.target_det_class_ids:
                filtered.append(det)

        return filtered

    def predict(
        self,
        img_bgr: np.ndarray,
        return_image: bool = False,
        det_conf: float | None = None,
        iou_thres: float | None = None,
    ) -> Dict:
        h, w = img_bgr.shape[:2]

        det_conf = settings.det_conf_thres if det_conf is None else det_conf
        iou_thres = settings.det_iou_thres if iou_thres is None else iou_thres

        detections = self.detector.infer(
            img_bgr,
            conf_thres=det_conf,
            iou_thres=iou_thres,
            max_det=settings.max_detections,
        )

        detections = self._filter_target_detections(detections)

        results: List[Dict] = []

        for det in detections:
            crop = crop_with_padding(
                img_bgr,
                det["bbox_xyxy"],
                padding_ratio=settings.crop_padding_ratio,
            )

            cls_result = self.classifier.infer(crop)

            cls_id = int(cls_result["cls_class_id"])
            cls_name = str(cls_result["cls_class_name"] or "")

            merged = {
                "bbox_xyxy": [round(float(v), 2) for v in det["bbox_xyxy"]],
                "det_confidence": round(float(det["det_confidence"]), 6),
                "det_class_id": int(det["det_class_id"]),
                "det_class_name": str(det["det_class_name"]),
                "cls_class_id": cls_id,
                "cls_class_name": cls_name,
                "cls_confidence": round(float(cls_result["cls_confidence"]), 6),
            }
            results.append(merged)

        annotated_image_base64 = None
        if return_image:
            annotated = draw_results(img_bgr, results)
            annotated_image_base64 = encode_image_to_base64(annotated)

        return {
            "success": True,
            "image_width": w,
            "image_height": h,
            "num_objects": len(results),
            "target_det_class_ids": self.target_det_class_ids,
            "results": results,
            "annotated_image_base64": annotated_image_base64,
            "message": "ok",
        }