#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
根据背景图 + 拼图块图计算滑块目标偏移，并可选使用 DrissionPage 执行拖动。

默认读取当前项目根目录：
- 下载 (1).png  (背景图)
- 下载.png      (拼图块，RGBA 透明图)
"""
from __future__ import annotations

import argparse
import random
import time
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw


def _load_images(bg_path: Path, piece_path: Path) -> tuple[np.ndarray, np.ndarray]:
    bg = np.array(Image.open(bg_path).convert("RGB"), dtype=np.int16)
    piece_rgba = np.array(Image.open(piece_path).convert("RGBA"), dtype=np.int16)
    return bg, piece_rgba


def _piece_bbox(alpha: np.ndarray, threshold: int = 12) -> tuple[int, int, int, int, np.ndarray]:
    mask = alpha > threshold
    if not mask.any():
        raise ValueError("拼图块图片 alpha 全透明，无法匹配")

    ys, xs = np.where(mask)
    y0, y1 = int(ys.min()), int(ys.max()) + 1
    x0, x1 = int(xs.min()), int(xs.max()) + 1
    return x0, y0, x1, y1, mask


def calc_drag_distance(
    bg_path: Path,
    piece_path: Path,
    alpha_threshold: int = 12,
) -> dict:
    """返回最佳匹配位置与建议拖动距离。"""
    bg, piece_rgba = _load_images(bg_path, piece_path)

    bh, bw = bg.shape[:2]
    ph, pw = piece_rgba.shape[:2]
    if ph != bh:
        raise ValueError(f"图片高度不一致：背景={bh}, 拼图块={ph}")

    alpha = piece_rgba[:, :, 3]
    x0, y0, x1, y1, _ = _piece_bbox(alpha, threshold=alpha_threshold)

    piece_crop = piece_rgba[y0:y1, x0:x1, :3]
    mask = alpha[y0:y1, x0:x1] > alpha_threshold
    ys, xs = np.where(mask)
    piece_pixels = piece_crop[ys, xs]

    patch_h, patch_w = piece_crop.shape[:2]
    if patch_w > bw or patch_h > bh:
        raise ValueError("拼图块裁剪后尺寸超过背景图")

    max_x = bw - patch_w
    best_x = 0
    best_score = float("inf")
    second_best = float("inf")

    # 固定 y，按 x 扫描。评分越小越好。
    for x in range(max_x + 1):
        patch_pixels = bg[y0 + ys, x + xs]
        score = float(np.abs(patch_pixels - piece_pixels).mean())
        if score < best_score:
            second_best = best_score
            best_score = score
            best_x = x
        elif score < second_best:
            second_best = score

    # 如果拼图初始从最左侧出发，建议拖动距离可用：best_x - x0
    drag_distance = best_x - x0
    confidence_ratio = (second_best / best_score) if best_score > 0 else float("inf")

    return {
        "target_x": int(best_x),
        "piece_bbox_x0": int(x0),
        "piece_bbox_y0": int(y0),
        "piece_bbox_w": int(patch_w),
        "piece_bbox_h": int(patch_h),
        "drag_distance": int(drag_distance),
        "best_score": best_score,
        "second_best": second_best,
        "confidence_ratio": confidence_ratio,
        "bg_width": int(bw),
        "bg_height": int(bh),
    }


def save_debug_overlay(bg_path: Path, out_path: Path, match: dict) -> None:
    img = Image.open(bg_path).convert("RGB")
    draw = ImageDraw.Draw(img)
    x = int(match["target_x"])
    y = int(match["piece_bbox_y0"])
    w = int(match["piece_bbox_w"])
    h = int(match["piece_bbox_h"])
    draw.rectangle([x, y, x + w, y + h], outline=(255, 0, 0), width=2)
    img.save(out_path)


def build_human_track(distance: int) -> list[int]:
    """生成拟人拖动轨迹（整数步长列表）。"""
    if distance == 0:
        return []

    direction = 1 if distance > 0 else -1
    remaining = abs(int(round(distance)))

    track: list[int] = []
    moved = 0
    velocity = 0.0
    interval = 0.02
    accelerate_until = remaining * 0.7

    while moved < remaining:
        if moved < accelerate_until:
            acc = random.uniform(2.0, 4.0)
        else:
            acc = -random.uniform(3.0, 6.0)

        step = max(1, int(velocity * interval + 0.5 * acc * interval * interval * 100))
        moved += step
        velocity = max(0.0, velocity + acc * interval)
        track.append(step)

    overflow = moved - remaining
    if overflow > 0 and track:
        track[-1] -= overflow
        if track[-1] <= 0:
            track.pop()

    # 小幅回拉，增加拟人性
    if remaining >= 30:
        back = random.randint(1, 3)
        track.extend([-back, back])

    return [direction * s for s in track if s != 0]


def drag_slider_with_dp(page, slider_selector: str, distance: int) -> None:
    slider = page.ele(slider_selector, timeout=8)
    if not slider:
        raise RuntimeError(f"未找到滑块元素：{slider_selector}")

    # 若元素支持直接 drag，优先用内置实现
    if hasattr(slider, "drag"):
        try:
            slider.drag(offset_x=distance, offset_y=0, duration=0.8)
            return
        except Exception:
            pass

    actions = page.actions
    actions.move_to(slider, duration=0.2)

    hold_ok = False
    for hold_call in (
        lambda: actions.hold(slider),
        lambda: actions.hold(),
    ):
        try:
            hold_call()
            hold_ok = True
            break
        except Exception:
            continue

    if not hold_ok:
        raise RuntimeError("无法按下滑块，请确认 DrissionPage 版本与页面状态")

    for step in build_human_track(distance):
        dy = random.randint(-1, 1)
        moved = False
        for move_call in (
            lambda: actions.move(step, dy, duration=random.uniform(0.01, 0.04)),
            lambda: actions.move(offset_x=step, offset_y=dy, duration=random.uniform(0.01, 0.04)),
            lambda: actions.move(step, dy),
            lambda: actions.move(offset_x=step, offset_y=dy),
        ):
            try:
                move_call()
                moved = True
                break
            except Exception:
                continue
        if not moved:
            raise RuntimeError("动作链 move 调用失败，请按当前 DrissionPage 版本调整 move 参数")

    released = False
    for release_call in (
        lambda: actions.release(),
        lambda: actions.release(on_ele=None),
    ):
        try:
            release_call()
            released = True
            break
        except Exception:
            continue

    if not released:
        raise RuntimeError("动作链 release 调用失败")


def maybe_run_dp_drag(args: argparse.Namespace, distance: int) -> None:
    if not args.url or not args.slider:
        return

    from DrissionPage import ChromiumOptions, ChromiumPage

    if args.port:
        co = ChromiumOptions().set_local_port(port=args.port)
        page = ChromiumPage(addr_or_opts=co)
    else:
        page = ChromiumPage()

    page.get(args.url)
    if args.wait_after_open > 0:
        time.sleep(args.wait_after_open)

    drag_slider_with_dp(page, args.slider, distance)


def parse_args() -> argparse.Namespace:
    root = Path(__file__).resolve().parent

    parser = argparse.ArgumentParser(description="DP 滑块验证辅助脚本")
    parser.add_argument("--bg", type=Path, default=root / "下载 (1).png", help="背景图路径")
    parser.add_argument("--piece", type=Path, default=root / "下载.png", help="拼图块图路径")
    parser.add_argument("--alpha-threshold", type=int, default=12, help="透明阈值(0-255)")
    parser.add_argument("--debug-out", type=Path, default=root / "slider_match_debug.png", help="匹配结果标注图")

    # 以下参数用于真实页面拖动（可选）
    parser.add_argument("--url", type=str, default="", help="验证码页面 URL")
    parser.add_argument("--slider", type=str, default="", help="滑块元素选择器")
    parser.add_argument("--port", type=int, default=0, help="连接已有浏览器的本地端口")
    parser.add_argument("--distance-adjust", type=int, default=0, help="对计算结果追加修正像素")
    parser.add_argument("--wait-after-open", type=float, default=1.0, help="打开页面后等待秒数")

    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if not args.bg.exists():
        raise FileNotFoundError(f"背景图不存在：{args.bg}")
    if not args.piece.exists():
        raise FileNotFoundError(f"拼图块图不存在：{args.piece}")

    match = calc_drag_distance(args.bg, args.piece, alpha_threshold=args.alpha_threshold)
    save_debug_overlay(args.bg, args.debug_out, match)

    distance = int(match["drag_distance"]) + int(args.distance_adjust)

    print("匹配结果：")
    print(f"  target_x={match['target_x']}")
    print(f"  piece_bbox_x0={match['piece_bbox_x0']}")
    print(f"  建议拖动距离 drag_distance={match['drag_distance']}")
    print(f"  调整后拖动距离 distance={distance}")
    print(f"  best_score={match['best_score']:.4f}, second_best={match['second_best']:.4f}")
    print(f"  confidence_ratio={match['confidence_ratio']:.4f} (越大越好)")
    print(f"  标注图已输出：{args.debug_out}")

    if args.url and args.slider:
        maybe_run_dp_drag(args, distance)
        print("DP 拖动已执行。")


if __name__ == "__main__":
    main()
