import os
import time
import cv2
import numpy as np
from ultralytics import YOLO
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont

from core.accel import resolve_yolo_weights


# ==================== 【配置区域】 ====================
_ROOT = Path(__file__).resolve().parents[1]

class Config:
    # 你的Plate检测模型路径（相对项目根目录）；若存在 best_openvino_model 会自动优先加载
    MODEL_PATH = str(_ROOT / "best.pt")

    CONF_THRESHOLD = 0.90
    IOU_THRESHOLD = 0.45
    # 显式 imgsz，便于 OpenVINO / CPU 控制耗时；可用环境变量 DISH_YOLO_IMGSZ 覆盖
    IMGSZ = int(os.environ.get("DISH_YOLO_IMGSZ", "640"))

    BOX_THICKNESS = 2
    FONT_SIZE = 22

    # 输出目录
    OUTPUT_DIR = r"detect"
    Cutting = r"cutting"

    GENERATE_VISUALIZATION = True  # 是否生成可视化图

    # 默认类别名（加载权重后会被 model.names 覆盖）
    LABEL_NAMES = {0: "Plate"}
    COLORS = {'Plate': (0, 255, 0)}
    # 检测到后直接丢弃，不进入裁剪/识别
    FILTER_CLASS_NAMES = ("手", "空盘")


# ========================================================

def _load_filter_names(config):
    names = getattr(config, "FILTER_CLASS_NAMES", None)
    if names is None:
        try:
            from dao.settings import DISH_DETECT_FILTER_NAMES
            names = DISH_DETECT_FILTER_NAMES
        except Exception:
            names = ("手", "空盘")
    return {str(n).strip() for n in (names or []) if str(n).strip()}


def _load_edge_filter_params(config):
    """贴边/过小框过滤参数；可用 Config 覆盖，否则读 settings。"""
    def _get(name, default):
        if hasattr(config, name):
            return getattr(config, name)
        try:
            from dao import settings as S
            return getattr(S, name, default)
        except Exception:
            return default

    return {
        "edge_margin": int(_get("DISH_DETECT_EDGE_MARGIN_PX", 8)),
        "edge_min_area_ratio": float(_get("DISH_DETECT_EDGE_MIN_AREA_RATIO", 0.02)),
        "edge_min_side": int(_get("DISH_DETECT_EDGE_MIN_SIDE_PX", 48)),
        "min_area_ratio": float(_get("DISH_DETECT_MIN_AREA_RATIO", 0.008)),
    }


class YOLOv11PlateExtractor:
    def __init__(self, config):
        self.config = config
        weights = resolve_yolo_weights(config.MODEL_PATH)
        self.model = YOLO(weights)
        self.weights_path = weights
        # 以权重内类别为准
        try:
            model_names = getattr(self.model, "names", None) or {}
            if isinstance(model_names, dict) and model_names:
                self.config.LABEL_NAMES = {int(k): str(v) for k, v in model_names.items()}
        except Exception:
            pass
        self.filter_class_names = _load_filter_names(config)
        self.edge_filter = _load_edge_filter_params(config)
        print(f"✅ Plate 模型加载成功: {weights}")
        print(f"[OK] detect classes: {self.config.LABEL_NAMES}")
        if self.filter_class_names:
            print(f"[OK] detect filter drop: {sorted(self.filter_class_names)}")
        ef = self.edge_filter
        print(
            f"[OK] edge/small box filter: margin={ef['edge_margin']}px "
            f"edge_area<{ef['edge_min_area_ratio']:.3f} "
            f"edge_side<{ef['edge_min_side']} "
            f"min_area<{ef['min_area_ratio']:.3f}"
        )

    def predict(self, source):
        """source 可为图片路径或 BGR ndarray。"""
        return self.model.predict(
            source=source,
            conf=self.config.CONF_THRESHOLD,
            iou=self.config.IOU_THRESHOLD,
            imgsz=getattr(self.config, "IMGSZ", 640),
            verbose=False,
        )[0]

    def _class_name(self, results, cls_id):
        cid = int(cls_id)
        names = getattr(results, "names", None) or self.config.LABEL_NAMES or {}
        if isinstance(names, dict):
            return str(names.get(cid, names.get(str(cid), cid)))
        return str(cid)

    def _should_drop_class(self, class_name):
        name = (class_name or "").strip()
        if not name:
            return False
        return name in self.filter_class_names

    def _should_drop_incomplete_box(self, box, img_w, img_h):
        """贴边且过小 / 面积过小 → 视为入画不完整，丢弃。"""
        try:
            x1, y1, x2, y2 = [float(v) for v in box[:4]]
        except Exception:
            return True
        bw = max(0.0, x2 - x1)
        bh = max(0.0, y2 - y1)
        if bw <= 1 or bh <= 1 or img_w <= 0 or img_h <= 0:
            return True
        area_ratio = (bw * bh) / float(img_w * img_h)
        ef = self.edge_filter or {}
        min_area = float(ef.get("min_area_ratio") or 0)
        if min_area > 0 and area_ratio < min_area:
            return True

        margin = max(0, int(ef.get("edge_margin") or 0))
        touches = (
            x1 <= margin
            or y1 <= margin
            or x2 >= (img_w - 1 - margin)
            or y2 >= (img_h - 1 - margin)
        )
        if not touches:
            return False
        edge_area = float(ef.get("edge_min_area_ratio") or 0)
        edge_side = int(ef.get("edge_min_side") or 0)
        # 贴边且「面积小」或「短边过短」→ 多半是只露一角
        if edge_area > 0 and area_ratio < edge_area:
            return True
        if edge_side > 0 and min(bw, bh) < edge_side:
            return True
        return False

    def iter_kept_detections(self, results, image_shape=None):
        """产出保留的检测框：(box, conf, class_name)。类别/贴边不完整已过滤。"""
        if results is None or results.boxes is None:
            return
        boxes = results.boxes.xyxy.cpu().numpy()
        confs = results.boxes.conf.cpu().numpy()
        if getattr(results.boxes, "cls", None) is not None:
            clss = results.boxes.cls.cpu().numpy()
        else:
            clss = [0] * len(boxes)
        img_h = img_w = None
        if image_shape is not None and len(image_shape) >= 2:
            img_h, img_w = int(image_shape[0]), int(image_shape[1])
        dropped_cls = 0
        dropped_edge = 0
        for box, conf, cls_id in zip(boxes, confs, clss):
            cname = self._class_name(results, cls_id)
            if self._should_drop_class(cname):
                dropped_cls += 1
                continue
            if img_w and img_h and self._should_drop_incomplete_box(box, img_w, img_h):
                dropped_edge += 1
                continue
            yield box, float(conf), cname
        if dropped_cls:
            print(f"[detect] filtered {dropped_cls} box(es) by class {sorted(self.filter_class_names)}")
        if dropped_edge:
            print(f"[detect] filtered {dropped_edge} incomplete/edge box(es)")

    def crop_plate(self, image, results):
        cropped_regions = []
        h, w = image.shape[:2]
        for box, conf, cname in self.iter_kept_detections(results, image_shape=image.shape):
            x1, y1, x2, y2 = map(int, box)
            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(w, x2)
            y2 = min(h, y2)
            if x2 > x1 and y2 > y1:
                crop = image[y1:y2, x1:x2]
                if crop.size > 0:
                    cropped_regions.append({
                        "image": crop.copy(),
                        "conf": float(conf),
                        "box": [x1, y1, x2, y2],
                        "class_name": cname,
                    })
        return cropped_regions

    # 标注默认：大库红框 + 白字；今日菜单由 style 覆盖为绿框绿字
    LABEL_BOX_COLOR = (229, 57, 53)
    LABEL_CHIP_BG = (229, 57, 53, 220)
    LABEL_CHIP_TEXT = (255, 255, 255, 255)

    def _load_font(self, size=None):
        size = size or self.config.FONT_SIZE
        for font_name in ("simhei.ttf", "simsun.ttc", "msyh.ttc", "msyh.ttf"):
            try:
                return ImageFont.truetype(font_name, size)
            except Exception:
                continue
        return ImageFont.load_default()

    def _draw_box_and_chip(self, base_img, box, text, font, style=None):
        """描边框 + 标签芯片；style 可覆盖颜色（今日菜单绿 / 大库红 / 未知橙）。"""
        style = style or {}
        outline = style.get("box") or (229, 57, 53, 180)
        chip_bg = style.get("chip_bg") or self.LABEL_CHIP_BG
        chip_fg = style.get("chip_fg") or self.LABEL_CHIP_TEXT
        thickness = int(style.get("box_width") or self.config.BOX_THICKNESS or 2)

        x1, y1, x2, y2 = map(int, box)
        rgba = base_img.convert("RGBA")
        overlay = Image.new("RGBA", rgba.size, (0, 0, 0, 0))
        odraw = ImageDraw.Draw(overlay)

        for t in range(max(1, thickness)):
            odraw.rectangle(
                [(x1 - t, y1 - t), (x2 + t, y2 + t)],
                outline=outline,
            )

        try:
            bbox = font.getbbox(str(text))
            tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
        except Exception:
            tw, th = len(str(text)) * 10, 16

        pad_x, pad_y = 8, 4
        chip_h = th + pad_y * 2
        chip_w = max(tw + pad_x * 2, chip_h)
        chip_y = y1 - chip_h - 4 if y1 - chip_h - 4 > 0 else y1 + 4
        chip_x = x1
        # 圆形编号更贴近参考图
        if str(text).isdigit() and len(str(text)) <= 2:
            r = max(chip_h, 22) // 2 + 2
            cx, cy = x1 + r + 2, max(y1 - r - 4, r + 2)
            odraw.ellipse([(cx - r, cy - r), (cx + r, cy + r)], fill=chip_bg)
            odraw.text((cx, cy), str(text), fill=chip_fg, font=font, anchor="mm")
        else:
            odraw.rectangle(
                [(chip_x, chip_y), (chip_x + chip_w, chip_y + chip_h)],
                fill=chip_bg,
            )
            odraw.text((chip_x + pad_x, chip_y + pad_y), str(text), fill=chip_fg, font=font)
        return Image.alpha_composite(rgba, overlay).convert("RGB")

    def visualize_results(self, image, results, save_path, labels=None):
        """
        在检测框上绘制标签（已过滤手/空盘）。
        labels 为 None 时画 类别+置信度；传入列表时画菜品名称（认菜不认盘展示）。
        """
        img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        font = self._load_font()
        kept = list(self.iter_kept_detections(results, image_shape=image.shape))
        for i, (box, conf, cname) in enumerate(kept):
            if labels is not None and i < len(labels):
                text = str(labels[i])
            else:
                text = f"{cname} {conf:.2f}"
            img = self._draw_box_and_chip(img, box, text, font)
        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
        img.save(save_path)
        return save_path

    def visualize_with_labels(self, image_or_path, boxes, labels, save_path, styles=None):
        """
        在原图上按框绘制菜品名称（不依赖 YOLO results）。
        :param image_or_path: BGR ndarray 或图片路径
        :param boxes: [[x1,y1,x2,y2], ...]
        :param labels: 与 boxes 对齐的菜品名称列表
        :param styles: 可选，与 boxes 对齐的颜色样式 dict 列表
        """
        if isinstance(image_or_path, str):
            image = cv2.imread(image_or_path)
            if image is None:
                raise ValueError(f"无法读取图片: {image_or_path}")
        else:
            image = image_or_path

        img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        font = self._load_font()
        styles = styles or []

        for i, (box, label) in enumerate(zip(boxes, labels)):
            style = styles[i] if i < len(styles) else None
            img = self._draw_box_and_chip(img, box, str(label), font, style=style)

        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
        img.save(save_path)
        return save_path

    def _load_bgr(self, image_or_path):
        if isinstance(image_or_path, np.ndarray):
            if image_or_path.size == 0:
                return None
            return image_or_path
        if isinstance(image_or_path, (str, Path)):
            image = cv2.imread(str(image_or_path))
            return image
        return None

    def process_image(
        self,
        image_or_path,
        stem=None,
        save_crops=True,
        visualize=None,
    ):
        """
        统一检测入口：支持路径或 BGR ndarray。

        :return: (visualize_save_path, crop_items, boxes)
                 crop_items: [{"image": ndarray, "path": str|None, "conf": float}, ...]
        """
        image = self._load_bgr(image_or_path)
        if image is None:
            print(f"❌ 读取失败: {image_or_path}")
            return None, [], []

        if stem is None:
            if isinstance(image_or_path, (str, Path)):
                stem = Path(image_or_path).stem
            else:
                stem = f"frame_{int(time.time() * 1000)}"

        # 内存直传，避免 YOLO 再读盘
        results = self.predict(image)
        cropped_regions = self.crop_plate(image, results)

        do_viz = self.config.GENERATE_VISUALIZATION if visualize is None else bool(visualize)
        visualize_save_path = None
        if do_viz:
            os.makedirs(self.config.OUTPUT_DIR, exist_ok=True)
            visualize_save_path = self.visualize_results(
                image, results,
                os.path.join(self.config.OUTPUT_DIR, f"{stem}.jpg"),
            )

        crop_items = []
        boxes = []
        if save_crops:
            os.makedirs(self.config.Cutting, exist_ok=True)
        for i, region in enumerate(cropped_regions):
            img_crop = region["image"]
            path = None
            if save_crops:
                filename = f"{stem}_Plate_{i + 1}.jpeg"
                path = os.path.join(self.config.Cutting, filename)
                cv2.imwrite(path, img_crop)
            crop_items.append({
                "image": img_crop,
                "path": path,
                "conf": region.get("conf"),
                "class_name": region.get("class_name"),
            })
            boxes.append(region["box"])

        print(f"✅ {stem} → 裁剪完成，共 {len(cropped_regions)} 张")
        return visualize_save_path, crop_items, boxes

    def process_single_image(self, image_path):
        """
        处理单张图片：检测餐盘、裁剪、生成初始可视化图。
        :return: (visualize_save_path, cropped_paths, boxes)
                 boxes 与 cropped_paths 一一对应，供后续标注菜品名
        """
        viz, crop_items, boxes = self.process_image(
            image_path, save_crops=True, visualize=True,
        )
        cropped_paths = [c["path"] for c in crop_items if c.get("path")]
        return viz, cropped_paths, boxes


def main():
    cfg = Config()
    app = YOLOv11PlateExtractor(cfg)
    img_path = r"C:\Users\王含波\Desktop\文件\classfiy\20260422095806.jpg"
    visualize_save_path, cropped_file_paths, boxes = app.process_single_image(img_path)
    print("\n📦 裁剪后的图片路径：")
    print(cropped_file_paths)
    print("boxes:", boxes)
    print("viz:", visualize_save_path)


if __name__ == "__main__":
    main()