import cv2
from detect_bag import detect_bag

THRESHOLD_X = 537  # min_x 阈值

def bag_judgment(img, return_conf=True, return_vis=False):
    """
    判断图片中的料袋状态，可动态返回置信度和可视化图像
    Args:
        img (np.ndarray): 待检测图片
        return_conf (bool): 是否返回置信度
        return_vis (bool): 是否返回可视化图像
    Returns:
        status_bool: True=到位, False=未到位, None=未检测到
        status_text: 中文状态
        conf: 最大置信度或 None
        min_x: 最左边 x 坐标或 None
        vis_img: 可视化图像或 None
    """
    # 调用 detect_bag
    outputs = detect_bag(img, return_conf=return_conf, return_vis=return_vis)

    # 初始化占位
    conf = None
    min_x = None
    vis_img = None

    # 根据返回值长度解析
    if return_conf and return_vis:
        if len(outputs) == 3:
            conf, min_x, vis_img = outputs
        elif len(outputs) == 2:
            conf, min_x = outputs
        elif len(outputs) == 1:
            min_x = outputs[0]
    elif return_conf and not return_vis:
        if len(outputs) >= 2:
            conf, min_x = outputs[:2]
        elif len(outputs) == 1:
            min_x = outputs[0]
    elif not return_conf and return_vis:
        if len(outputs) == 2:
            min_x, vis_img = outputs
        elif len(outputs) == 1:
            min_x = outputs[0]
    else:
        min_x = outputs if isinstance(outputs, (int, float, np.number)) else outputs[0]

    # 判断状态
    if min_x is None:
        status_bool = None
        status_text = "没有料袋"
    elif min_x >= THRESHOLD_X:
        status_bool = True
        status_text = "料袋到位"
    else:
        status_bool = False
        status_text = "料袋未到位"

    return status_bool, status_text, conf, min_x, vis_img


# ====================== 测试 ======================
if __name__ == "__main__":
    IMG_PATH = "3.jpg"
    img = cv2.imread(IMG_PATH)
    if img is None:
        raise FileNotFoundError(f"图片无法读取: {IMG_PATH}")

    status_bool, status_text, conf, min_x, vis_img = bag_judgment(img, return_conf=True, return_vis=True)
    print(f"判断结果: {status_bool}, 中文状态: {status_text}, conf={conf}, min_x={min_x}")

    if vis_img is not None:
        cv2.imshow("Vis", vis_img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()