import cv2
import numpy as np
from PIL import Image
import os
import sys


def detect_slider_button(image_path, output_path=None):
    """
    识别滑块验证码底部拖动滑块按钮的位置，返回其四边坐标，并用红框标记。
    改进版：只识别滑块轨道内的橙色小方块按钮，避免把轨道边框识别进去。

    参数:
        image_path: 输入图片路径
        output_path: 输出图片路径，默认在原图目录下生成 <原图名>_slider_red.png

    返回:
        coords: dict, 包含 top_left, top_right, bottom_left, bottom_right
        output_path: 红框标记图路径
    """
    img = Image.open(image_path)
    img_np = np.array(img)
    h, w = img_np.shape[:2]

    # 只检测图片下半部分
    start_y = int(h * 0.55)
    lower_half = img_np[start_y:, :]

    # 转换为HSV颜色空间
    hsv = cv2.cvtColor(lower_half, cv2.COLOR_RGB2HSV)

    # 检测橙色滑块按钮（严格范围，只保留高饱和高亮度的按钮内部）
    lower = np.array([8, 230, 230])
    upper = np.array([22, 255, 255])
    mask = cv2.inRange(hsv, lower, upper)

    # 形态学操作去噪
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

    # 查找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if not contours:
        raise ValueError("未检测到滑块按钮")

    # 筛选轮廓：近似正方形、面积适中、位于轨道区域内
    candidates = []
    for cnt in contours:
        x, y, bw, bh = cv2.boundingRect(cnt)
        area = bw * bh
        if area < 300 or bw < 30 or bh < 30:
            continue
        aspect_ratio = min(bw, bh) / max(bw, bh) if max(bw, bh) > 0 else 0
        if aspect_ratio < 0.5:
            continue
        candidates.append((cnt, area, aspect_ratio, x, y, bw, bh))

    if not candidates:
        raise ValueError("未找到合适的滑块按钮")

    # 选择最像滑块按钮的：面积适中、宽高比接近1
    best = max(candidates, key=lambda c: c[1] * c[2])
    _, _, _, x, y, bw, bh = best

    # 转换回全图坐标
    x = x
    y = y + start_y

    # 四边坐标
    top_left = (x, y)
    top_right = (x + bw, y)
    bottom_left = (x, y + bh)
    bottom_right = (x + bw, y + bh)

    coords = {
        "top_left": top_left,
        "top_right": top_right,
        "bottom_left": bottom_left,
        "bottom_right": bottom_right,
    }

    # 在原图上画红框
    canvas = img_np.copy()
    cv2.rectangle(canvas, top_left, bottom_right, (255, 0, 0), 2)

    # 保存红框标记图
    if output_path is None:
        base, ext = os.path.splitext(image_path)
        output_path = f"{base}_slider_red.png"

    cv2.imwrite(output_path, cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR))

    return coords, output_path


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("用法: python detect_slider_button.py <图片路径> [输出路径]")
        sys.exit(1)

    image_path = sys.argv[1]
    output_path = sys.argv[2] if len(sys.argv) > 2 else None

    coords, out = detect_slider_button(image_path, output_path)
    print(f"滑块按钮四边坐标:")
    print(f"  左上: {coords['top_left']}")
    print(f"  右上: {coords['top_right']}")
    print(f"  左下: {coords['bottom_left']}")
    print(f"  右下: {coords['bottom_right']}")
    print(f"红框标记图: {out}")